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.MPXLimiter
	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]
}

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

// 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 }
	if g.cfg.FM.LimiterEnabled {
		g.limiter = dsp.NewMPXLimiter(ceiling, 0.1, 50, 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.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}
	}

	// Apply live combiner gains
	g.combiner.PilotGain = lp.PilotLevel
	g.combiner.RDSGain = lp.RDSInjection

	ceiling := lp.LimiterCeiling
	if ceiling <= 0 { ceiling = 1.0 }

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

		comps := g.stereoEncoder.Encode(in)
		if !lp.StereoEnabled {
			comps.Stereo = 0; comps.Pilot = 0
		}

		rdsValue := 0.0
		if g.rdsEnc != nil && lp.RDSEnabled {
			rdsCarrier := g.stereoEncoder.RDSCarrier()
			rdsValue = g.rdsEnc.NextSampleWithCarrier(rdsCarrier)
		}

		composite := g.combiner.Combine(comps.Mono, comps.Stereo, comps.Pilot, rdsValue)
		composite *= lp.OutputDrive

		if lp.LimiterEnabled && g.limiter != nil {
			composite = g.limiter.Process(composite)
			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)
}