package offline

import (
	"math"
	"testing"
	"time"

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

// TestWatermarkE2EFloat32 tests the FULL path including float32 storage
// (as happens in IQSample.I) and FMUpsampler FM modulation + demodulation.
func TestWatermarkE2EFloat32(t *testing.T) {
	const key = "test-key-e2e-f32"
	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 mode
	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)

	frame := gen.GenerateFrame(duration)
	nSamples := len(frame.Samples)
	compositeRate := frame.SampleRateHz
	t.Logf("Generated %d samples @ %.0f Hz", nSamples, compositeRate)

	// Test 1: float32 truncation
	t.Run("float32_storage", func(t *testing.T) {
		// Simulate what IQSample does: float32(composite)
		composite := make([]float64, nSamples)
		for i, s := range frame.Samples {
			composite[i] = float64(s.I) // s.I is float32
		}
		testDecode(t, composite, compositeRate, key)
	})

	// Test 2: FM modulate + demodulate
	t.Run("fm_mod_demod", func(t *testing.T) {
		maxDev := 75000.0
		// FM modulate (same as FMUpsampler phase accumulation)
		phases := make([]float64, nSamples)
		phaseInc := 2 * math.Pi * maxDev / compositeRate
		phase := 0.0
		for i, s := range frame.Samples {
			phase += float64(s.I) * phaseInc
			phases[i] = phase
		}
		// FM demodulate: instantaneous frequency = dphase/dt
		demod := make([]float64, nSamples)
		for i := 1; i < nSamples; i++ {
			dp := phases[i] - phases[i-1]
			demod[i] = dp / phaseInc // recover composite
		}
		demod[0] = demod[1]
		testDecode(t, demod, compositeRate, key)
	})

	// Test 3: FM mod + upsample 10× + downsample + demod (full SDR path)
	t.Run("fm_upsample_downsample", func(t *testing.T) {
		maxDev := 75000.0
		deviceRate := 2280000.0
		upsampler := dsp.NewFMUpsampler(compositeRate, deviceRate, maxDev)
		upFrame := upsampler.Process(frame)
		t.Logf("Upsampled: %d IQ samples @ %.0f Hz", len(upFrame.Samples), upFrame.SampleRateHz)

		// FM demodulate the IQ output: phase = atan2(Q, I), freq = dphase/dt
		nUp := len(upFrame.Samples)
		demod := make([]float64, nUp)
		prevPhase := 0.0
		for i, s := range upFrame.Samples {
			p := math.Atan2(float64(s.Q), float64(s.I))
			dp := p - prevPhase
			// Unwrap
			for dp > math.Pi { dp -= 2 * math.Pi }
			for dp < -math.Pi { dp += 2 * math.Pi }
			demod[i] = dp * deviceRate / (2 * math.Pi * maxDev)
			prevPhase = p
		}

		// Downsample back to composite rate
		ratio := int(deviceRate / compositeRate)
		nDown := nUp / ratio
		downsampled := make([]float64, nDown)
		for i := 0; i < nDown; i++ {
			// Simple decimate (use average over ratio samples)
			sum := 0.0
			for j := 0; j < ratio; j++ {
				idx := i*ratio + j
				if idx < nUp { sum += demod[idx] }
			}
			downsampled[i] = sum / float64(ratio)
		}
		t.Logf("Downsampled: %d samples @ %.0f Hz", nDown, compositeRate)
		testDecode(t, downsampled, compositeRate, key)
	})
}

func testDecode(t *testing.T, composite []float64, rate float64, key string) {
	t.Helper()
	chipRate := float64(watermark.ChipRate)
	samplesPerBit := int(float64(watermark.PnChips) * rate / chipRate)
	nSamples := len(composite)

	// Phase search
	bestPhase := 0
	bestEnergy := 0.0
	step := max(1, samplesPerBit/200)
	for phase := 0; phase < samplesPerBit; phase += step {
		var energy float64
		for b := 0; b < min(100, nSamples/samplesPerBit); b++ {
			start := phase + b*samplesPerBit
			if start+samplesPerBit > nSamples { break }
			c := watermark.CorrelateAt(composite, start, rate)
			energy += c * c
		}
		if energy > bestEnergy {
			bestEnergy = energy; bestPhase = phase
		}
	}

	// Correlate
	nComplete := (nSamples - bestPhase) / samplesPerBit
	nFrames := nComplete / 128
	if nFrames < 1 { nFrames = 1 }
	corrs := make([]float64, 128)
	for i := 0; i < 128; i++ {
		for f := 0; f < nFrames; f++ {
			start := bestPhase + (f*128+i)*samplesPerBit
			if start+samplesPerBit > nSamples { break }
			corrs[i] += watermark.CorrelateAt(composite, start, rate)
		}
	}

	var nStrong, nDead int
	var sumAbs float64
	for _, c := range corrs {
		ac := math.Abs(c)
		sumAbs += ac
		if ac > 50 { nStrong++ }
		if ac < 5 { nDead++ }
	}
	t.Logf("phase=%d, frames=%d, avg|c|=%.1f, strong=%d, dead=%d",
		bestPhase, nFrames, sumAbs/128, nStrong, nDead)

	if nStrong < 100 {
		t.Errorf("Only %d/128 strong bits — watermark degraded", nStrong)
	}
}