package dsp

import (
	"math"

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

// FMUpsampler converts a composite baseband signal at a low source rate into
// FM-modulated IQ samples at a higher device rate, via phase-domain
// interpolation.
//
// Architecture: accumulate FM phase at source rate (cheap, few trig ops),
// then linearly interpolate the phase to device rate and emit sin/cos.
// This is mathematically equivalent to running the full FMModulator at device
// rate, but needs trig only at the output rate — saving all the DSP that
// would otherwise run at the higher rate (stereo encode, RDS, limiter etc.).
//
// Cross-chunk boundary: the upsampler carries over prevPhase and srcPos
// between calls. The interpolation coordinate system places prevPhase at
// virtual index 0 and srcPhases[0..N-1] at indices 1..N. This guarantees
// smooth phase transitions at every chunk boundary with zero discontinuity.
//
// Zero-allocation in steady state: all buffers are pre-allocated on first
// call and reused. The returned CompositeFrame is an internal buffer —
// valid only until the next Process() call.
type FMUpsampler struct {
	srcRate      float64 // composite rate (e.g. 228000)
	dstRate      float64 // device rate (e.g. 2280000)
	maxDeviation float64 // peak FM deviation in Hz (e.g. 75000)
	step         float64 // source-samples per output-sample = srcRate/dstRate

	// Persistent state across Process() calls
	phase     float64 // accumulated FM phase in radians, continuous across chunks
	prevPhase float64 // phase at end of previous chunk (virtual index 0)
	srcPos    float64 // fractional source position carry-over into next chunk
	seeded    bool    // true after first Process() call

	// Pre-allocated buffers — grown once, never shrunk
	srcPhases []float64
	outBuf    []output.IQSample
	outFrame  output.CompositeFrame
}

// NewFMUpsampler creates a phase-domain upsampler.
//
// srcRate:      composite DSP rate (typ. 228000 Hz)
// dstRate:      device output rate (typ. 2280000 Hz)
// maxDeviation: FM peak deviation (typ. 75000 Hz)
func NewFMUpsampler(srcRate, dstRate, maxDeviation float64) *FMUpsampler {
	return &FMUpsampler{
		srcRate:      srcRate,
		dstRate:      dstRate,
		maxDeviation: maxDeviation,
		step:         srcRate / dstRate,
	}
}

// Process takes a CompositeFrame where Samples[i].I contains the composite
// baseband value (FM modulation must be OFF in the generator; Q is ignored).
// Returns an FM-modulated IQ frame at dstRate.
//
// The returned frame is an internal buffer — valid until the next Process()
// call. The caller must consume or copy the data before calling again.
func (u *FMUpsampler) Process(frame *output.CompositeFrame) *output.CompositeFrame {
	if frame == nil || len(frame.Samples) == 0 {
		return frame
	}

	srcLen := len(frame.Samples)

	// --- Phase accumulation at source rate ---

	// Grow srcPhases buffer if needed
	if cap(u.srcPhases) < srcLen {
		u.srcPhases = make([]float64, srcLen)
	}
	srcPhases := u.srcPhases[:srcLen]

	phaseInc := 2 * math.Pi * u.maxDeviation / u.srcRate
	for i, s := range frame.Samples {
		u.phase += float64(s.I) * phaseInc
		srcPhases[i] = u.phase
	}

	// Phase wrapping — symmetric, shift prevPhase in lockstep
	if u.phase > math.Pi || u.phase < -math.Pi {
		offset := 2 * math.Pi * math.Floor((u.phase+math.Pi)/(2*math.Pi))
		u.phase -= offset
		for i := range srcPhases {
			srcPhases[i] -= offset
		}
		if u.seeded {
			u.prevPhase -= offset
		}
	}

	// Seed prevPhase on very first call
	if !u.seeded {
		// Extrapolate backwards from first two phases to get a virtual "previous"
		// phase, so the first chunk's boundary interpolation is smooth.
		if srcLen >= 2 {
			u.prevPhase = 2*srcPhases[0] - srcPhases[1]
		} else {
			u.prevPhase = srcPhases[0]
		}
		u.srcPos = 0
		u.seeded = true
	}

	// --- Interpolation coordinate system ---
	//
	// Virtual index 0 = prevPhase (end of previous chunk)
	// Virtual index 1 = srcPhases[0]
	// Virtual index 2 = srcPhases[1]
	// ...
	// Virtual index N = srcPhases[N-1]
	//
	// srcPos ranges from 0 (carry-over) to N (= srcLen).
	// We generate output samples while srcPos < srcLen (virtual index srcLen).

	// phaseAt returns the phase at a virtual index.
	phaseAt := func(vi int) float64 {
		if vi <= 0 {
			return u.prevPhase
		}
		if vi > srcLen {
			return srcPhases[srcLen-1]
		}
		return srcPhases[vi-1]
	}

	// Calculate output count: from srcPos to srcLen, stepping by u.step.
	// +1 for safety margin; we clamp below.
	maxOut := int(math.Ceil(float64(srcLen)-u.srcPos)/u.step) + 1
	if maxOut < 0 {
		maxOut = 0
	}

	// Grow output buffer if needed
	if cap(u.outBuf) < maxOut {
		u.outBuf = make([]output.IQSample, maxOut)
	}
	out := u.outBuf[:maxOut]

	// --- Generate output samples ---
	pos := u.srcPos
	n := 0
	for pos < float64(srcLen) && n < maxOut {
		vi := int(pos)     // virtual index (integer part)
		frac := pos - float64(vi)

		pA := phaseAt(vi)
		pB := phaseAt(vi + 1)
		p := pA + frac*(pB-pA)

		out[n] = output.IQSample{
			I: float32(math.Cos(p)),
			Q: float32(math.Sin(p)),
		}
		n++
		pos += u.step
	}

	// Carry state for next chunk
	u.prevPhase = srcPhases[srcLen-1]
	u.srcPos = pos - float64(srcLen)

	// Package output
	u.outFrame.Samples = out[:n]
	u.outFrame.SampleRateHz = u.dstRate
	u.outFrame.Timestamp = frame.Timestamp
	u.outFrame.Sequence = frame.Sequence

	return &u.outFrame
}

// Reset clears all accumulated state, as if freshly constructed.
func (u *FMUpsampler) Reset() {
	u.phase = 0
	u.prevPhase = 0
	u.srcPos = 0
	u.seeded = false
}

// Stats returns internal state for diagnostics/testing.
func (u *FMUpsampler) Stats() (phase, prevPhase, srcPos float64) {
	return u.phase, u.prevPhase, u.srcPos
}