// Package cwdecode is a real-time CW (Morse) decoder: it turns a stream of
// mono PCM samples into decoded text. The pipeline is the classic one — a bank
// of Goertzel tone detectors, a pitch LOCK that follows a single tone (so QRM
// at other pitches is ignored), an adaptive envelope/threshold on the LOCKED
// tone (level-independent, so weak or strong signals both key cleanly), an
// adaptive dot-length (WPM) estimate, and a timing state machine that maps
// marks/spaces to Morse and then to characters.
//
// It is deliberately self-contained and dependency-free so it can be unit
// tested with synthetic signals. As with every audio CW decoder, weak signals
// and very heavy QRM still degrade it; the pitch lock keeps QRM on other tones
// out of the decode.
package cwdecode

import (
	"math"
	"sort"
	"sync/atomic"
)

// Status is a periodic snapshot for the UI (pitch lock, speed, signal).
type Status struct {
	WPM    int     `json:"wpm"`
	Pitch  int     `json:"pitch"`  // Hz of the locked tone (0 = not locked)
	Level  float64 `json:"level"`  // 0..1 input audio level (RMS) for the meter
	Active bool    `json:"active"` // a tone is currently keyed down
}

// Decoder consumes PCM and emits decoded characters via onChar (one or more
// characters at a time, including " " for word gaps) and periodic onStatus.
type Decoder struct {
	fs     int
	hop    int // samples between updates
	win    int // Goertzel window length
	freqs  []float64
	coeffs []float64 // precomputed 2*cos(w) per freq

	ring []float64 // last win samples
	acc  int       // samples since last hop
	mags []float64 // per-bin magnitude this hop
	nbuf []float64 // scratch for the noise percentile

	// Fixed-pitch target (Hz). 0 = auto-search; >0 = lock to the nearest bin and
	// ignore everything else (e.g. follow the radio's CW pitch). Set live from
	// another goroutine, so it's atomic.
	targetHz atomic.Int32

	// Pitch lock.
	lockIdx    int     // index of the locked tone bin, -1 = unlocked
	candIdx    int     // current argmax candidate while unlocked
	candHops   int     // consecutive hops the candidate has been dominant
	quietHops  int     // consecutive key-up hops while locked
	noise      float64 // broadband noise estimate (percentile of bins)
	relockHops int     // quiet hops before the lock is released
	acqSNR     float64 // tone/noise ratio to acquire after a few stable hops
	strongSNR  float64 // tone/noise ratio to lock immediately (1 hop)

	// Adaptive keying envelope, on the LOCKED bin's magnitude.
	peak, floor float64
	state       bool // true = mark (key down)
	stateHops   int
	dotHops     float64 // adaptive dot length, in hops
	elem        []byte  // current "." / "-" run for the in-progress character
	charEmitted bool
	wordEmitted bool

	lastPitch float64
	lastRMS   float64

	statusEvery int
	sinceStatus int

	onChar   func(string)
	onStatus func(Status)
}

var morse = map[string]byte{
	".-": 'A', "-...": 'B', "-.-.": 'C', "-..": 'D', ".": 'E', "..-.": 'F',
	"--.": 'G', "....": 'H', "..": 'I', ".---": 'J', "-.-": 'K', ".-..": 'L',
	"--": 'M', "-.": 'N', "---": 'O', ".--.": 'P', "--.-": 'Q', ".-.": 'R',
	"...": 'S', "-": 'T', "..-": 'U', "...-": 'V', ".--": 'W', "-..-": 'X',
	"-.--": 'Y', "--..": 'Z',
	"-----": '0', ".----": '1', "..---": '2', "...--": '3', "....-": '4',
	".....": '5', "-....": '6', "--...": '7', "---..": '8', "----.": '9',
	".-.-.-": '.', "--..--": ',', "..--..": '?', "-..-.": '/', "-...-": '=',
	".-.-.": '+', "-.-.--": '!', "---...": ':', "-....-": '-', ".--.-.": '@',
}

