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

package offline

import (
	"context"
	"encoding/binary"
	"fmt"
	"log"
	"path/filepath"
	"sync/atomic"
	"time"

	"github.com/jan/fm-rds-tx/internal/audio"
	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"
	"github.com/jan/fm-rds-tx/internal/dsp"
	"github.com/jan/fm-rds-tx/internal/mpx"
	"github.com/jan/fm-rds-tx/internal/output"
	"github.com/jan/fm-rds-tx/internal/rds"
	"github.com/jan/fm-rds-tx/internal/stereo"
)

type frameSource interface {
	NextFrame() audio.Frame
}

// LiveParams carries DSP parameters that can be hot-swapped at runtime.
// Loaded once per chunk via atomic pointer — zero per-sample overhead.
type LiveParams struct {
	OutputDrive    float64
	StereoEnabled  bool
	PilotLevel     float64
	RDSInjection   float64
	RDSEnabled     bool
	LimiterEnabled bool
	LimiterCeiling float64
	MpxGain        float64 // hardware calibration factor for composite output
	// Tone + gain: live-patchable without DSP chain reinit.
	ToneLeftHz    float64
	ToneRightHz   float64
	ToneAmplitude float64
	AudioGain     float64
	// Composite clipper: live-toggleable without DSP chain reinit.
	CompositeClipperEnabled bool
}

// PreEmphasizedSource wraps an audio source and applies pre-emphasis.
// The source is expected to already output at composite rate (resampled
// upstream). Pre-emphasis is applied per-sample at that rate.
type PreEmphasizedSource struct {
	src  frameSource
	preL *dsp.PreEmphasis
	preR *dsp.PreEmphasis
	gain float64
}

func NewPreEmphasizedSource(src frameSource, tauUS, sampleRate, gain float64) *PreEmphasizedSource {
	p := &PreEmphasizedSource{src: src, gain: gain}
	if tauUS > 0 {
		p.preL = dsp.NewPreEmphasis(tauUS, sampleRate)
		p.preR = dsp.NewPreEmphasis(tauUS, sampleRate)
	}
	return p
}

func (p *PreEmphasizedSource) NextFrame() audio.Frame {
	f := p.src.NextFrame()
	l := float64(f.L) * p.gain
	r := float64(f.R) * p.gain
	if p.preL != nil {
		l = p.preL.Process(l)
		r = p.preR.Process(r)
	}
	return audio.NewFrame(audio.Sample(l), audio.Sample(r))
}

type SourceInfo struct {
	Kind       string
	SampleRate float64
	Detail     string
}

type Generator struct {
	cfg cfgpkg.Config

	// Persistent DSP state across GenerateFrame calls
	source        *PreEmphasizedSource
	stereoEncoder stereo.StereoEncoder
	rdsEnc        *rds.Encoder
	combiner      mpx.DefaultCombiner
	fmMod         *dsp.FMModulator
	sampleRate    float64
	initialized   bool
	frameSeq      uint64

	// Broadcast-standard clip-filter-clip chain (per channel L/R):
	//
	//   PreEmph → LPF₁(14kHz) → Notch(19kHz) → ×Drive
	//   → StereoLimiter (slow AGC: raises average level)
	//   → Clip₁ → LPF₂(14kHz) [cleanup] → Clip₂ [catches LPF overshoots]
	//   → Stereo Encode → Composite Clip → Notch₁₉ → Notch₅₇
	//   → + Pilot → + RDS → FM
	//
	audioLPF_L     *dsp.FilterChain  // 14kHz 8th-order (pre-clip)
	audioLPF_R     *dsp.FilterChain
	pilotNotchL    *dsp.FilterChain  // 19kHz double-notch (guard band)
	pilotNotchR    *dsp.FilterChain
	limiter        *dsp.StereoLimiter // slow compressor (raises average, clips catch peaks)
	cleanupLPF_L   *dsp.FilterChain  // 14kHz 8th-order (post-clip cleanup)
	cleanupLPF_R   *dsp.FilterChain
	mpxNotch19     *dsp.FilterChain  // composite clipper protection
	mpxNotch57     *dsp.FilterChain
	bs412          *dsp.BS412Limiter // ITU-R BS.412 MPX power limiter (optional)
	compositeClip  *dsp.CompositeClipper // ITU-R SM.1268 iterative composite clipper (optional)

	// Pre-allocated frame buffer — reused every GenerateFrame call.
	frameBuf *output.CompositeFrame
	bufCap   int

	// Live-updatable DSP parameters — written by control API, read per chunk.
	liveParams atomic.Pointer[LiveParams]

	// Optional external audio source (e.g. StreamResampler for live audio).
	// When set, takes priority over WAV/tones in sourceFor().
	externalSource frameSource

	// Tone source reference — non-nil when a ToneSource is the active audio input.
	// Allows live-updating tone parameters via LiveParams each chunk.
	toneSource *audio.ToneSource

	// License: jingle injection when unlicensed.
	licenseState *license.State
	jingleFrames []license.JingleFrame

	// Watermark: STFT-domain spread-spectrum (Kirovski & Malvar 2003).
	stftEmbedder *watermark.STFTEmbedder
	wmDecimLPF   *dsp.FilterChain // anti-alias LPF for 228k→12k decimation
}

