package main

import (
	"encoding/json"
	"fmt"
	"log"
	"math"
	"math/cmplx"
	"runtime"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
)

// ============================================================
// FFT: 纯 Go radix-2 Cooley-Tukey (1024 点)
// ============================================================

const fftSize = 1024
const targetBins = 128

func fft(x []complex128) {
	n := len(x)
	if n <= 1 {
		return
	}
	for i, j := 1, 0; i < n; i++ {
		bit := n >> 1
		for j&bit != 0 {
			j ^= bit
			bit >>= 1
		}
		j ^= bit
		if i < j {
			x[i], x[j] = x[j], x[i]
		}
	}
	for length := 2; length <= n; length <<= 1 {
		half := length >> 1
		angle := -2 * math.Pi / float64(length)
		wn := complex(math.Cos(angle), math.Sin(angle))
		for i := 0; i < n; i += length {
			w := complex(1, 0)
			for j := 0; j < half; j++ {
				u := x[i+j]
				v := x[i+j+half] * w
				x[i+j] = u + v
				x[i+j+half] = u - v
				w *= wn
			}
		}
	}
}

func fftMagnitude(input []float32) []float64 {
	n := len(input)
	if n > fftSize {
		n = fftSize
		input = input[:fftSize]
	}
	x := make([]complex128, fftSize)
	for i := 0; i < n; i++ {
		w := 0.5 * (1 - math.Cos(2*math.Pi*float64(i)/float64(n-1)))
		x[i] = complex(float64(input[i])*w, 0)
	}
	fft(x)
	mag := make([]float64, fftSize/2)
	for i := range mag {
		mag[i] = cmplx.Abs(x[i]) / float64(fftSize)
	}
	return mag
}

func downsampleBins(magnitudes []float64, target int) []float64 {
	srcLen := len(magnitudes)
	if srcLen == 0 || target == 0 {
		return make([]float64, target)
	}
	bins := make([]float64, target)
	for i := 0; i < target; i++ {
		frac := float64(i) / float64(target)
		start := int(math.Pow(float64(srcLen), frac))
		fracNext := float64(i+1) / float64(target)
		end := int(math.Pow(float64(srcLen), fracNext))
		if end <= start {
			end = start + 1
		}
		if end > srcLen {
			end = srcLen
		}
		sum := 0.0
		for j := start; j < end; j++ {
			sum += magnitudes[j]
		}
		bins[i] = sum / float64(end-start)
	}
	return bins
}

func smoothBins(prev, curr []float64, factor float64) []float64 {
	n := len(curr)
	if len(prev) != n {
		return curr
	}
	out := make([]float64, n)
	for i := 0; i < n; i++ {
		out[i] = prev[i]*factor + curr[i]*(1-factor)
	}
	return out
}

// ============================================================
// COM vtable 调用辅助
// ============================================================

func comCall(obj uintptr, vtableIdx uintptr, args ...uintptr) uintptr {
	if obj == 0 {
		return 0x80004003 // E_POINTER
	}
	vtable := *(*uintptr)(unsafe.Pointer(obj))
	method := *(*uintptr)(unsafe.Pointer(vtable + vtableIdx*unsafe.Sizeof(uintptr(0))))
	// 构建参数数组 (this + args)，最多支持 12 参数
	a := make([]uintptr, 12)
	a[0] = obj
	for i, arg := range args {
		a[i+1] = arg
	}
	n := uintptr(len(args) + 1)
	switch {
	case n <= 3:
		r, _, _ := syscall.Syscall(method, n, a[0], a[1], a[2])
		return r
	case n <= 6:
		r, _, _ := syscall.Syscall6(method, n, a[0], a[1], a[2], a[3], a[4], a[5])
		return r
	case n <= 9:
		r, _, _ := syscall.Syscall9(method, n, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8])
		return r
	case n <= 12:
		r, _, _ := syscall.Syscall12(method, n, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8], a[9], a[10], a[11])
		return r
	default:
		return 0x80070057 // E_INVALIDARG
	}
}

