// cmd/wmdecode — fm-rds-tx spread-spectrum watermark recovery tool.
//
// Approach: downsample to chip rate (12 kHz), correlate at 1 sample/chip.
// No fractional stepping, no clock drift issues. FFT-free phase search.
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))

	chipRate := float64(watermark.ChipRate) // 12000
	pnChips := watermark.PnChips            // 2048

	// Step 1: LPF at ChipRate/2 then downsample to ChipRate.
	// At chip rate: 1 sample = 1 chip. No fractional stepping.
	decimFactor := int(recRate / chipRate) // 192000/12000 = 16
	if decimFactor < 1 {
		decimFactor = 1
	}
	actualChipRate := recRate / float64(decimFactor) // should be exactly chipRate

	fmt.Printf("Downsample: %d:1 (%.0f Hz → %.0f Hz)\n", decimFactor, recRate, actualChipRate)

	// Anti-alias LPF (8th-order IIR at 5.5 kHz)
	lpfCoeffs := designLPF8(5500, recRate)
	filtered := applyIIR(samples, lpfCoeffs)

	// Decimate
	nDown := len(filtered) / decimFactor
	down := make([]float64, nDown)
	for i := 0; i < nDown; i++ {
		down[i] = filtered[i*decimFactor]
	}
	rmsDown := rmsLevel(down)
	fmt.Printf("Downsampled: %d samples @ %.0f Hz, RMS %.1f dBFS\n",
		nDown, actualChipRate, 20*math.Log10(rmsDown+1e-9))

	// Step 2: Phase search — slide 1-bit PN template across [0, pnChips).
	// At chip rate this is a simple 2048-element dot product per offset.
	// Test all 2048 phases, accumulate energy over many bits.
	fmt.Printf("Phase search: %d candidates\n", pnChips)

	nSearchBits := nDown / pnChips
	if nSearchBits > 500 {
		nSearchBits = 500
	}
	bestPhase := 0
	bestEnergy := 0.0
	for phase := 0; phase < pnChips; phase++ {
		var energy float64
		for b := 0; b < nSearchBits; b++ {
			start := phase + b*pnChips
			if start+pnChips > nDown {
				break
			}
			c := corrChipRate(down, start, pnChips)
			energy += c * c
		}
		if energy > bestEnergy {
			bestEnergy = energy
			bestPhase = phase
		}
	}
	fmt.Printf("Phase: offset=%d (%.2fms), energy=%.0f\n",
		bestPhase, float64(bestPhase)/actualChipRate*1000, bestEnergy)

	// Step 3: Per-bit correlation with ±4 sample sliding (handles residual drift).
	// At chip rate, ±4 samples = ±0.33ms — covers ~±40 ppm over 22s frame.
	nCompleteBits := (nDown - bestPhase) / pnChips
	nFrames := nCompleteBits / watermark.PayloadBits
	if nFrames < 1 {
		nFrames = 1
	}
	fmt.Printf("Sync:  %d complete bits, %d frames\n", nCompleteBits, nFrames)

	const slideWindow = 200 // ±200 chips — handles phase errors + drift

	corrs := make([]float64, watermark.PayloadBits)
	for i := 0; i < watermark.PayloadBits; i++ {
		// For each offset, sum correlation across ALL frames first.
		// Signal adds coherently (×nFrames), noise adds as √nFrames.
		// Then pick the offset with maximum |sum|.
		bestAbs := 0.0
		bestVal := 0.0
		for off := -slideWindow; off <= slideWindow; off++ {
			var sum float64
			for f := 0; f < nFrames; f++ {
				nominal := bestPhase + (f*watermark.PayloadBits+i)*pnChips + off
				if nominal < 0 || nominal+pnChips > nDown {
					continue
				}
				sum += corrChipRate(down, nominal, pnChips)
			}
			if math.Abs(sum) > bestAbs {
				bestAbs = math.Abs(sum)
				bestVal = sum
			}
		}
		corrs[i] = bestVal
	}

	// Diagnostics
	var corrMin, corrMax, sumAbs float64
	var nStrong, nDead int
	for i, c := range corrs {
		ac := math.Abs(c)
		sumAbs += ac
		if i == 0 || ac < corrMin {
			corrMin = ac
		}
		if ac > corrMax {
			corrMax = ac
		}
		if ac > sumAbs/float64(i+1)*2 {
			nStrong++
		}
		if ac < 3 {
			nDead++
		}
	}
	avgCorr := sumAbs / 128
	nStrong = 0
	for _, c := range corrs {
		if math.Abs(c) > avgCorr*0.5 {
			nStrong++
		}
	}
	fmt.Printf("Corrs: min|c|=%.1f, max|c|=%.1f, avg|c|=%.1f (strong=%d, dead=%d)\n",
		corrMin, corrMax, avgCorr, nStrong, nDead)

	// Step 4: Frame sync — 128 rotations × byte-level erasure + bit-flipping.

	// Verbose: compute BER at each rotation against the known key (if supplied)
	knownPayload := [watermark.RsDataBytes]byte{}
	hasKnown := false
	if len(os.Args) >= 3 {
		hasKnown = true
		knownPayload = watermark.KeyToPayload(os.Args[2])
		knownCW := watermark.RSEncode(knownPayload)
		var knownBits [watermark.PayloadBits]int
		for i := 0; i < watermark.PayloadBits; i++ {
			knownBits[i] = int((knownCW[i/8] >> uint(7-(i%8))) & 1)
		}
		fmt.Println("\nRotation sweep (top 10 by BER):")
		type rotBER struct{ rot, ber int }
		var results []rotBER
		for rot := 0; rot < watermark.PayloadBits; rot++ {
			nerr := 0
			for i := 0; i < watermark.PayloadBits; i++ {
				srcBit := (i + rot) % watermark.PayloadBits
				hard := 0
				if corrs[srcBit] < 0 {
					hard = 1
				}
				if hard != knownBits[i] {
					nerr++
				}
			}
			results = append(results, rotBER{rot, nerr})
		}
		sort.Slice(results, func(a, b int) bool { return results[a].ber < results[b].ber })
		for j := 0; j < 10 && j < len(results); j++ {
			r := results[j]
			fmt.Printf("  rot=%3d: BER=%d/128 (%4.1f%%)\n", r.rot, r.ber, 100*float64(r.ber)/128)
		}

		// Show byte error pattern at best rotation
		bestRot := results[0].rot
		fmt.Printf("\nByte errors at rot=%d:\n  ", bestRot)
		for b := 0; b < watermark.RsTotalBytes; b++ {
			nerr := 0
			for bit := 0; bit < 8; bit++ {
				srcBit := (b*8 + bit + bestRot) % watermark.PayloadBits
				hard := 0
				if corrs[srcBit] < 0 {
					hard = 1
				}
				if hard != knownBits[b*8+bit] {
					nerr++
				}
			}
			fmt.Printf("B%d:%d ", b, nerr)
		}
		fmt.Println()

		// Show received vs expected codeword at best rotation
		var recv [watermark.RsTotalBytes]byte
		for i := 0; i < watermark.PayloadBits; i++ {
			srcBit := (i + bestRot) % watermark.PayloadBits
			if corrs[srcBit] < 0 {
				recv[i/8] |= 1 << uint(7-(i%8))
			}
		}
		fmt.Printf("  recv: %x\n", recv)
		fmt.Printf("  want: %x\n", knownCW)
	}
	_ = hasKnown
	_ = knownPayload
	type decodeResult struct {
		rotation int
		payload  [watermark.RsDataBytes]byte
		flips    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))
			}
		}

		// Brute-force RS decode: try ALL possible erasure subsets of size 1..8.
		// With sliding correlation, confidence values are unreliable for erasure
		// selection (all bits look "strong"). Instead, let RS tell us which
		// subsets produce a valid codeword. This is fast: sum(C(16,k), k=1..8)
		// = ~39k RS decodes per rotation, ~5M total. Each takes <1µs.
		decoded := false
		for nErase := 1; nErase <= watermark.RsCheckBytes; nErase++ {
			if decoded { break }
			indices := make([]int, nErase)
			for i := range indices { indices[i] = i }
			for {
				erasePos := make([]int, nErase)
				copy(erasePos, indices)
				payload, ok := watermark.RSDecode(recv, erasePos)
				if ok {
					if best == nil {
						best = &decodeResult{rot, payload, nErase}
					}
					decoded = true
					break
				}
				// Next combination
				i := nErase - 1
				for i >= 0 && indices[i] == watermark.RsTotalBytes-nErase+i {
					i--
				}
				if i < 0 { break }
				indices[i]++
				for j := i + 1; j < nErase; j++ {
					indices[j] = indices[j-1] + 1
				}
			}
		}
		if decoded && best != nil && best.flips <= 4 {
			break // clean decode with few erasures — stop early
		}
	}

	if best == nil {
		fmt.Println("RS decode: FAILED — no valid frame alignment found.")
		fmt.Println("Watermark may not be present, or recording is too noisy/short.")
		os.Exit(1)
	}

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

	keys := os.Args[2:]
	if len(keys) == 0 {
		fmt.Println("No keys supplied — payload shown above.")
		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.")
	}
}

