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
	lpfL, lpfR    *dsp.BiquadLPF    // 15kHz lowpass after limiter, protects RDS band
	fmMod         *dsp.FMModulator
	sampleRate    float64
	initialized   bool
	frameSeq      uint64

	// 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)
	}
	// 15kHz lowpass after limiter — removes limiter gain-step intermodulation
	// products that would otherwise fall into pilot/stereo/RDS bands.
	g.lpfL = dsp.NewBiquadLPF(15000, g.sampleRate)
	g.lpfR = dsp.NewBiquadLPF(15000, 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}
	}

	// Signal path (matches professional broadcast processors):
	//   Audio L/R → × Drive → Stereo-linked limiter → Stereo encoder
	//   → Mono + Stereo sub (from limited audio, natural levels)
	//   → + Pilot (fixed) → + RDS (fixed) → FM modulator
	//
	// The limiter never sees the 38kHz subcarrier, so it can't pump
	// the stereo difference signal. Pilot and RDS are post-encoder
	// at fixed amplitudes, unaffected by audio dynamics.
	//
	// Audio ceiling is auto-reduced to leave headroom for pilot + RDS,
	// so total composite stays within ±ceiling (= ±75kHz deviation).
	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 } // safety floor

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

		// --- Stage 1: Band-limit pre-emphasized audio ---
		// The 15kHz LPF goes BEFORE drive+limiter. Pre-emphasis boosts
		// HF by up to +13.5dB. Without the LPF, the limiter would waste
		// gain reduction on HF peaks that get filtered later, causing
		// wild modulation swings (30-163%). With LPF first, the limiter
		// sees the final audio bandwidth and sets gain correctly.
		l := g.lpfL.Process(float64(in.L))
		r := g.lpfR.Process(float64(in.R))

		// --- Stage 2: Scale and limit ---
		l *= lp.OutputDrive
		r *= lp.OutputDrive

		if lp.LimiterEnabled && g.limiter != nil {
			l, r = g.limiter.Process(l, r)
		}

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

		// --- Stage 3: Combine at fixed levels ---
		composite := float64(comps.Mono)
		if lp.StereoEnabled {
			composite += float64(comps.Stereo)
			composite += pilotAmp * comps.Pilot
		}
		if g.rdsEnc != nil && lp.RDSEnabled {
			rdsCarrier := g.stereoEncoder.RDSCarrier()
			rdsValue := g.rdsEnc.NextSampleWithCarrier(rdsCarrier)
			composite += rdsAmp * rdsValue
		}

		// Final composite safety clip — only fires on brief limiter
		// overshoots during fast transients. Clips the entire composite,
		// not individual audio bands, so harmonics don't target RDS.
		if lp.LimiterEnabled {
			composite = dsp.HardClip(composite, ceiling)
		}

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