fm-rds-tx/internal/offline/watermark_e2e_test.go

package offline

import (
	"math"
	"sort"
	"testing"
	"time"

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

// TestWatermarkE2E runs a FULL generator pipeline (audio source → pre-emphasis →
// LPF → drive → clip → cleanup → watermark injection → stereo encode → composite
// clip → notch → pilot → RDS → FM mod) and then tries to decode the watermark
// from the composite output. This tests the real code path, not just the embedder.
func TestWatermarkE2E(t *testing.T) {
	const key = "test-key-e2e"
	const duration = 45 * time.Second

	cfg := cfgpkg.Default()
	cfg.FM.CompositeRateHz = 228000
	cfg.FM.StereoEnabled = true
	cfg.FM.OutputDrive = 0.5
	cfg.FM.LimiterEnabled = true
	cfg.FM.LimiterCeiling = 1.0
	cfg.FM.FMModulationEnabled = false // split-rate: composite output, no IQ
	cfg.Audio.ToneLeftHz = 1000
	cfg.Audio.ToneRightHz = 1600
	cfg.Audio.ToneAmplitude = 0.4
	cfg.Audio.Gain = 1.0
	cfg.FM.PreEmphasisTauUS = 50

	gen := NewGenerator(cfg)
	licState := license.NewState("")
	gen.SetLicense(licState, key)

	// Generate composite
	frame := gen.GenerateFrame(duration)
	if frame == nil {
		t.Fatal("GenerateFrame returned nil")
	}
	t.Logf("Generated %d composite samples @ %.0f Hz (%.2fs)",
		len(frame.Samples), frame.SampleRateHz, float64(len(frame.Samples))/frame.SampleRateHz)

	// Extract mono composite (I channel = composite baseband in non-FM mode)
	compositeRate := frame.SampleRateHz
	nSamples := len(frame.Samples)
	composite := make([]float64, nSamples)
	for i, s := range frame.Samples {
		composite[i] = float64(s.I)
	}

	// RMS
	var rmsAcc float64
	for _, s := range composite {
		rmsAcc += s * s
	}
	rms := math.Sqrt(rmsAcc / float64(nSamples))
	t.Logf("Composite RMS: %.1f dBFS", 20*math.Log10(rms+1e-12))

	// Now decode the watermark from the composite
	chipRate := float64(watermark.ChipRate)
	samplesPerBit := int(float64(watermark.PnChips) * compositeRate / chipRate)
	frameLen := samplesPerBit * watermark.PayloadBits
	nFrames := nSamples / frameLen
	t.Logf("Decode: samplesPerBit=%d, frameLen=%d, nFrames=%d", samplesPerBit, frameLen, nFrames)

	if nFrames < 1 {
		t.Fatalf("Need at least 1 frame (%d samples), have %d", frameLen, nSamples)
	}

	// Phase search (should be 0 since we start from sample 0)
	bestPhase := 0
	bestEnergy := 0.0
	step := max(1, samplesPerBit/500)
	for phase := 0; phase < samplesPerBit; phase += step {
		var energy float64
		nBits := min(200, (nSamples-phase)/samplesPerBit)
		for b := 0; b < nBits; b++ {
			start := phase + b*samplesPerBit
			if start+samplesPerBit > nSamples { break }
			c := watermark.CorrelateAt(composite, start, compositeRate)
			energy += c * c
		}
		if energy > bestEnergy {
			bestEnergy = energy
			bestPhase = phase
		}
	}
	t.Logf("Phase: %d (energy=%.1f)", bestPhase, bestEnergy)

	// Correlate all 128 bits with frame averaging
	nCompleteBits := (nSamples - bestPhase) / samplesPerBit
	nAvgFrames := nCompleteBits / watermark.PayloadBits
	if nAvgFrames < 1 { nAvgFrames = 1 }

	corrs := make([]float64, watermark.PayloadBits)
	for i := 0; i < watermark.PayloadBits; i++ {
		for f := 0; f < nAvgFrames; f++ {
			bitGlobal := f*watermark.PayloadBits + i
			start := bestPhase + bitGlobal*samplesPerBit
			if start+samplesPerBit > nSamples { break }
			corrs[i] += watermark.CorrelateAt(composite, start, compositeRate)
		}
	}

	var minC, maxC, sumC float64
	var nStrong, nDead int
	for i, c := range corrs {
		ac := math.Abs(c)
		sumC += ac
		if i == 0 || ac < minC { minC = ac }
		if ac > maxC { maxC = ac }
		if ac > 50 { nStrong++ }
		if ac < 5 { nDead++ }
	}
	t.Logf("Correlations: min=%.1f max=%.1f avg=%.1f strong=%d dead=%d",
		minC, maxC, sumC/128, nStrong, nDead)

	if nStrong < 64 {
		t.Errorf("Too few strong bits: %d/128 (expected >64)", nStrong)
	}

	// Frame sync: try all 128 rotations with byte-level erasure
	type decResult struct {
		rot     int
		payload [watermark.RsDataBytes]byte
		ok      bool
	}
	var bestDec *decResult

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

		// Try with 0..8 byte erasures
		type bc struct{ idx int; conf float64 }
		ranked := make([]bc, watermark.RsTotalBytes)
		for i := range ranked { ranked[i] = bc{i, byteConfs[i]} }
		sort.Slice(ranked, func(a, b int) bool { return ranked[a].conf < ranked[b].conf })

		for ne := 0; ne <= watermark.RsCheckBytes; ne++ {
			erasePos := make([]int, ne)
			for i := 0; i < ne; i++ { erasePos[i] = ranked[i].idx }
			sort.Ints(erasePos)
			payload, ok := watermark.RSDecode(recv, erasePos)
			if ok {
				bestDec = &decResult{rot, payload, true}
				break
			}
		}
		if bestDec != nil { break }
	}

	if bestDec == nil {
		t.Fatal("RS decode FAILED — watermark not recoverable from generator output")
	}

	t.Logf("Decoded: rotation=%d, payload=%x", bestDec.rot, bestDec.payload)
	if !watermark.KeyMatchesPayload(key, bestDec.payload) {
		t.Errorf("Key mismatch: %q does not match payload %x", key, bestDec.payload)
	} else {
		t.Logf("Key %q MATCHES ✓", key)
	}
}

func min(a, b int) int { if a < b { return a }; return b }
func max(a, b int) int { if a > b { return a }; return b }