package dsp

import (
	"math"
	"testing"

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

// makeCompositeFrame builds a CompositeFrame with composite values in I, Q=0.
func makeCompositeFrame(values []float64, rate float64) *output.CompositeFrame {
	samples := make([]output.IQSample, len(values))
	for i, v := range values {
		samples[i] = output.IQSample{I: float32(v), Q: 0}
	}
	return &output.CompositeFrame{
		Samples:      samples,
		SampleRateHz: rate,
		Sequence:     1,
	}
}

// iqMagnitude returns sqrt(I²+Q²) for an IQ sample.
func iqMagnitude(s output.IQSample) float64 {
	return math.Sqrt(float64(s.I)*float64(s.I) + float64(s.Q)*float64(s.Q))
}

// iqPhase returns atan2(Q, I) for an IQ sample.
func iqPhase(s output.IQSample) float64 {
	return math.Atan2(float64(s.Q), float64(s.I))
}

// phaseDiff returns the shortest signed angular difference between two angles.
func phaseDiff(a, b float64) float64 {
	d := b - a
	for d > math.Pi {
		d -= 2 * math.Pi
	}
	for d < -math.Pi {
		d += 2 * math.Pi
	}
	return d
}

// --- Test: zero composite produces no frequency offset ---

func TestFMUpsampler_ZeroComposite(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)
	n := 11400 // 50ms at 228kHz
	vals := make([]float64, n)
	// All zeros — FM deviation = 0, carrier should be constant phase.
	frame := makeCompositeFrame(vals, 228000)
	out := u.Process(frame)

	if len(out.Samples) == 0 {
		t.Fatal("no output samples")
	}
	if out.SampleRateHz != 2280000 {
		t.Fatalf("expected rate 2280000, got %f", out.SampleRateHz)
	}

	// All output IQ should have magnitude ≈ 1.0 and constant phase (≈ 0)
	for i, s := range out.Samples {
		mag := iqMagnitude(s)
		if math.Abs(mag-1.0) > 0.001 {
			t.Fatalf("sample %d: magnitude %.6f, expected ~1.0", i, mag)
		}
	}

	// Phase should not advance (zero deviation)
	firstPhase := iqPhase(out.Samples[0])
	for i := 1; i < len(out.Samples); i++ {
		p := iqPhase(out.Samples[i])
		if math.Abs(phaseDiff(firstPhase, p)) > 0.01 {
			t.Fatalf("sample %d: phase drifted to %.4f (first was %.4f)", i, p, firstPhase)
		}
	}
}

// --- Test: DC composite produces constant frequency ---

func TestFMUpsampler_DCComposite(t *testing.T) {
	srcRate := 228000.0
	dstRate := 2280000.0
	maxDev := 75000.0

	u := NewFMUpsampler(srcRate, dstRate, maxDev)

	dc := 0.5 // half deviation = +37.5 kHz
	n := 11400
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = dc
	}
	frame := makeCompositeFrame(vals, srcRate)
	out := u.Process(frame)

	if len(out.Samples) < 100 {
		t.Fatalf("too few output samples: %d", len(out.Samples))
	}

	// Expected frequency: dc * maxDev = 37500 Hz
	// Expected phase step per output sample: 2π * 37500 / 2280000
	expectedStep := 2 * math.Pi * dc * maxDev / dstRate

	// Measure average phase step over a stable region (skip first 100)
	var totalDiff float64
	count := 0
	for i := 101; i < len(out.Samples) && i < 10000; i++ {
		d := phaseDiff(iqPhase(out.Samples[i-1]), iqPhase(out.Samples[i]))
		totalDiff += d
		count++
	}
	avgStep := totalDiff / float64(count)

	if math.Abs(avgStep-expectedStep) > 0.001 {
		t.Fatalf("average phase step = %.6f, expected %.6f (error %.6f)",
			avgStep, expectedStep, avgStep-expectedStep)
	}
}

// --- Test: IQ magnitude is always ≈ 1.0 ---