func NewGenerator(cfg cfgpkg.Config) *Generator {
	return &Generator{cfg: cfg}
}

// SetLicense configures license state (jingle) and creates the STFT watermark
// embedder. Must be called before the first GenerateFrame.
func (g *Generator) SetLicense(state *license.State, key string) {
	g.licenseState = state
	g.stftEmbedder = watermark.NewSTFTEmbedder(key)
}

// SetExternalSource sets a live audio source (e.g. StreamResampler) that
// takes priority over WAV/tone sources. Must be called before the first
// GenerateFrame() call; calling it after init() has no effect because
// g.source is already wired to the old source.
func (g *Generator) SetExternalSource(src frameSource) error {
	if g.initialized {
		return fmt.Errorf("generator: SetExternalSource called after GenerateFrame; call it before the engine starts")
	}
	g.externalSource = src
	return nil
}

// UpdateLive hot-swaps DSP parameters. Thread-safe — called from control API,
// applied at the next chunk boundary by the DSP goroutine.
func (g *Generator) UpdateLive(p LiveParams) {
	g.liveParams.Store(&p)
}

// CurrentLiveParams returns the current live parameter snapshot.
// Used by Engine.UpdateConfig to do read-modify-write on the params.
func (g *Generator) CurrentLiveParams() LiveParams {
	if lp := g.liveParams.Load(); lp != nil {
		return *lp
	}
	return LiveParams{OutputDrive: 1.0, LimiterCeiling: 1.0, MpxGain: 1.0}
}

// RDSEncoder returns the live RDS encoder, or nil if RDS is disabled.
// Used by the Engine to forward text updates.
func (g *Generator) RDSEncoder() *rds.Encoder {
	return g.rdsEnc
}

