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
}

// downsampleBins: 分段映射 — 低频线性保细节，中高频对数保覆盖
// 不再做额外 tilt 补偿，映射本身就是均衡的
func downsampleBins(magnitudes []float64, target int) []float64 {
	srcLen := len(magnitudes)
	if srcLen == 0 || target == 0 {
		return make([]float64, target)
	}

	// 低频线性区间：前20%的 display bins 覆盖 FFT 1-80 (约47Hz-3.7kHz)
	// 中高频对数区间：后80%覆盖 FFT 80-512 (3.7kHz-24kHz)
	const lowEnd = 80 // 低频截止 FFT 索引
	const lowFrac = 0.2

	bins := make([]float64, target)
	for i := 0; i < target; i++ {
		frac := float64(i) / float64(target)
		fracNext := float64(i+1) / float64(target)

		var start, end int
		if frac < lowFrac {
			start = int(frac / lowFrac * float64(lowEnd))
			end = int(fracNext / lowFrac * float64(lowEnd))
		} else {
			f := (frac - lowFrac) / (1 - lowFrac)
			fNext := (fracNext - lowFrac) / (1 - lowFrac)
			start = lowEnd + int(float64(srcLen-lowEnd)*math.Pow(f, 0.7))
			end = lowEnd + int(float64(srcLen-lowEnd)*math.Pow(fNext, 0.7))
		}

		if end <= start {
			end = start + 1
		}
		if end > srcLen {
			end = srcLen
		}

		// 取最大值保留峰值冲击力
		maxVal := 0.0
		for j := start; j < end; j++ {
			if magnitudes[j] > maxVal {
				maxVal = magnitudes[j]
			}
		}
		bins[i] = maxVal
	}
	return bins
}

func smoothBinsGo(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))))
	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}}
)

const (
	enumeratorGetDefaultAudioEndpoint = 4
)

const (
	deviceActivate = 3
)

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

const (
	captureGetBuffer     = 3
	captureReleaseBuffer = 4
)

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

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

func startAudioCapture() {
	audioMu.Lock()
	defer audioMu.Unlock()
	if audioActive {
		return
	}
	audioStop = make(chan struct{})
	audioDone = make(chan struct{})
	audioActive = true
	smoothed = nil
	runningMax = 0
	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()

	ole32dll.NewProc("CoInitializeEx").Call(0, 0)
	defer ole32dll.NewProc("CoUninitialize").Call()

	var enumerator uintptr
	hr, _, _ := ole32dll.NewProc("CoCreateInstance").Call(
		uintptr(unsafe.Pointer(&clsidMMDeviceEnumerator)),
		0, 1,
		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)

	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)

	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)

	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)

	hr = comCall(audioClient, audioClientInitialize,
		0,
		0x00020000,
		10000000,
		0,
		pwfx,
		0,
	)
	if hr != 0 {
		log.Printf("音频捕获: Initialize 失败 hr=0x%x", hr)
		close(audioDone)
		return
	}

	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)

	comCall(audioClient, audioClientStart)
	defer comCall(audioClient, audioClientStop)

	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
	}

	// 预分配 mono slice
	mono := make([]float32, fftSize)

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

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

		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))
		}

		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 = smoothBinsGo(smoothed, bins, smoothingFactor)
		}
		smoothed = bins

		// 全局归一化（慢衰减 runningMax）
		curMax := 0.0
		for _, v := range bins {
			if v > curMax {
				curMax = v
			}
		}
		if curMax > runningMax {
			runningMax = curMax
		} else {
			runningMax *= 0.985
		}
		if runningMax > 0.001 {
			for i := range bins {
				bins[i] /= runningMax
			}
		}

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

func comRelease(obj uintptr) {
	if obj != 0 {
		comCall(obj, 2)
	}
}