func TestFMUpsampler_UnitMagnitude(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	// Varying composite: a 1 kHz tone at full deviation
	n := 11400
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = math.Sin(2 * math.Pi * 1000 * float64(i) / 228000)
	}
	frame := makeCompositeFrame(vals, 228000)
	out := u.Process(frame)

	for i, s := range out.Samples {
		mag := iqMagnitude(s)
		if math.Abs(mag-1.0) > 0.002 {
			t.Fatalf("sample %d: magnitude %.6f, expected ~1.0", i, mag)
		}
	}
}

// --- Test: output sample count for exact integer ratio ---

func TestFMUpsampler_OutputCountExactRatio(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	srcLen := 11400 // 50ms
	vals := make([]float64, srcLen)
	frame := makeCompositeFrame(vals, 228000)

	// With exact 10:1 ratio, each chunk should produce exactly srcLen * 10 samples.
	// Small tolerance for boundary effects on first chunk.
	for chunk := 0; chunk < 10; chunk++ {
		out := u.Process(frame)
		expected := srcLen * 10
		if out == nil || len(out.Samples) < expected-1 || len(out.Samples) > expected+1 {
			t.Fatalf("chunk %d: got %d output samples, expected ~%d",
				chunk, len(out.Samples), expected)
		}
	}
}

// --- Test: output sample count for non-integer ratio ---

func TestFMUpsampler_OutputCountNonIntegerRatio(t *testing.T) {
	// PlutoSDR minimum: 228000 → 2084000, ratio ≈ 9.14
	u := NewFMUpsampler(228000, 2084000, 75000)

	srcLen := 11400
	vals := make([]float64, srcLen)
	frame := makeCompositeFrame(vals, 228000)

	var totalOut int
	for chunk := 0; chunk < 20; chunk++ {
		out := u.Process(frame)
		totalOut += len(out.Samples)
	}

	// Over 20 chunks, total output should be close to 20 * srcLen * ratio
	expectedTotal := int(20 * float64(srcLen) * 2084000 / 228000)
	drift := math.Abs(float64(totalOut-expectedTotal)) / float64(expectedTotal)
	if drift > 0.001 {
		t.Fatalf("total output %d, expected ~%d (drift %.4f%%)", totalOut, expectedTotal, drift*100)
	}
}

// --- CRITICAL TEST: phase continuity across chunk boundaries ---

func TestFMUpsampler_PhaseContinuityAcrossChunks(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	srcLen := 11400
	dc := 0.3 // constant deviation
	vals := make([]float64, srcLen)
	for i := range vals {
		vals[i] = dc
	}
	frame := makeCompositeFrame(vals, 228000)

	var prevLastPhase float64
	for chunk := 0; chunk < 20; chunk++ {
		out := u.Process(frame)
		n := len(out.Samples)
		if n == 0 {
			t.Fatalf("chunk %d: no output", chunk)
		}

		firstPhase := iqPhase(out.Samples[0])

		// Check continuity from previous chunk
		if chunk > 0 {
			jump := math.Abs(phaseDiff(prevLastPhase, firstPhase))
			// For DC composite at dstRate, the expected step between
			// consecutive output samples is 2π * dc * maxDev / dstRate ≈ 0.062 rad.
			// The boundary jump should be close to this, not larger.
			expectedStep := 2 * math.Pi * dc * 75000 / 2280000
			if jump > expectedStep*3 {
				t.Fatalf("chunk %d: boundary phase jump %.6f rad (expected ~%.6f)",
					chunk, jump, expectedStep)
			}
		}

		// Also check that phase increments WITHIN the chunk are smooth
		maxJump := 0.0
		for i := 1; i < n; i++ {
			d := math.Abs(phaseDiff(iqPhase(out.Samples[i-1]), iqPhase(out.Samples[i])))
			if d > maxJump {
				maxJump = d
			}
		}
		expectedIntra := 2 * math.Pi * dc * 75000 / 2280000
		if maxJump > expectedIntra*2 {
			t.Fatalf("chunk %d: intra-chunk max phase jump %.6f (expected ~%.6f)",
				chunk, maxJump, expectedIntra)
		}

		prevLastPhase = iqPhase(out.Samples[n-1])
	}
}

// --- Test: non-integer ratio boundary continuity ---