func (g *Generator) init() {
	if g.initialized {
		return
	}
	g.sampleRate = float64(g.cfg.FM.CompositeRateHz)
	if g.sampleRate <= 0 {
		g.sampleRate = 228000
	}
	rawSource, _ := g.sourceFor(g.sampleRate)
	g.source = NewPreEmphasizedSource(rawSource, g.cfg.FM.PreEmphasisTauUS, g.sampleRate, g.cfg.Audio.Gain)
	g.stereoEncoder = stereo.NewStereoEncoder(g.sampleRate)
	g.combiner = mpx.DefaultCombiner{
		MonoGain: 1.0, StereoGain: 1.0,
		PilotGain: g.cfg.FM.PilotLevel, RDSGain: g.cfg.FM.RDSInjection,
	}
	if g.cfg.RDS.Enabled {
		piCode, _ := cfgpkg.ParsePI(g.cfg.RDS.PI)
		g.rdsEnc, _ = rds.NewEncoder(rds.RDSConfig{
			PI: piCode, PS: g.cfg.RDS.PS, RT: g.cfg.RDS.RadioText,
			PTY: uint8(g.cfg.RDS.PTY), SampleRate: g.sampleRate,
		})
	}
	ceiling := g.cfg.FM.LimiterCeiling
	if ceiling <= 0 { ceiling = 1.0 }

	// Broadcast clip-filter-clip chain:
	// Pre-clip: 14kHz LPF (8th-order) + 19kHz double-notch (per channel)
	g.audioLPF_L = dsp.NewAudioLPF(g.sampleRate)
	g.audioLPF_R = dsp.NewAudioLPF(g.sampleRate)
	g.pilotNotchL = dsp.NewPilotNotch(g.sampleRate)
	g.pilotNotchR = dsp.NewPilotNotch(g.sampleRate)
	// Slow compressor: 5ms attack / 200ms release. Brings average level UP.
	// The clips after it catch the peaks the limiter's attack time misses.
	// This is the "slow-to-fast progression" from broadcast processing:
	// slow limiter → fast clips.
	g.limiter = dsp.NewStereoLimiter(ceiling, 5, 200, g.sampleRate)
	// Post-clip cleanup: second 14kHz LPF pass (removes clip harmonics)
	g.cleanupLPF_L = dsp.NewAudioLPF(g.sampleRate)
	g.cleanupLPF_R = dsp.NewAudioLPF(g.sampleRate)
	// Composite clipper protection: double-notch at 19kHz + 57kHz
	g.mpxNotch19, g.mpxNotch57 = dsp.NewCompositeProtection(g.sampleRate)
	// ITU-R SM.1268 iterative composite clipper (optional, replaces simple clip+notch)
	// Always created so it can be live-toggled via CompositeClipperEnabled.
	g.compositeClip = dsp.NewCompositeClipper(dsp.CompositeClipperConfig{
		Ceiling:     ceiling,
		Iterations:  g.cfg.FM.CompositeClipper.Iterations,
		SoftKnee:    g.cfg.FM.CompositeClipper.SoftKnee,
		LookaheadMs: g.cfg.FM.CompositeClipper.LookaheadMs,
		SampleRate:  g.sampleRate,
	})
	// BS.412 MPX power limiter (EU/CH requirement for licensed FM)
	if g.cfg.FM.BS412Enabled {
		// chunkSec is not known at init time (Engine.chunkDuration may differ).
		// Pass 0 here; GenerateFrame computes the actual chunk duration from
		// the real sample count and updates BS.412 accordingly.
		g.bs412 = dsp.NewBS412Limiter(
			g.cfg.FM.BS412ThresholdDBr,
			g.cfg.FM.PilotLevel,
			g.cfg.FM.RDSInjection,
			0,
		)
	}
	if g.cfg.FM.FMModulationEnabled {
		g.fmMod = dsp.NewFMModulator(g.sampleRate)
		maxDev := g.cfg.FM.MaxDeviationHz
		if maxDev > 0 {
			if g.cfg.FM.MpxGain > 0 && g.cfg.FM.MpxGain != 1.0 {
				maxDev *= g.cfg.FM.MpxGain
			}
			g.fmMod.MaxDeviation = maxDev
		}
	}

	// Seed initial live params from config
	g.liveParams.Store(&LiveParams{
		OutputDrive:    g.cfg.FM.OutputDrive,
		StereoEnabled:  g.cfg.FM.StereoEnabled,
		PilotLevel:     g.cfg.FM.PilotLevel,
		RDSInjection:   g.cfg.FM.RDSInjection,
		RDSEnabled:     g.cfg.RDS.Enabled,
		LimiterEnabled: g.cfg.FM.LimiterEnabled,
		LimiterCeiling: ceiling,
		MpxGain:        g.cfg.FM.MpxGain,
		ToneLeftHz:     g.cfg.Audio.ToneLeftHz,
		ToneRightHz:    g.cfg.Audio.ToneRightHz,
		ToneAmplitude:  g.cfg.Audio.ToneAmplitude,
		AudioGain:      g.cfg.Audio.Gain,
		CompositeClipperEnabled: g.cfg.FM.CompositeClipper.Enabled,
	})

	if g.licenseState != nil {
		frames, err := license.LoadJingleFrames(license.JingleWAV(), g.sampleRate)
		if err != nil {
			log.Printf("license: jingle load failed: %v", err)
		} else {
			g.jingleFrames = frames
		}
	}

	// STFT watermark: anti-alias LPF for decimation to WMRate (12 kHz).
	// Nyquist at 12 kHz = 6 kHz. Cut at 5.5 kHz with margin.
	if g.stftEmbedder != nil {
		g.wmDecimLPF = dsp.NewLPF4(5500, g.sampleRate)
	}

	g.initialized = true
}