// COM GUID 常量
var (
	clsidMMDeviceEnumerator = windows.GUID{Data1: 0xBCDE0395, Data2: 0xE52F, Data3: 0x467C, Data4: [8]byte{0x8E, 0x3D, 0xC4, 0x57, 0x92, 0x91, 0x69, 0x2E}}
	iidIMMDeviceEnumerator  = windows.GUID{Data1: 0xA95664D2, Data2: 0x9614, Data3: 0x4F35, Data4: [8]byte{0xA7, 0x46, 0xDE, 0x8D, 0xB6, 0x36, 0x17, 0xE6}}
	iidIAudioClient         = windows.GUID{Data1: 0x1CB9AD4C, Data2: 0xDBFA, Data3: 0x4C32, Data4: [8]byte{0xB1, 0x78, 0xC2, 0xF5, 0x68, 0xA7, 0x03, 0xB2}}
	iidIAudioCaptureClient   = windows.GUID{Data1: 0xC8ADBD64, Data2: 0xE71E, Data3: 0x48A0, Data4: [8]byte{0xA4, 0xDE, 0x18, 0x5C, 0x39, 0x5C, 0xD3, 0x17}}
)

// IMMDeviceEnumerator vtable (IUnknown 3 + 自身方法)
const (
	enumeratorGetDefaultAudioEndpoint = 4 // 第 4 个方法 (0-based from vtable start)
)

// IMMDevice vtable
const (
	deviceActivate = 3 // 第 3 个方法
)

// IAudioClient vtable
const (
	audioClientInitialize   = 3
	audioClientGetBufferSize = 4
	audioClientGetMixFormat  = 8
	audioClientStart         = 10
	audioClientStop          = 11
	audioClientGetService    = 14
)

// IAudioCaptureClient vtable
const (
	captureGetBuffer     = 3
	captureReleaseBuffer = 4
)

// ============================================================
// WASAPI 捕获生命周期
// ============================================================

var (
	audioMu     sync.Mutex
	audioStop   chan struct{}
	audioDone   chan struct{}
	audioActive bool
	smoothed    []float64
)

func startAudioCapture() {
	audioMu.Lock()
	defer audioMu.Unlock()
	if audioActive {
		return
	}
	audioStop = make(chan struct{})
	audioDone = make(chan struct{})
	audioActive = true
	smoothed = nil
	go audioCaptureLoop()
	log.Println("音频捕获: 已启动")
}

func stopAudioCapture() {
	audioMu.Lock()
	stop := audioStop
	audioMu.Unlock()
	if stop == nil {
		return
	}
	close(stop)
	if audioDone != nil {
		<-audioDone
	}
	audioMu.Lock()
	audioActive = false
	audioStop = nil
	audioDone = nil
	audioMu.Unlock()
	log.Println("音频捕获: 已停止")
}

func audioActiveState() bool {
	audioMu.Lock()
	defer audioMu.Unlock()
	return audioActive
}