func TestFMUpsampler_BoundaryContinuityNonIntegerRatio(t *testing.T) {
	// 228000 → 2084000, ratio ≈ 9.14 — carry-over is non-zero every chunk
	u := NewFMUpsampler(228000, 2084000, 75000)

	srcLen := 11400
	dc := 0.5
	vals := make([]float64, srcLen)
	for i := range vals {
		vals[i] = dc
	}
	frame := makeCompositeFrame(vals, 228000)

	expectedStep := 2 * math.Pi * dc * 75000 / 2084000
	var prevLastPhase float64

	for chunk := 0; chunk < 50; chunk++ {
		out := u.Process(frame)
		n := len(out.Samples)
		if n == 0 {
			t.Fatalf("chunk %d: no output", chunk)
		}

		if chunk > 0 {
			jump := math.Abs(phaseDiff(prevLastPhase, iqPhase(out.Samples[0])))
			if jump > expectedStep*3 {
				t.Fatalf("chunk %d: boundary jump %.6f rad (expected ~%.6f)",
					chunk, jump, expectedStep)
			}
		}

		prevLastPhase = iqPhase(out.Samples[n-1])
	}
}

// --- Test: phase wrapping doesn't introduce discontinuity ---

func TestFMUpsampler_PhaseWrapping(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	// Full deviation for many chunks — phase will wrap many times
	srcLen := 11400
	vals := make([]float64, srcLen)
	for i := range vals {
		vals[i] = 1.0 // full positive deviation
	}
	frame := makeCompositeFrame(vals, 228000)

	expectedStep := 2 * math.Pi * 1.0 * 75000 / 2280000
	var prevLastPhase float64

	for chunk := 0; chunk < 100; chunk++ {
		out := u.Process(frame)
		n := len(out.Samples)

		if chunk > 0 {
			jump := math.Abs(phaseDiff(prevLastPhase, iqPhase(out.Samples[0])))
			if jump > expectedStep*3 {
				t.Fatalf("chunk %d: post-wrap boundary jump %.6f (limit %.6f)",
					chunk, jump, expectedStep*3)
			}
		}

		// Spot-check magnitudes (should still be 1.0 after wrapping)
		for i := 0; i < n; i += n / 10 {
			mag := iqMagnitude(out.Samples[i])
			if math.Abs(mag-1.0) > 0.002 {
				t.Fatalf("chunk %d sample %d: magnitude %.6f after wrapping", chunk, i, mag)
			}
		}

		prevLastPhase = iqPhase(out.Samples[n-1])
	}
}

// --- Test: negative composite (phase wraps negative direction) ---

func TestFMUpsampler_NegativeDeviation(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	srcLen := 11400
	vals := make([]float64, srcLen)
	for i := range vals {
		vals[i] = -0.8
	}
	frame := makeCompositeFrame(vals, 228000)

	expectedStep := 2 * math.Pi * 0.8 * 75000 / 2280000 // magnitude
	var prevLastPhase float64

	for chunk := 0; chunk < 100; chunk++ {
		out := u.Process(frame)
		n := len(out.Samples)

		if chunk > 0 {
			jump := math.Abs(phaseDiff(prevLastPhase, iqPhase(out.Samples[0])))
			if jump > expectedStep*3 {
				t.Fatalf("chunk %d: negative-dev boundary jump %.6f", chunk, jump)
			}
		}
		prevLastPhase = iqPhase(out.Samples[n-1])
	}
}

// --- Test: equivalence with direct FMModulator (frequency comparison) ---

