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

package offline

import (
	"context"
	"encoding/binary"
	"fmt"
	"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/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
}

// 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
	limiter       *dsp.StereoLimiter  // stereo-linked, operates on L/R BEFORE stereo encoding
	fmMod         *dsp.FMModulator
	sampleRate    float64
	initialized   bool
	frameSeq      uint64

	// Broadcast-standard audio filter chain (per channel, L and R):
	//   Pre-emphasis → 15kHz LPF (4th-order) → 19kHz Notch → Drive → Limiter
	audioLPF_L    *dsp.FilterChain  // 4th-order Butterworth 15kHz
	audioLPF_R    *dsp.FilterChain
	pilotNotchL   *dsp.Biquad       // 19kHz notch (guard band protection)
	pilotNotchR   *dsp.Biquad

	// Composite clipper protection (post-clip notch filters):
	//   Audio composite → clip → notch 19kHz → notch 57kHz → + pilot → + RDS
	mpxNotch19    *dsp.Biquad       // removes clip harmonics at pilot freq
	mpxNotch57    *dsp.Biquad       // removes clip harmonics at RDS freq

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

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

// SetExternalSource sets a live audio source (e.g. StreamResampler) that
// takes priority over WAV/tone sources. Must be called before the first
// GenerateFrame() call (i.e. before init).
func (g *Generator) SetExternalSource(src frameSource) {
	g.externalSource = src
}

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

// 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 }
	// Audio ceiling leaves headroom for pilot + RDS so total ≤ ceiling
	pilotAmp := g.cfg.FM.PilotLevel * g.cfg.FM.OutputDrive
	rdsAmp := g.cfg.FM.RDSInjection * g.cfg.FM.OutputDrive
	audioCeiling := ceiling - pilotAmp - rdsAmp
	if audioCeiling < 0.3 { audioCeiling = 0.3 }
	if g.cfg.FM.LimiterEnabled {
		g.limiter = dsp.NewStereoLimiter(audioCeiling, 0.5, 100, g.sampleRate)
	}

	// Broadcast-standard filter chain:
	// 1) 15kHz 4th-order Butterworth LPF — steep guard band, -14dB@19kHz, -54dB@57kHz
	g.audioLPF_L = dsp.NewAudioLPF(g.sampleRate)
	g.audioLPF_R = dsp.NewAudioLPF(g.sampleRate)
	// 2) 19kHz notch — kills residual audio at pilot freq, >40dB rejection
	g.pilotNotchL = dsp.NewPilotNotch(g.sampleRate)
	g.pilotNotchR = dsp.NewPilotNotch(g.sampleRate)
	// 3) Composite clipper protection notches at 19kHz + 57kHz
	g.mpxNotch19, g.mpxNotch57 = dsp.NewCompositeProtection(g.sampleRate)
	if g.cfg.FM.FMModulationEnabled {
		g.fmMod = dsp.NewFMModulator(g.sampleRate)
		if g.cfg.FM.MaxDeviationHz > 0 { g.fmMod.MaxDeviation = g.cfg.FM.MaxDeviationHz }
	}

	// 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,
	})

	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}
		}
		return audio.NewConfiguredToneSource(sampleRate, g.cfg.Audio.ToneLeftHz, g.cfg.Audio.ToneRightHz, g.cfg.Audio.ToneAmplitude), SourceInfo{Kind: "tone-fallback", SampleRate: sampleRate, Detail: g.cfg.Audio.InputPath}
	}
	return audio.NewConfiguredToneSource(sampleRate, g.cfg.Audio.ToneLeftHz, g.cfg.Audio.ToneRightHz, g.cfg.Audio.ToneAmplitude), 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

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

	// Broadcast-standard FM MPX signal chain:
	//
	//   Audio L/R
	//     → PreEmphasis (50µs EU / 75µs US)
	//     → 15kHz LPF (4th-order Butterworth, -14dB@19kHz, -54dB@57kHz)
	//     → 19kHz Notch (>40dB rejection, guard band protection)
	//     → × OutputDrive
	//     → StereoLimiter (instant attack, smooth release)
	//     → Stereo Encode → Mono (L+R)/2 + Stereo Sub (L-R)/2 × 38kHz
	//   Audio MPX composite
	//     → HardClip at audioCeiling (catches limiter overshoots)
	//     → 19kHz Notch (removes clip harmonics at pilot freq)
	//     → 57kHz Notch (removes clip harmonics at RDS freq)
	//   + Pilot 19kHz (fixed amplitude, post-clip)
	//   + RDS 57kHz (fixed amplitude, post-clip)
	//   → FM Modulator
	//
	// Key: Pilot and RDS are NEVER clipped or filtered. They're added
	// after all audio processing at constant amplitude.
	ceiling := lp.LimiterCeiling
	if ceiling <= 0 { ceiling = 1.0 }
	pilotAmp := lp.PilotLevel * lp.OutputDrive
	rdsAmp := lp.RDSInjection * lp.OutputDrive
	audioCeiling := ceiling - pilotAmp - rdsAmp
	if audioCeiling < 0.3 { audioCeiling = 0.3 }

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

		// --- Stage 1: Audio filtering (per-channel) ---
		// 15kHz LPF removes out-of-band pre-emphasis energy.
		// 19kHz notch kills residual energy at pilot frequency.
		// Both run BEFORE drive+limiter so the limiter sees the
		// actual audio bandwidth, not wasted HF energy.
		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 + Limit ---
		l *= lp.OutputDrive
		r *= lp.OutputDrive
		if lp.LimiterEnabled && g.limiter != nil {
			l, r = g.limiter.Process(l, r)
		}

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

		// --- Stage 4: Audio composite clip + protection ---
		// Clip the audio-only composite (mono + stereo sub) to budget.
		// Then notch-filter the clip harmonics out of the pilot (19kHz)
		// and RDS (57kHz) bands before adding the real pilot and RDS.
		audioMPX := float64(comps.Mono)
		if lp.StereoEnabled {
			audioMPX += float64(comps.Stereo)
		}
		audioMPX = dsp.HardClip(audioMPX, audioCeiling)
		audioMPX = g.mpxNotch19.Process(audioMPX)
		audioMPX = g.mpxNotch57.Process(audioMPX)

		// --- Stage 5: Add protected components at fixed levels ---
		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
		}

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