package main

import (
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/gen2brain/malgo"
	"github.com/go-audio/audio"
	"github.com/go-audio/wav"
	hook "github.com/robotn/gohook"
)

const (
	sampleRate   = 44100
	channelCount = 1
	bitDepth     = 16
	bufferSize   = 1024
)

type Recorder struct {
	context     *malgo.AllocatedContext
	device      *malgo.Device
	isRecording bool
	audioData   []byte
	lock        sync.Mutex
	wg          sync.WaitGroup
}

func NewRecorder() (*Recorder, error) {
	ctx, err := malgo.InitContext(nil, malgo.ContextConfig{}, func(message string) {
		fmt.Printf("MALGO LOG: %s\n", message)
	})
	if err != nil {
		return nil, err
	}

	return &Recorder{
		context: ctx,
	}, nil
}

func (r *Recorder) Start() {
	r.lock.Lock()
	defer r.lock.Unlock()

	if r.isRecording {
		return
	}

	r.audioData = make([]byte, 0, 1024*1024) // 预分配1MB
	r.isRecording = true

	deviceConfig := malgo.DefaultDeviceConfig(malgo.Capture)
	deviceConfig.Capture.Format = malgo.FormatS16
	deviceConfig.Capture.Channels = channelCount
	deviceConfig.SampleRate = sampleRate
	deviceConfig.Alsa.NoMMap = 1

	// 回调函数处理音频数据
	onRecvFrames := func(pSample2, pSample []byte, framecount uint32) {
		if !r.isRecording {
			return
		}
		r.lock.Lock()
		r.audioData = append(r.audioData, pSample...)
		r.lock.Unlock()
	}

	var err error
	r.device, err = malgo.InitDevice(r.context.Context, deviceConfig, malgo.DeviceCallbacks{
		Data: onRecvFrames,
	})
	if err != nil {
		fmt.Println("设备初始化失败:", err)
		return
	}

	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		if err := r.device.Start(); err != nil {
			fmt.Println("录音设备启动失败:", err)
		}
	}()

	fmt.Println("录音开始...")
}

func (r *Recorder) Stop(filename string) {
	r.lock.Lock()
	defer r.lock.Unlock()

	if !r.isRecording {
		return
	}

	r.isRecording = false
	time.Sleep(100 * time.Millisecond) // 等待最后数据写入

	if r.device != nil {
		r.device.Uninit()
		r.device = nil
	}

	if len(r.audioData) == 0 {
		fmt.Println("无录音数据")
		return
	}

	// 保存为WAV文件
	if err := saveWAV(filename, r.audioData); err != nil {
		fmt.Println("保存失败:", err)
		return
	}

	fmt.Printf("录音已保存: %s (大小: %.2f KB)\n",
		filename, float64(len(r.audioData))/1024)
	r.audioData = nil
}

func (r *Recorder) Close() {
	if r.context != nil {
		r.context.Uninit()
		r.context = nil
	}
}

func saveWAV(filename string, data []byte) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	// 创建WAV编码器
	enc := wav.NewEncoder(file,
		sampleRate,
		bitDepth,
		channelCount,
		1) // PCM格式

	// 将byte转换为int16
	numSamples := len(data) / 2
	intData := make([]int, numSamples)
	for i := 0; i < numSamples; i++ {
		// 小端字节序转换
		val := int16(data[i*2]) | int16(data[i*2+1])<<8
		intData[i] = int(val)
	}

	// 创建音频Buffer
	buf := &audio.IntBuffer{
		Format: &audio.Format{
			SampleRate:  sampleRate,
			NumChannels: channelCount,
		},
		Data:           intData,
		SourceBitDepth: bitDepth,
	}

	// 写入WAV文件
	if err := enc.Write(buf); err != nil {
		return err
	}
	return enc.Close()
}

func main() {
	recorder, err := NewRecorder()
	if err != nil {
		fmt.Println("录音器初始化失败:", err)
		return
	}
	defer recorder.Close()

	fmt.Println("按下 T 键开始录音，松开 T 键保存录音")
	fmt.Println("按 ESC 退出程序")

	// 创建事件通道
	evChan := hook.Start()
	defer hook.End()

	// 状态跟踪
	isTDown := false
	fileCounter := 1

	// 事件处理循环
	for ev := range evChan {
		// 只处理键盘事件
		if ev.Kind != hook.KeyDown && ev.Kind != hook.KeyUp {
			continue
		}
		// 检查 T 键
		if ev.Rawcode == 84 { // T 键的键码

			if ev.Kind == hook.KeyDown && !isTDown {
				isTDown = true
				recorder.Start()
			} else if ev.Kind == hook.KeyUp && isTDown {
				isTDown = false
				filename := fmt.Sprintf("recording_%d.wav", fileCounter)
				fileCounter++
				recorder.Stop(filename)
			}
		}

		// 检查 ESC 键
		if ev.Rawcode == 27 && ev.Kind == hook.KeyDown {
			fmt.Println("退出程序")
			return
		}
	}
}