func TestFMUpsampler_EquivalenceWithFMModulator(t *testing.T) {
	srcRate := 228000.0
	maxDev := 75000.0

	// Generate a 1kHz tone composite signal
	srcLen := 11400
	composite := make([]float64, srcLen)
	for i := range composite {
		composite[i] = 0.6 * math.Sin(2*math.Pi*1000*float64(i)/srcRate)
	}

	// Path A: FMUpsampler at 1:1 ratio (no upsampling, just FM modulation)
	up := NewFMUpsampler(srcRate, srcRate, maxDev)
	frame := makeCompositeFrame(composite, srcRate)
	outA := up.Process(frame)

	// Path B: Direct FMModulator
	mod := NewFMModulator(srcRate)
	mod.MaxDeviation = maxDev
	outB := make([]output.IQSample, srcLen)
	for i, c := range composite {
		oi, oq := mod.Modulate(c)
		outB[i] = output.IQSample{I: float32(oi), Q: float32(oq)}
	}

	// The upsampler has a 1-sample offset due to the virtual prevPhase index.
	// Instead of comparing absolute phase, compare instantaneous frequency
	// (phase step per sample), which is invariant to time shifts.

	freqA := make([]float64, 0, len(outA.Samples)-1)
	for i := 1; i < len(outA.Samples); i++ {
		freqA = append(freqA, phaseDiff(iqPhase(outA.Samples[i-1]), iqPhase(outA.Samples[i])))
	}

	freqB := make([]float64, 0, len(outB)-1)
	for i := 1; i < len(outB); i++ {
		freqB = append(freqB, phaseDiff(iqPhase(outB[i-1]), iqPhase(outB[i])))
	}

	// Compare frequency profiles (skip edges, allow 1-sample alignment offset)
	minLen := len(freqA)
	if len(freqB) < minLen {
		minLen = len(freqB)
	}

	maxFreqDiff := 0.0
	for i := 20; i < minLen-20; i++ {
		// Try both aligned and 1-sample-shifted
		d0 := math.Abs(freqA[i] - freqB[i])
		d1 := math.Abs(freqA[i] - freqB[i-1])
		d := math.Min(d0, d1)
		if d > maxFreqDiff {
			maxFreqDiff = d
		}
	}

	// Instantaneous frequency should match within 0.02 rad/sample
	if maxFreqDiff > 0.02 {
		t.Fatalf("max instantaneous frequency difference: %.6f rad/sample (too large)", maxFreqDiff)
	}

	// All magnitudes should be ~1.0
	maxMagDiff := 0.0
	for i := range outA.Samples {
		mag := iqMagnitude(outA.Samples[i])
		d := math.Abs(mag - 1.0)
		if d > maxMagDiff {
			maxMagDiff = d
		}
	}
	if maxMagDiff > 0.01 {
		t.Fatalf("max magnitude deviation: %.6f", maxMagDiff)
	}

	t.Logf("equivalence: maxFreqDiff=%.6f rad/sample, maxMagDiff=%.6f", maxFreqDiff, maxMagDiff)
}

// --- Test: buffer reuse (no allocation after first call) ---

func TestFMUpsampler_BufferReuse(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	srcLen := 11400
	vals := make([]float64, srcLen)
	frame := makeCompositeFrame(vals, 228000)

	// First call allocates
	out1 := u.Process(frame)
	ptr1 := &out1.Samples[0]

	// Second call should reuse the same buffer
	out2 := u.Process(frame)
	ptr2 := &out2.Samples[0]

	if ptr1 != ptr2 {
		t.Fatal("buffer was reallocated on second call — should reuse")
	}
}

// --- Test: Reset clears state properly ---

func TestFMUpsampler_Reset(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	vals := make([]float64, 11400)
	for i := range vals {
		vals[i] = 0.5
	}
	frame := makeCompositeFrame(vals, 228000)

	// Run a few chunks to accumulate state
	for i := 0; i < 5; i++ {
		u.Process(frame)
	}

	phase1, prev1, pos1 := u.Stats()
	if phase1 == 0 && prev1 == 0 && pos1 == 0 {
		t.Fatal("state should be non-zero after processing")
	}

	u.Reset()
	phase2, prev2, pos2 := u.Stats()
	if phase2 != 0 || prev2 != 0 || pos2 != 0 {
		t.Fatalf("state not zeroed after Reset: phase=%f prev=%f pos=%f", phase2, prev2, pos2)
	}
}

// --- Test: tiny chunks (edge case: 1 source sample) ---