func (g *Generator) sourceFor(sampleRate float64) (frameSource, SourceInfo) {
	if g.externalSource != nil {
		return g.externalSource, SourceInfo{Kind: "stream", SampleRate: sampleRate, Detail: "live audio"}
	}
	if g.cfg.Audio.InputPath != "" {
		if src, err := audio.LoadWAVSource(g.cfg.Audio.InputPath); err == nil {
			return audio.NewResampledSource(src, sampleRate), SourceInfo{Kind: "wav", SampleRate: float64(src.SampleRate), Detail: g.cfg.Audio.InputPath}
		}
		ts := audio.NewConfiguredToneSource(sampleRate, g.cfg.Audio.ToneLeftHz, g.cfg.Audio.ToneRightHz, g.cfg.Audio.ToneAmplitude)
		g.toneSource = ts
		return ts, SourceInfo{Kind: "tone-fallback", SampleRate: sampleRate, Detail: g.cfg.Audio.InputPath}
	}
	ts := audio.NewConfiguredToneSource(sampleRate, g.cfg.Audio.ToneLeftHz, g.cfg.Audio.ToneRightHz, g.cfg.Audio.ToneAmplitude)
	g.toneSource = ts
	return ts, SourceInfo{Kind: "tones", SampleRate: sampleRate, Detail: "generated"}
}

