package offline

import (
	"math"
	"testing"
	"time"

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

func TestWatermarkE2EFloat32(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping long watermark E2E float32 test in -short mode")
	}
	const key = "test-key-e2e-f32"
	const duration = 45 * time.Second

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

	t.Run("float32_storage", func(t *testing.T) {
		composite := extractCompositeFrame(frame)
		payload, _, ok := decodeWatermarkFromComposite(t, composite, compositeRate)
		if !ok {
			t.Fatal("decode failed after float32 composite storage")
		}
		if !watermark.KeyMatchesPayload(key, payload) {
			t.Fatalf("payload mismatch after float32 storage: %x", payload)
		}
	})

	t.Run("fm_mod_demod", func(t *testing.T) {
		maxDev := 75000.0
		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
		}
		demod := make([]float64, nSamples)
		for i := 1; i < nSamples; i++ {
			dp := phases[i] - phases[i-1]
			demod[i] = dp / phaseInc
		}
		demod[0] = demod[1]
		payload, _, ok := decodeWatermarkFromComposite(t, demod, compositeRate)
		if !ok {
			t.Fatal("decode failed after FM mod/demod")
		}
		if !watermark.KeyMatchesPayload(key, payload) {
			t.Fatalf("payload mismatch after FM mod/demod: %x", payload)
		}
	})

	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)

		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
			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
		}

		ratio := int(deviceRate / compositeRate)
		nDown := nUp / ratio
		downsampled := make([]float64, nDown)
		for i := 0; i < nDown; i++ {
			sum := 0.0
			for j := 0; j < ratio; j++ {
				idx := i*ratio + j
				if idx < nUp {
					sum += demod[idx]
				}
			}
			downsampled[i] = sum / float64(ratio)
		}
		payload, _, ok := decodeWatermarkFromComposite(t, downsampled, compositeRate)
		if !ok {
			t.Fatal("decode failed after FM upsample/downsample path")
		}
		if !watermark.KeyMatchesPayload(key, payload) {
			t.Fatalf("payload mismatch after FM upsample/downsample path: %x", payload)
		}
	})
}