func audioCaptureLoop() {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	// COM 初始化
	ole32dll.NewProc("CoInitializeEx").Call(0, 0)
	defer ole32dll.NewProc("CoUninitialize").Call()

	// CoCreateInstance -> IMMDeviceEnumerator
	var enumerator uintptr
	hr, _, _ := ole32dll.NewProc("CoCreateInstance").Call(
		uintptr(unsafe.Pointer(&clsidMMDeviceEnumerator)),
		0, 1, // CLSCTX_INPROC_SERVER
		uintptr(unsafe.Pointer(&iidIMMDeviceEnumerator)),
		uintptr(unsafe.Pointer(&enumerator)),
	)
	if hr != 0 || enumerator == 0 {
		log.Printf("音频捕获: 创建枚举器失败 hr=0x%x", hr)
		close(audioDone)
		return
	}
	defer comRelease(enumerator)

	// enumerator->GetDefaultAudioEndpoint(eRender=0, eConsole=0)
	var device uintptr
	hr = comCall(enumerator, enumeratorGetDefaultAudioEndpoint, 0, 0, uintptr(unsafe.Pointer(&device)))
	if hr != 0 || device == 0 {
		log.Printf("音频捕获: 获取默认音频设备失败 hr=0x%x", hr)
		close(audioDone)
		return
	}
	defer comRelease(device)

	// device->Activate(IID_IAudioClient, CLSCTX_INPROC=1, NULL, &audioClient)
	var audioClient uintptr
	hr = comCall(device, deviceActivate,
		uintptr(unsafe.Pointer(&iidIAudioClient)),
		1, 0,
		uintptr(unsafe.Pointer(&audioClient)),
	)
	if hr != 0 || audioClient == 0 {
		log.Printf("音频捕获: Activate(IAudioClient) 失败 hr=0x%x", hr)
		close(audioDone)
		return
	}
	defer comRelease(audioClient)

	// audioClient->GetMixFormat(&pwfx)
	var pwfx uintptr
	hr = comCall(audioClient, audioClientGetMixFormat, uintptr(unsafe.Pointer(&pwfx)))
	if hr != 0 || pwfx == 0 {
		log.Printf("音频捕获: GetMixFormat 失败 hr=0x%x", hr)
		close(audioDone)
		return
	}
	defer ole32dll.NewProc("CoTaskMemFree").Call(pwfx)

	nChannels := *(*uint16)(unsafe.Pointer(pwfx + 2))
	nSamplesPerSec := *(*uint32)(unsafe.Pointer(pwfx + 4))
	log.Printf("音频捕获: 格式 ch=%d rate=%d", nChannels, nSamplesPerSec)

	// audioClient->Initialize(SHARED, LOOPBACK, 1s buffer, 0, pwfx, NULL)
	hr = comCall(audioClient, audioClientInitialize,
		0, // AUDCLNT_SHAREMODE_SHARED
		0x00020000, // AUDCLNT_STREAMFLAGS_LOOPBACK
		10000000,   // 1 second (100ns units)
		0,
		pwfx,
		0,
	)
	if hr != 0 {
		log.Printf("音频捕获: Initialize 失败 hr=0x%x", hr)
		close(audioDone)
		return
	}

	// audioClient->GetService(IID_IAudioCaptureClient, &captureClient)
	var captureClient uintptr
	hr = comCall(audioClient, audioClientGetService,
		uintptr(unsafe.Pointer(&iidIAudioCaptureClient)),
		uintptr(unsafe.Pointer(&captureClient)),
	)
	if hr != 0 || captureClient == 0 {
		log.Printf("音频捕获: GetService(IAudioCaptureClient) 失败 hr=0x%x", hr)
		close(audioDone)
		return
	}
	defer comRelease(captureClient)

	// audioClient->Start()
	comCall(audioClient, audioClientStart)
	defer comCall(audioClient, audioClientStop)

	// 环形缓冲区 (~100ms)
	ch := int(nChannels)
	bufSize := int(nSamplesPerSec)/10*ch + fftSize*ch
	ringBuf := make([]float32, bufSize)
	writePos := 0

	cfg := loadConfig()
	smoothingFactor := cfg.AudioSmoothing
	if smoothingFactor <= 0 {
		smoothingFactor = 0.7
	}

	ticker := time.NewTicker(33 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-audioStop:
			close(audioDone)
			return
		case <-ticker.C:
		}

		// 读取 WASAPI 缓冲区
		for {
			var pData, nFrames, flags uintptr
			hr = comCall(captureClient, captureGetBuffer,
				uintptr(unsafe.Pointer(&pData)),
				uintptr(unsafe.Pointer(&nFrames)),
				uintptr(unsafe.Pointer(&flags)),
				0, 0,
			)
			if hr != 0 || nFrames == 0 {
				break
			}
			samples := unsafe.Slice((*float32)(unsafe.Pointer(pData)), nFrames*uintptr(ch))
			for _, s := range samples {
				ringBuf[writePos] = s
				writePos = (writePos + 1) % bufSize
			}
			comCall(captureClient, captureReleaseBuffer, uintptr(nFrames))
		}

		// 取 fftSize 个单声道样本
		mono := make([]float32, fftSize)
		for i := 0; i < fftSize; i++ {
			pos := (writePos - fftSize*ch + i*ch + bufSize) % bufSize
			var sum float32
			for c := 0; c < ch; c++ {
				sum += ringBuf[(pos+c)%bufSize]
			}
			mono[i] = sum / float32(ch)
		}

		mag := fftMagnitude(mono)
		bins := downsampleBins(mag, targetBins)

		if smoothed != nil {
			bins = smoothBins(smoothed, bins, smoothingFactor)
		}
		smoothed = bins

		// 归一化
		maxVal := 0.0
		for _, v := range bins {
			if v > maxVal {
				maxVal = v
			}
		}
		if maxVal > 0.001 {
			for i := range bins {
				bins[i] /= maxVal
			}
		}

		data, _ := json.Marshal(bins)
		evalJS(fmt.Sprintf(`if(window.updateAudioDataFromGo) updateAudioDataFromGo(%s)`, string(data)))
	}
}

func comRelease(obj uintptr) {
	if obj != 0 {
		comCall(obj, 2) // IUnknown::Release = vtable index 2
	}
}