func (g *Generator) GenerateFrame(duration time.Duration) *output.CompositeFrame {
	g.init()

	samples := int(duration.Seconds() * g.sampleRate)
	if samples <= 0 { samples = int(g.sampleRate / 10) }

	// Reuse buffer — grow only if needed, never shrink
	if g.frameBuf == nil || g.bufCap < samples {
		g.frameBuf = &output.CompositeFrame{
			Samples: make([]output.IQSample, samples),
		}
		g.bufCap = samples
	}
	frame := g.frameBuf
	frame.Samples = frame.Samples[:samples]
	frame.SampleRateHz = g.sampleRate
	frame.Timestamp = time.Now().UTC()
	g.frameSeq++
	frame.Sequence = g.frameSeq

	// L/R buffers for two-pass processing (STFT watermark between stages 3 and 4)
	lBuf := make([]float64, samples)
	rBuf := make([]float64, samples)

	// Load live params once per chunk — single atomic read, zero per-sample cost
	lp := g.liveParams.Load()
	if lp == nil {
		// Fallback: should never happen after init(), but be safe
		lp = &LiveParams{OutputDrive: 1.0, LimiterCeiling: 1.0, MpxGain: 1.0}
	}

	// Apply live tone and gain updates each chunk. GenerateFrame runs on a
	// single goroutine so these field writes are safe without additional locking.
	if g.toneSource != nil {
		g.toneSource.LeftFreq  = lp.ToneLeftHz
		g.toneSource.RightFreq = lp.ToneRightHz
		g.toneSource.Amplitude = lp.ToneAmplitude
	}
	if g.source != nil {
		g.source.gain = lp.AudioGain
	}

	// Broadcast clip-filter-clip FM MPX signal chain:
	//
	//   Audio L/R → PreEmphasis
	//     → LPF₁ (14kHz, 8th-order)  → 19kHz Notch (double)
	//     → × OutputDrive             → HardClip₁ (ceiling)
	//     → LPF₂ (14kHz, 8th-order)  [removes clip₁ harmonics]
	//     → HardClip₂ (ceiling)       [catches LPF₂ overshoots]
	//     → Stereo Encode
	//   Audio MPX (mono + stereo sub)
	//     → HardClip₃ (ceiling)       [composite deviation control]
	//     → 19kHz Notch (double)      [protect pilot band]
	//     → 57kHz Notch (double)      [protect RDS band]
	//   + Pilot 19kHz (fixed, NEVER clipped)
	//   + RDS 57kHz (fixed, NEVER clipped)
	//   → FM Modulator
	//
	// Guard band depth at 19kHz: LPF₁(-21dB) + Notch(-60dB) + LPF₂(-21dB)
	//   + CompNotch(-60dB) → broadband floor -42dB, exact 19kHz >-90dB
	ceiling := lp.LimiterCeiling
	if ceiling <= 0 { ceiling = 1.0 }
	// Pilot and RDS are FIXED injection levels, independent of OutputDrive.
	// Config values directly represent percentage of ±75kHz deviation:
	//   pilotLevel: 0.09 = 9% = ±6.75kHz (ITU standard)
	//   rdsInjection: 0.04 = 4% = ±3.0kHz (typical)
	pilotAmp := lp.PilotLevel
	rdsAmp := lp.RDSInjection

	// BS.412 MPX power limiter: uses previous chunk's measurement to set gain.
	// Power is measured during this chunk and fed back at the end.
	bs412Gain := 1.0
	var bs412PowerAccum float64
	if g.bs412 != nil {
		bs412Gain = g.bs412.CurrentGain()
	}

	if g.licenseState != nil {
		g.licenseState.Tick()
	}

	for i := 0; i < samples; i++ {
		in := g.source.NextFrame()

		// --- Stage 1: Band-limit pre-emphasized audio ---
		l := g.audioLPF_L.Process(float64(in.L))
		l = g.pilotNotchL.Process(l)
		r := g.audioLPF_R.Process(float64(in.R))
		r = g.pilotNotchR.Process(r)

		// --- Stage 2: Drive + Compress + Clip₁ ---
		l *= lp.OutputDrive
		r *= lp.OutputDrive
		if g.limiter != nil {
			l, r = g.limiter.Process(l, r)
		}
		l = dsp.HardClip(l, ceiling)
		r = dsp.HardClip(r, ceiling)

		// --- Stage 3: Cleanup LPF + Clip₂ (overshoot compensator) ---
		l = g.cleanupLPF_L.Process(l)
		r = g.cleanupLPF_R.Process(r)
		l = dsp.HardClip(l, ceiling)
		r = dsp.HardClip(r, ceiling)

		lBuf[i] = l
		rBuf[i] = r
	}

	// --- STFT Watermark: decimate → embed → upsample → add to L/R ---
	if g.stftEmbedder != nil {
		decimFactor := int(g.sampleRate) / watermark.WMRate // 228000/12000 = 19
		if decimFactor < 1 {
			decimFactor = 1
		}
		nDown := samples / decimFactor

		// Anti-alias: LPF ALL composite-rate samples, THEN decimate.
		// The LPF must see every sample for correct IIR state update.
		mono12k := make([]float64, nDown)
		lpfState := 0.0
		decimCount := 0
		outIdx := 0
		for i := 0; i < samples && outIdx < nDown; i++ {
			mono := (lBuf[i] + rBuf[i]) / 2
			if g.wmDecimLPF != nil {
				lpfState = g.wmDecimLPF.Process(mono)
			} else {
				lpfState = mono
			}
			decimCount++
			if decimCount >= decimFactor {
				decimCount = 0
				mono12k[outIdx] = lpfState
				outIdx++
			}
		}

		// STFT embed at 12 kHz
		embedded := g.stftEmbedder.ProcessBlock(mono12k)

		// Extract watermark signal (difference) and upsample via ZOH
		for i := 0; i < samples; i++ {
			wmIdx := i / decimFactor
			if wmIdx >= nDown {
				wmIdx = nDown - 1
			}
			wmSig := embedded[wmIdx] - mono12k[wmIdx]
			lBuf[i] += wmSig
			rBuf[i] += wmSig
		}
	}

	// --- Pass 2: Stereo encode + composite processing ---
	for i := 0; i < samples; i++ {
		l := lBuf[i]
		r := rBuf[i]

		// --- Stage 4: Stereo encode ---
		limited := audio.NewFrame(audio.Sample(l), audio.Sample(r))
		comps := g.stereoEncoder.Encode(limited)

		// --- Stage 5: Composite clip + protection ---
		audioMPX := float64(comps.Mono)
		if lp.StereoEnabled {
			audioMPX += float64(comps.Stereo)
		}
		if lp.CompositeClipperEnabled && g.compositeClip != nil {
			// ITU-R SM.1268 iterative clipper: look-ahead + N×(clip→notch→notch) + final clip
			audioMPX = g.compositeClip.Process(audioMPX)
		} else {
			// Legacy single-pass: one clip, then notch, no final safety clip
			audioMPX = dsp.HardClip(audioMPX, ceiling)
			audioMPX = g.mpxNotch19.Process(audioMPX)
			audioMPX = g.mpxNotch57.Process(audioMPX)
		}

		// BS.412: apply gain and measure power
		if bs412Gain < 1.0 {
			audioMPX *= bs412Gain
		}
		bs412PowerAccum += audioMPX * audioMPX

		// --- Stage 6: Add protected components ---
		composite := audioMPX
		if lp.StereoEnabled {
			composite += pilotAmp * comps.Pilot
		}
		if g.rdsEnc != nil && lp.RDSEnabled {
			rdsCarrier := g.stereoEncoder.RDSCarrier()
			rdsValue := g.rdsEnc.NextSampleWithCarrier(rdsCarrier)
			composite += rdsAmp * rdsValue
		}

		// Jingle: injected when unlicensed, bypasses drive/gain controls.
		if g.licenseState != nil && len(g.jingleFrames) > 0 {
			composite += g.licenseState.NextSample(g.jingleFrames)
		}
		if g.fmMod != nil {
			iq_i, iq_q := g.fmMod.Modulate(composite)
			frame.Samples[i] = output.IQSample{I: float32(iq_i), Q: float32(iq_q)}
		} else {
			frame.Samples[i] = output.IQSample{I: float32(composite), Q: 0}
		}
	}

	// BS.412: feed this chunk's actual duration and average audio power for
	// the next chunk's gain calculation. Using the real sample count avoids
	// the error that occurred when chunkSec was hardcoded to 0.05 — any
	// SetChunkDuration() call from the engine would silently miscalibrate
	// the ITU-R BS.412 power measurement window.
	if g.bs412 != nil && samples > 0 {
		chunkSec := float64(samples) / g.sampleRate
		g.bs412.UpdateChunkDuration(chunkSec)
		g.bs412.ProcessChunk(bs412PowerAccum / float64(samples))
	}

	// STFT watermark diagnostic
	if g.stftEmbedder != nil && g.frameSeq%100 == 1 {
		log.Printf("watermark stft: frame=%d, active", g.frameSeq)
	}

	return frame
}