// New builds a decoder for the given sample rate. onChar receives decoded text
// incrementally; onStatus receives ~10 snapshots/second. Either may be nil.
func New(sampleRate int, onChar func(string), onStatus func(Status)) *Decoder {
	if sampleRate <= 0 {
		sampleRate = 16000
	}
	d := &Decoder{
		fs:          sampleRate,
		hop:         sampleRate / 250, // ~4 ms resolution
		win:         sampleRate / 72,  // ~14 ms Goertzel window (selective, fairly snappy)
		dotHops:     15,               // ~20 WPM seed
		acqSNR:      1.9,              // ignore noise spikes (looser locked onto noise = garbage)
		strongSNR:   3.2,              // only a genuinely strong tone locks in 1 hop
		lockIdx:     -1,
		candIdx:     -1,
		statusEvery: 25, // ~10 Hz
		onChar:      onChar,
		onStatus:    onStatus,
	}
	if d.hop < 1 {
		d.hop = 1
	}
	d.relockHops = int(0.8 * float64(d.fs) / float64(d.hop)) // release lock after ~0.8 s quiet
	// Candidate CW tones: 400–1000 Hz every 25 Hz. Deliberately NOT lower: strong
	// low-frequency noise/hum (pink/red noise rises toward DC) would otherwise win
	// the argmax and lock the decoder onto ~250 Hz junk instead of the signal.
	for f := 400.0; f <= 1000.0; f += 25 {
		d.freqs = append(d.freqs, f)
		d.coeffs = append(d.coeffs, 2*math.Cos(2*math.Pi*f/float64(d.fs)))
	}
	d.mags = make([]float64, len(d.freqs))
	d.nbuf = make([]float64, len(d.freqs))
	return d
}

// SetTarget fixes the decode pitch to hz (lock to the nearest bin, ignore other
// tones), or returns to auto-search when hz <= 0. Safe to call concurrently.
func (d *Decoder) SetTarget(hz int) { d.targetHz.Store(int32(hz)) }

// nearestBin returns the bin index closest to hz.
func (d *Decoder) nearestBin(hz float64) int {
	best, bestD := 0, math.Inf(1)
	for i, f := range d.freqs {
		if dd := math.Abs(f - hz); dd < bestD {
			bestD, best = dd, i
		}
	}
	return best
}

// Reset clears decode state (e.g. when the user re-arms the decoder).
func (d *Decoder) Reset() {
	d.ring = d.ring[:0]
	d.acc = 0
	d.lockIdx, d.candIdx, d.candHops, d.quietHops = -1, -1, 0, 0
	d.peak, d.floor = 0, 0
	d.state = false
	d.stateHops = 0
	d.dotHops = 15
	d.elem = d.elem[:0]
	d.charEmitted, d.wordEmitted = false, false
}

// Process feeds a block of mono samples through the decoder.
func (d *Decoder) Process(samples []int16) {
	for _, s := range samples {
		d.ring = append(d.ring, float64(s))
		if len(d.ring) > d.win {
			d.ring = d.ring[len(d.ring)-d.win:]
		}
		d.acc++
		if d.acc >= d.hop && len(d.ring) >= d.win {
			d.acc = 0
			d.analyze()
			d.step()
		}
	}
}

// analyze runs the Goertzel bank, estimates the noise floor, and maintains the
// pitch lock (which tone the envelope detector then follows).
func (d *Decoder) analyze() {
	n := float64(len(d.ring))
	var sumSq float64
	maxIdx, maxMag := 0, -1.0
	for i, coeff := range d.coeffs {
		var s1, s2 float64
		for _, x := range d.ring {
			s0 := x + coeff*s1 - s2
			s2 = s1
			s1 = s0
		}
		m := math.Sqrt(math.Max(s1*s1+s2*s2-coeff*s1*s2, 0)) / n
		d.mags[i] = m
		if m > maxMag {
			maxMag = m
			maxIdx = i
		}
	}
	for _, x := range d.ring {
		sumSq += x * x
	}
	d.lastRMS = math.Min(1, math.Sqrt(sumSq/n)/32768*4)

	// Fixed-pitch mode: lock straight to the target bin, skip the auto search.
	// A narrow filter at the known pitch is exactly how a skimmer avoids QRM.
	if th := int(d.targetHz.Load()); th > 0 {
		d.lockIdx = d.nearestBin(float64(th))
		d.lastPitch = d.freqs[d.lockIdx]
		return
	}

	// Noise floor = 40th percentile of the bins (robust to a few strong tones).
	copy(d.nbuf, d.mags)
	sort.Float64s(d.nbuf)
	d.noise = d.nbuf[int(0.4*float64(len(d.nbuf)-1)+0.5)]

	if d.lockIdx < 0 {
		if maxIdx == d.candIdx {
			d.candHops++
		} else {
			d.candIdx, d.candHops = maxIdx, 1
		}
		snr := maxMag / (d.noise + 1e-9)
		// Tiered acquisition: a clearly strong tone locks on the FIRST hop (so we
		// don't eat the first element of a strong signal), a marginal/weak tone
		// locks after a couple of stable hops (so we don't lock onto pure noise).
		if snr > d.strongSNR || (d.candHops >= 3 && snr > d.acqSNR) {
			d.lockIdx = maxIdx
			d.peak, d.floor = maxMag, d.noise // seed the envelope to this bin
			d.quietHops = 0
		}
	}
	if d.lockIdx >= 0 {
		d.lastPitch = d.freqs[d.lockIdx]
	} else {
		d.lastPitch = 0
	}
}