func TestFMUpsampler_TinyChunk(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	vals := []float64{0.5}
	frame := makeCompositeFrame(vals, 228000)

	for chunk := 0; chunk < 20; chunk++ {
		out := u.Process(frame)
		if len(out.Samples) == 0 {
			t.Fatalf("chunk %d: no output from single-sample input", chunk)
		}
		for i, s := range out.Samples {
			mag := iqMagnitude(s)
			if math.Abs(mag-1.0) > 0.01 {
				t.Fatalf("chunk %d sample %d: magnitude %.4f", chunk, i, mag)
			}
		}
	}
}

// --- Test: alternating sign composite (stress boundary interpolation) ---

func TestFMUpsampler_AlternatingComposite(t *testing.T) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	srcLen := 11400
	vals := make([]float64, srcLen)
	for i := range vals {
		if i%2 == 0 {
			vals[i] = 0.8
		} else {
			vals[i] = -0.8
		}
	}
	frame := makeCompositeFrame(vals, 228000)

	var prevLastPhase float64
	for chunk := 0; chunk < 20; chunk++ {
		out := u.Process(frame)
		n := len(out.Samples)

		if chunk > 0 {
			jump := math.Abs(phaseDiff(prevLastPhase, iqPhase(out.Samples[0])))
			// Alternating composite: the phase meanders. Boundary jump
			// should still be bounded by the maximum single-sample phase change.
			maxSingleStep := 2 * math.Pi * 0.8 * 75000 / 228000.0 // at source rate
			maxOutputJump := maxSingleStep * 0.1                    // step = srcRate/dstRate
			if jump > maxOutputJump*3 {
				t.Fatalf("chunk %d: alternating boundary jump %.4f (limit %.4f)",
					chunk, jump, maxOutputJump*3)
			}
		}

		// Check all magnitudes
		for i, s := range out.Samples {
			mag := iqMagnitude(s)
			if math.Abs(mag-1.0) > 0.01 {
				t.Fatalf("chunk %d sample %d: magnitude %.4f", chunk, i, mag)
			}
		}

		prevLastPhase = iqPhase(out.Samples[n-1])
	}
}

// --- Test: long-running stability (simulates 1 hour) ---

func TestFMUpsampler_LongRun(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping long-run test in short mode")
	}

	u := NewFMUpsampler(228000, 2280000, 75000)

	srcLen := 11400
	vals := make([]float64, srcLen)
	for i := range vals {
		vals[i] = 0.3 * math.Sin(2*math.Pi*1000*float64(i)/228000)
	}
	frame := makeCompositeFrame(vals, 228000)

	// 20 chunks/sec × 3600 sec = 72000 chunks per hour.
	// We test 72000 chunks (simulating 1 hour).
	chunks := 72000
	var prevLastPhase float64

	for chunk := 0; chunk < chunks; chunk++ {
		out := u.Process(frame)
		n := len(out.Samples)

		if n == 0 {
			t.Fatalf("chunk %d: no output", chunk)
		}

		// Check boundary every 1000 chunks
		if chunk > 0 && chunk%1000 == 0 {
			jump := math.Abs(phaseDiff(prevLastPhase, iqPhase(out.Samples[0])))
			if jump > 0.5 {
				t.Fatalf("chunk %d: boundary jump %.4f rad (phase drift?)", chunk, jump)
			}

			// Check magnitude of first sample
			mag := iqMagnitude(out.Samples[0])
			if math.Abs(mag-1.0) > 0.01 {
				t.Fatalf("chunk %d: magnitude %.4f after long run", chunk, mag)
			}
		}

		prevLastPhase = iqPhase(out.Samples[n-1])
	}

	// Check phase is still bounded
	phase, _, _ := u.Stats()
	if math.Abs(phase) > 2*math.Pi {
		t.Fatalf("phase unbounded after %d chunks: %.2f", chunks, phase)
	}
}

// --- Benchmark ---

func BenchmarkFMUpsampler_Process(b *testing.B) {
	u := NewFMUpsampler(228000, 2280000, 75000)

	srcLen := 11400
	vals := make([]float64, srcLen)
	for i := range vals {
		vals[i] = 0.5 * math.Sin(2*math.Pi*1000*float64(i)/228000)
	}
	frame := makeCompositeFrame(vals, 228000)

	// Warm up (trigger buffer allocation)
	u.Process(frame)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		u.Process(frame)
	}
}