func (g *Generator) WriteFile(path string, duration time.Duration) error {
	if path == "" {
		path = g.cfg.Backend.OutputPath
	}
	if path == "" {
		path = filepath.Join("build", "offline", "composite.iqf32")
	}
	backend, err := output.NewFileBackend(path, binary.LittleEndian, output.BackendInfo{
		Name:        "offline-file",
		Description: "offline composite file backend",
	})
	if err != nil {
		return err
	}
	defer backend.Close(context.Background())

	if err := backend.Configure(context.Background(), output.BackendConfig{
		SampleRateHz: float64(g.cfg.FM.CompositeRateHz),
		Channels:     2,
		IQLevel:      float32(g.cfg.FM.OutputDrive),
	}); err != nil {
		return err
	}

	frame := g.GenerateFrame(duration)
	if _, err := backend.Write(context.Background(), frame); err != nil {
		return err
	}
	return backend.Flush(context.Background())
}

func (g *Generator) Summary(duration time.Duration) string {
	sampleRate := float64(g.cfg.FM.CompositeRateHz)
	if sampleRate <= 0 {
		sampleRate = 228000
	}
	_, info := g.sourceFor(sampleRate)
	preemph := "off"
	if g.cfg.FM.PreEmphasisTauUS > 0 {
		preemph = fmt.Sprintf("%.0fµs", g.cfg.FM.PreEmphasisTauUS)
	}
	modMode := "composite"
	if g.cfg.FM.FMModulationEnabled {
		modMode = fmt.Sprintf("FM-IQ(±%.0fHz)", g.cfg.FM.MaxDeviationHz)
	}
	return fmt.Sprintf("offline frame: freq=%.1fMHz rate=%d duration=%s drive=%.2f stereo=%t rds=%t preemph=%s limiter=%t output=%s source=%s detail=%s",
		g.cfg.FM.FrequencyMHz, g.cfg.FM.CompositeRateHz, duration.String(),
		g.cfg.FM.OutputDrive, g.cfg.FM.StereoEnabled, g.cfg.RDS.Enabled,
		preemph, g.cfg.FM.LimiterEnabled, modMode, info.Kind, info.Detail)
}