// cmd/wmdecode — fm-rds-tx spread-spectrum watermark recovery tool.
//
// Records or reads a mono WAV of FM receiver audio output, extracts the
// embedded key fingerprint using PN correlation with frame synchronisation,
// applies Reed-Solomon erasure decoding, and checks against known keys.
//
// Usage:
//
//	wmdecode <file.wav> [key ...]
//
// Examples:
//
//	wmdecode aufnahme.wav
//	wmdecode aufnahme.wav free studio@sender.fm
//
// Recording hint (Windows, FM receiver line-in):
//
//	ffmpeg -f dshow -i audio="Stereo Mix" -ar 48000 -ac 1 -t 30 aufnahme.wav
package main

import (
	"encoding/binary"
	"fmt"
	"math"
	"os"
	"sort"

	"github.com/jan/fm-rds-tx/internal/watermark"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "usage: wmdecode <file.wav> [key ...]")
		os.Exit(1)
	}

	samples, recRate, err := readMonoWAV(os.Args[1])
	if err != nil {
		fmt.Fprintf(os.Stderr, "read WAV: %v\n", err)
		os.Exit(1)
	}

	rms := rmsLevel(samples)
	fmt.Printf("WAV:   %d samples @ %.0f Hz = %.2fs, RMS %.1f dBFS\n",
		len(samples), recRate, float64(len(samples))/recRate, 20*math.Log10(rms+1e-9))

	samplesPerBit := int(float64(watermark.PnChips) * recRate / float64(watermark.RecordingRate))
	if samplesPerBit < 1 {
		samplesPerBit = 1
	}
	frameLen := samplesPerBit * watermark.PayloadBits
	fmt.Printf("Frame: %d samples/bit, %d samples/frame (%.3fs), %d frames in recording\n",
		samplesPerBit, frameLen, float64(frameLen)/recRate, len(samples)/frameLen)

	if len(samples) < samplesPerBit*2 {
		fmt.Fprintln(os.Stderr, "recording too short for even 2 bits")
		os.Exit(1)
	}

	// ---------------------------------------------------------------
	// Step 1: Phase search — find sample offset of bit boundaries.
	//
	// Coarse pass: test every 8th offset in [0, samplesPerBit).
	// Fine pass:   refine ±8 around the coarse peak.
	// For each candidate offset, average |correlation| over several bits.
	// ---------------------------------------------------------------
	const coarseStep = 8
	const syncBits = 64

	bestPhase := 0
	bestMag := 0.0

	for phase := 0; phase < samplesPerBit; phase += coarseStep {
		mag := avgCorrMag(samples, phase, samplesPerBit, syncBits, recRate)
		if mag > bestMag {
			bestMag = mag
			bestPhase = phase
		}
	}

	fineStart := bestPhase - coarseStep
	if fineStart < 0 {
		fineStart = 0
	}
	fineEnd := bestPhase + coarseStep
	if fineEnd > samplesPerBit {
		fineEnd = samplesPerBit
	}
	for phase := fineStart; phase < fineEnd; phase++ {
		mag := avgCorrMag(samples, phase, samplesPerBit, syncBits, recRate)
		if mag > bestMag {
			bestMag = mag
			bestPhase = phase
		}
	}

	fmt.Printf("Phase: offset=%d (%.3fms into recording), avg|corr|=%.4f\n",
		bestPhase, float64(bestPhase)/recRate*1000, bestMag)

	// ---------------------------------------------------------------
	// Step 2: Extract bit correlations at found phase, averaged over frames.
	// ---------------------------------------------------------------
	nCompleteBits := (len(samples) - bestPhase) / samplesPerBit
	nFrames := nCompleteBits / watermark.PayloadBits
	if nFrames == 0 {
		nFrames = 1
	}

	fmt.Printf("Sync:  %d complete bits, %d usable frames\n", nCompleteBits, nFrames)

	corrs := make([]float64, watermark.PayloadBits)
	for i := 0; i < watermark.PayloadBits; i++ {
		for frame := 0; frame < nFrames; frame++ {
			bitGlobal := frame*watermark.PayloadBits + i
			start := bestPhase + bitGlobal*samplesPerBit
			if start+samplesPerBit > len(samples) {
				break
			}
			corrs[i] += watermark.CorrelateAt(samples, start, recRate)
		}
	}

	// ---------------------------------------------------------------
	// Step 3: Frame sync — try all 128 cyclic rotations.
	// The correct rotation yields a valid RS codeword.
	// ---------------------------------------------------------------
	type decodeResult struct {
		rotation int
		payload  [watermark.RsDataBytes]byte
		erasures int
	}

	var best *decodeResult

	for rot := 0; rot < watermark.PayloadBits; rot++ {
		var recv [watermark.RsTotalBytes]byte
		confs := make([]float64, watermark.PayloadBits)

		for i := 0; i < watermark.PayloadBits; i++ {
			srcBit := (i + rot) % watermark.PayloadBits
			c := corrs[srcBit]
			confs[i] = math.Abs(c)
			if c < 0 {
				recv[i/8] |= 1 << uint(7-(i%8))
			}
		}

		// Sort by confidence ascending for erasure selection
		type bitConf struct {
			idx  int
			conf float64
		}
		ranked := make([]bitConf, watermark.PayloadBits)
		for i := range ranked {
			ranked[i] = bitConf{i, confs[i]}
		}
		sort.Slice(ranked, func(a, b int) bool {
			return ranked[a].conf < ranked[b].conf
		})

		for nErase := 0; nErase <= watermark.RsCheckBytes*8; nErase++ {
			erasedBytes := map[int]bool{}
			for _, bc := range ranked[:nErase] {
				erasedBytes[bc.idx/8] = true
			}
			if len(erasedBytes) > watermark.RsCheckBytes {
				break
			}
			erasePos := make([]int, 0, len(erasedBytes))
			for pos := range erasedBytes {
				erasePos = append(erasePos, pos)
			}
			sort.Ints(erasePos)

			payload, ok := watermark.RSDecode(recv, erasePos)
			if ok {
				if best == nil || len(erasePos) < best.erasures {
					best = &decodeResult{
						rotation: rot,
						payload:  payload,
						erasures: len(erasePos),
					}
				}
				break
			}
		}

		if best != nil && best.erasures == 0 {
			break
		}
	}

	if best == nil {
		fmt.Println("\nRS decode: FAILED — no valid frame alignment found.")
		fmt.Println("Watermark may not be present, or recording is too noisy/short.")
		var maxCorr, minCorr float64
		for _, c := range corrs {
			ac := math.Abs(c)
			if ac > maxCorr {
				maxCorr = ac
			}
			if minCorr == 0 || ac < minCorr {
				minCorr = ac
			}
		}
		fmt.Printf("Correlation range: min |c|=%.4f, max |c|=%.4f\n", minCorr, maxCorr)
		os.Exit(1)
	}

	fmt.Printf("\nFrame sync: rotation=%d, RS erasures=%d\n", best.rotation, best.erasures)
	fmt.Printf("Payload:    %x\n\n", best.payload)

	keys := os.Args[2:]
	if len(keys) == 0 {
		fmt.Println("No keys supplied — payload shown above.")
		fmt.Println("Usage: wmdecode <file.wav> free [other-keys...]")
		return
	}

	fmt.Println("Key check:")
	matched := false
	for _, key := range keys {
		if watermark.KeyMatchesPayload(key, best.payload) {
			fmt.Printf("  ✓ MATCH: %q\n", key)
			matched = true
		} else {
			fmt.Printf("  ✗      : %q\n", key)
		}
	}
	if !matched {
		fmt.Println("\nNo key matched.")
	}
}