// step runs the adaptive envelope on the locked bin and the timing state
// machine, one hop. The envelope adapts to the signal level (not an absolute
// threshold), so weak and strong signals both key correctly.
func (d *Decoder) step() {
	on := false
	if d.lockIdx >= 0 {
		m := d.mags[d.lockIdx]
		// Peak: fast attack, slow release.
		if m > d.peak {
			d.peak += (m - d.peak) * 0.4
		} else {
			d.peak += (m - d.peak) * 0.02
		}
		// Floor: drops fast toward the signal, but only RISES between marks (when
		// keyed up). Letting the floor rise during a long dash would shrink the
		// span until the dash drops below the threshold and fragments into dots —
		// the cause of the "all dots" garbage on a strong clean signal.
		if m < d.floor {
			d.floor += (m - d.floor) * 0.4
		} else if !d.state {
			d.floor += (m - d.floor) * 0.02
		}
		span := d.peak - d.floor
		// The frozen floor already stops dashes fragmenting, so keep balanced
		// thresholds: low enough that short inter-element GAPS are still seen
		// (otherwise elements merge into >7-symbol runs that decode to nothing).
		if span > d.floor*0.3+1e-9 {
			onTh := d.floor + 0.55*span
			offTh := d.floor + 0.35*span
			if d.state {
				on = m > offTh
			} else {
				on = m > onTh
			}
		}
		// Release the lock after a long quiet so we can retune to a new signal.
		if on {
			d.quietHops = 0
		} else {
			d.quietHops++
			if d.quietHops > d.relockHops {
				d.lockIdx, d.candIdx, d.candHops = -1, -1, 0
			}
		}
	}

	if on == d.state {
		d.stateHops++
		if !d.state {
			d.spaceProgress()
		}
	} else {
		if d.state {
			d.endMark(d.stateHops)
		}
		d.state = on
		d.stateHops = 1
		if on {
			d.charEmitted, d.wordEmitted = false, false
		}
	}
	d.emitStatus(on)
}

// endMark classifies a finished key-down run as a dot or dash and adapts the
// dot-length estimate. Runs shorter than a third of a dot are rejected as
// clicks/noise.
func (d *Decoder) endMark(hops int) {
	h := float64(hops)
	// Reject clicks/noise: shorter than a third of a dot AND an absolute floor
	// of ~4 hops (~16 ms, i.e. faster than ~75 WPM) so noise can't drag the
	// dot-length estimate down to the clamp (which produced 100 WPM garbage).
	if h < d.dotHops*0.35 || h < 4 {
		return
	}
	if h > d.dotHops*2 {
		d.elem = append(d.elem, '-')
		d.adaptDot(h / 3)
	} else {
		d.elem = append(d.elem, '.')
		d.adaptDot(h)
	}
}

// adaptDot nudges the dot-length estimate toward an observation (EMA, clamped
// to ~5–100 WPM).
func (d *Decoder) adaptDot(obs float64) {
	d.dotHops = d.dotHops*0.8 + obs*0.2 // slower: a few odd marks can't yank it
	if d.dotHops < 5 {                  // 5 hops ≈ 60 WPM ceiling — never 100
		d.dotHops = 5
	}
	if d.dotHops > 55 {
		d.dotHops = 55
	}
}

// spaceProgress flushes the current character once the gap exceeds a character
// gap, and a word space once it exceeds a word gap.
func (d *Decoder) spaceProgress() {
	g := float64(d.stateHops)
	if !d.charEmitted && g > d.dotHops*2 {
		d.flushChar()
		d.charEmitted = true
	}
	if !d.wordEmitted && g > d.dotHops*5 {
		if d.onChar != nil {
			d.onChar(" ")
		}
		d.wordEmitted = true
	}
}

// flushChar looks up the accumulated element string and emits the character.
func (d *Decoder) flushChar() {
	if len(d.elem) == 0 {
		return
	}
	if c, ok := morse[string(d.elem)]; ok {
		if d.onChar != nil {
			d.onChar(string(c))
		}
	} else if d.onChar != nil && len(d.elem) <= 7 {
		// Only flag a genuinely Morse-shaped but unknown char with "?". An
		// over-long element run is noise — drop it silently rather than spam "?".
		d.onChar("?")
	}
	d.elem = d.elem[:0]
}

func (d *Decoder) emitStatus(on bool) {
	d.sinceStatus++
	if d.sinceStatus < d.statusEvery || d.onStatus == nil {
		return
	}
	d.sinceStatus = 0
	hopMs := float64(d.hop) / float64(d.fs) * 1000
	wpm := 0
	if d.dotHops > 0 {
		wpm = int(math.Round(1200 / (d.dotHops * hopMs)))
	}
	d.onStatus(Status{WPM: wpm, Pitch: int(math.Round(d.lastPitch)), Level: d.lastRMS, Active: on})
}