// corrChipRate correlates at chip rate (1 sample = 1 chip).
func corrChipRate(down []float64, start, pnChips int) float64 {
	var acc float64
	for i := 0; i < pnChips; i++ {
		acc += down[start+i] * float64(watermark.PNSequence[i])
	}
	return acc
}

// --- 8th-order Butterworth LPF (4 cascaded biquads) ---
type biquad struct{ b0, b1, b2, a1, a2 float64 }
type iirCoeffs []biquad

func designLPF8(cutoffHz, sampleRate float64) iirCoeffs {
	// 8th-order Butterworth = 4 biquad sections
	angles := []float64{math.Pi / 16, 3 * math.Pi / 16, 5 * math.Pi / 16, 7 * math.Pi / 16}
	coeffs := make(iirCoeffs, 4)
	for i, angle := range angles {
		q := 1.0 / (2 * math.Cos(angle))
		omega := 2 * math.Pi * cutoffHz / sampleRate
		cosW := math.Cos(omega)
		sinW := math.Sin(omega)
		alpha := sinW / (2 * q)
		a0 := 1 + alpha
		coeffs[i] = biquad{
			b0: (1 - cosW) / 2 / a0,
			b1: (1 - cosW) / a0,
			b2: (1 - cosW) / 2 / a0,
			a1: (-2 * cosW) / a0,
			a2: (1 - alpha) / a0,
		}
	}
	return coeffs
}

func applyIIR(samples []float64, coeffs iirCoeffs) []float64 {
	out := make([]float64, len(samples))
	copy(out, samples)
	for _, bq := range coeffs {
		var z1, z2 float64
		for i, x := range out {
			y := bq.b0*x + z1
			z1 = bq.b1*x - bq.a1*y + z2
			z2 = bq.b2*x - bq.a2*y
			out[i] = y
		}
	}
	return out
}

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
}