package watermark

import (
	"math"
	"sort"
	"testing"
)

// TestRoundTrip verifies the full embed → downsample → phase-search → rotation → RS-decode chain.
func TestRoundTrip(t *testing.T) {
	const key = "test-key-42"
	const recRate = float64(RecordingRate) // 48000
	const compRate = float64(CompositeRate) // 228000
	const duration = 15.0 // seconds — should give ~2 full frames

	nRecSamples := int(duration * recRate)

	// === Embed ===
	emb := NewEmbedder(key)
	// No gate — test pure watermark signal
	samples := make([]float64, 0, nRecSamples)

	// Drive embedder at CompositeRate, collect at RecordingRate via Bresenham
	accum := 0
	var last float64
	for len(samples) < nRecSamples {
		last = emb.NextSample()
		accum += RecordingRate
		if accum >= CompositeRate {
			accum -= CompositeRate
			samples = append(samples, last)
		}
	}

	t.Logf("Embedded: %d samples @ %.0f Hz = %.2fs", len(samples), recRate, float64(len(samples))/recRate)

	// RMS check
	var rmsAcc float64
	for _, s := range samples {
		rmsAcc += s * s
	}
	rms := math.Sqrt(rmsAcc / float64(len(samples)))
	rmsDBFS := 20 * math.Log10(rms+1e-12)
	t.Logf("Watermark RMS: %.1f dBFS (expect ~-48)", rmsDBFS)
	if rmsDBFS < -52 || rmsDBFS > -44 {
		t.Errorf("RMS %.1f dBFS out of expected range [-52, -44]", rmsDBFS)
	}

	// === Decode: Phase search ===
	samplesPerBit := int(float64(PnChips) * recRate / float64(RecordingRate))
	t.Logf("samplesPerBit=%d, frameLen=%d", samplesPerBit, samplesPerBit*PayloadBits)

	const coarseStep = 8
	const syncBits = 64

	bestPhase := 0
	bestMag := 0.0
	for phase := 0; phase < samplesPerBit; phase += coarseStep {
		mag := testAvgCorrMag(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 := testAvgCorrMag(samples, phase, samplesPerBit, syncBits, recRate)
		if mag > bestMag {
			bestMag = mag
			bestPhase = phase
		}
	}
	t.Logf("Phase search: bestPhase=%d, avgCorr=%.4f", bestPhase, bestMag)

	// Phase should be 0 for clean signal starting at sample 0
	if bestPhase != 0 {
		t.Errorf("expected bestPhase=0, got %d", bestPhase)
	}

	// === Decode: Extract correlations ===
	nCompleteBits := (len(samples) - bestPhase) / samplesPerBit
	nFrames := nCompleteBits / PayloadBits
	if nFrames == 0 { nFrames = 1 }
	t.Logf("Complete bits: %d, frames: %d", nCompleteBits, nFrames)

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

	// Log correlation stats
	var minAbs, maxAbs float64
	for i, c := range corrs {
		ac := math.Abs(c)
		if i == 0 || ac < minAbs { minAbs = ac }
		if ac > maxAbs { maxAbs = ac }
	}
	t.Logf("Correlation range: min|c|=%.2f, max|c|=%.2f", minAbs, maxAbs)

	// === Decode: Frame sync via rotation ===
	type decodeResult struct {
		rotation int
		payload  [RsDataBytes]byte
		erasures int
	}
	var best *decodeResult

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

		type bitConf struct { idx int; conf float64 }
		ranked := make([]bitConf, 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 <= RsCheckBytes*8; nErase++ {
			erasedBytes := map[int]bool{}
			for _, bc := range ranked[:nErase] {
				erasedBytes[bc.idx/8] = true
			}
			if len(erasedBytes) > RsCheckBytes { break }
			erasePos := make([]int, 0, len(erasedBytes))
			for pos := range erasedBytes { erasePos = append(erasePos, pos) }
			sort.Ints(erasePos)

			payload, ok := 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 {
		t.Fatal("RS decode FAILED — no valid rotation found")
	}

	t.Logf("Decoded: rotation=%d, erasures=%d, payload=%x", best.rotation, best.erasures, best.payload)

	// Rotation should be 0 for clean signal
	if best.rotation != 0 {
		t.Errorf("expected rotation=0, got %d", best.rotation)
	}
	if best.erasures != 0 {
		t.Errorf("expected 0 erasures, got %d", best.erasures)
	}

	// Key match
	if !KeyMatchesPayload(key, best.payload) {
		t.Errorf("key %q does NOT match decoded payload %x", key, best.payload)
	} else {
		t.Logf("Key %q MATCHES", key)
	}
}

func testAvgCorrMag(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 := CorrelateAt(samples, start, recRate)
		total += math.Abs(c)
		count++
	}
	if count == 0 { return 0 }
	return total / float64(count)
}