package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"time"
	"voicesnap/internal/engine"
)

func main() {
	wavPath := flag.String("wav", "", "16kHz mono PCM16 WAV file")
	runs := flag.Int("runs", 3, "number of recognition runs")
	flag.Parse()

	if *wavPath == "" {
		fmt.Fprintln(os.Stderr, "missing -wav")
		os.Exit(2)
	}
	if *runs < 1 {
		*runs = 1
	}

	samples, sampleRate, err := readPCM16WAV(*wavPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "read wav: %v\n", err)
		os.Exit(1)
	}
	if sampleRate != 16000 {
		fmt.Fprintf(os.Stderr, "unsupported sample rate %d: provide 16kHz mono PCM16 WAV\n", sampleRate)
		os.Exit(1)
	}

	eng, err := engine.New()
	if err != nil {
		fmt.Fprintf(os.Stderr, "init engine: %v\n", err)
		os.Exit(1)
	}
	defer eng.Close()

	audioMS := int64(len(samples)) * 1000 / 16000
	fmt.Printf("engine=%q samples=%d audio_ms=%d runs=%d\n", eng.HardwareInfo(), len(samples), audioMS, *runs)
	var total int64
	for i := 1; i <= *runs; i++ {
		start := time.Now()
		text, err := eng.Recognize(samples)
		elapsedMS := time.Since(start).Milliseconds()
		if err != nil {
			fmt.Fprintf(os.Stderr, "run=%d error=%v\n", i, err)
			os.Exit(1)
		}
		total += elapsedMS
		fmt.Printf("run=%d recognize_ms=%d rtf=%.3f text_runes=%d\n", i, elapsedMS, float64(elapsedMS)/float64(audioMS), len([]rune(text)))
	}
	fmt.Printf("avg_recognize_ms=%d\n", total/int64(*runs))
}

func readPCM16WAV(path string) ([]float32, int, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, 0, err
	}
	if len(data) < 12 || string(data[0:4]) != "RIFF" || string(data[8:12]) != "WAVE" {
		return nil, 0, fmt.Errorf("not a RIFF/WAVE file")
	}

	var audioFormat, channels, bitsPerSample uint16
	var sampleRate uint32
	var pcm []byte
	for offset := 12; offset+8 <= len(data); {
		chunkID := string(data[offset : offset+4])
		chunkSize := int(binary.LittleEndian.Uint32(data[offset+4 : offset+8]))
		chunkStart := offset + 8
		chunkEnd := chunkStart + chunkSize
		if chunkEnd > len(data) {
			return nil, 0, fmt.Errorf("truncated %s chunk", chunkID)
		}

		switch chunkID {
		case "fmt ":
			if chunkSize < 16 {
				return nil, 0, fmt.Errorf("fmt chunk too small")
			}
			audioFormat = binary.LittleEndian.Uint16(data[chunkStart : chunkStart+2])
			channels = binary.LittleEndian.Uint16(data[chunkStart+2 : chunkStart+4])
			sampleRate = binary.LittleEndian.Uint32(data[chunkStart+4 : chunkStart+8])
			bitsPerSample = binary.LittleEndian.Uint16(data[chunkStart+14 : chunkStart+16])
		case "data":
			pcm = data[chunkStart:chunkEnd]
		}

		offset = chunkEnd
		if offset%2 == 1 {
			offset++
		}
	}

	if audioFormat != 1 || channels != 1 || bitsPerSample != 16 {
		return nil, 0, fmt.Errorf("unsupported WAV format: audio_format=%d channels=%d bits_per_sample=%d; want PCM16 mono", audioFormat, channels, bitsPerSample)
	}
	if len(pcm) == 0 {
		return nil, 0, fmt.Errorf("missing data chunk")
	}
	if len(pcm)%2 != 0 {
		return nil, 0, fmt.Errorf("odd PCM data length")
	}

	samples := make([]float32, len(pcm)/2)
	for i := range samples {
		v := int16(binary.LittleEndian.Uint16(pcm[i*2 : i*2+2]))
		samples[i] = float32(v) / 32768.0
	}
	return samples, int(sampleRate), nil
}