func avgCorrMag(samples []float64, phase, samplesPerBit, nBits int, recRate float64) float64 {
	var total float64
	var count int
	for b := 0; b < nBits; b++ {
		start := phase + b*samplesPerBit
		if start+samplesPerBit > len(samples) {
			break
		}
		c := watermark.CorrelateAt(samples, start, recRate)
		total += math.Abs(c)
		count++
	}
	if count == 0 {
		return 0
	}
	return total / float64(count)
}

func rmsLevel(s []float64) float64 {
	var acc float64
	for _, v := range s {
		acc += v * v
	}
	return math.Sqrt(acc / float64(len(s)))
}

func readMonoWAV(path string) ([]float64, float64, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, 0, err
	}
	if len(data) < 44 || string(data[0:4]) != "RIFF" || string(data[8:12]) != "WAVE" {
		return nil, 0, fmt.Errorf("not a RIFF/WAVE file")
	}
	var channels, bitsPerSample uint16
	var sampleRate uint32
	var dataStart, dataLen int
	i := 12
	for i+8 <= len(data) {
		id := string(data[i : i+4])
		sz := int(binary.LittleEndian.Uint32(data[i+4 : i+8]))
		i += 8
		switch id {
		case "fmt ":
			if sz >= 16 {
				channels = binary.LittleEndian.Uint16(data[i+2 : i+4])
				sampleRate = binary.LittleEndian.Uint32(data[i+4 : i+8])
				bitsPerSample = binary.LittleEndian.Uint16(data[i+14 : i+16])
			}
		case "data":
			dataStart, dataLen = i, sz
		}
		i += sz
		if sz%2 != 0 {
			i++
		}
		if dataStart > 0 && channels > 0 {
			break
		}
	}
	if dataStart == 0 || bitsPerSample != 16 || channels == 0 {
		return nil, 0, fmt.Errorf("unsupported WAV (need 16-bit PCM, got bits=%d ch=%d)", bitsPerSample, channels)
	}
	if dataStart+dataLen > len(data) {
		dataLen = len(data) - dataStart
	}
	step := int(channels) * 2
	nFrames := dataLen / step
	out := make([]float64, nFrames)
	for j := 0; j < nFrames; j++ {
		off := dataStart + j*step
		l := float64(int16(binary.LittleEndian.Uint16(data[off : off+2])))
		r := l
		if channels >= 2 {
			r = float64(int16(binary.LittleEndian.Uint16(data[off+2 : off+4])))
		}
		out[j] = (l + r) / 2.0 / 32768.0
	}
	return out, float64(sampleRate), nil
}