// 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 (auto-picking the dominant pitch), an adaptive
// envelope/threshold to recover key-down/key-up, 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. Robustness on weak/QRM/QSB signals is limited
// (as with every audio CW decoder); it does well on clean signals.
package cwdecode

import "math"

// 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
	Level  float64 `json:"level"`  // 0..1 rough signal strength (SNR proxy)
	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 envelope 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

	// Adaptive envelope (relative, so absolute gain is irrelevant).
	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   // current space already flushed a character
	wordEmitted  bool   // current space already flushed a word gap
	lastPitch    float64
	lastRMS      float64 // 0..1 input level of the current window (for the UI meter)

	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 envelope resolution
		win:         sampleRate / 62,   // ~16 ms Goertzel window
		dotHops:     15,                // ~20 WPM seed (15 hops * 4 ms = 60 ms)
		statusEvery: 25,                // ~10 Hz status
		onChar:      onChar,
		onStatus:    onStatus,
	}
	if d.hop < 1 {
		d.hop = 1
	}
	// Candidate CW tones: 250–1200 Hz every 25 Hz (wide enough for most rigs'
	// audio offset). The dominant bin is the pitch (auto), and its magnitude
	// drives the envelope.
	for f := 250.0; f <= 1200.0; f += 25 {
		d.freqs = append(d.freqs, f)
		d.coeffs = append(d.coeffs, 2*math.Cos(2*math.Pi*f/float64(d.fs)))
	}
	return d
}

// 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.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
			mag, pitch := d.toneMag()
			d.step(mag, pitch)
		}
	}
}

// toneMag runs the Goertzel bank over the current window and returns the
// strongest bin's magnitude and its frequency (the auto-detected pitch).
func (d *Decoder) toneMag() (float64, float64) {
	best, bestF := 0.0, d.lastPitch
	n := float64(len(d.ring))
	var sumSq float64
	for i, coeff := range d.coeffs {
		var s1, s2 float64
		for _, x := range d.ring {
			s0 := x + coeff*s1 - s2
			s2 = s1
			s1 = s0
		}
		power := s1*s1 + s2*s2 - coeff*s1*s2
		if power > best {
			best = power
			bestF = d.freqs[i]
		}
	}
	for _, x := range d.ring {
		sumSq += x * x
	}
	d.lastRMS = math.Min(1, math.Sqrt(sumSq/n)/32768*4) // ×4 so quiet audio is visible
	// Normalise by window length so the magnitude scale is rate-independent.
	return math.Sqrt(math.Max(best, 0)) / n, bestF
}

// step advances the envelope follower + timing state machine by one hop.
func (d *Decoder) step(mag, pitch float64) {
	// Envelope: fast attack / slow release for the peak, fast drop / slow rise
	// for the noise floor. Tracks the signal even through QSB.
	if mag > d.peak {
		d.peak += (mag - d.peak) * 0.4
	} else {
		d.peak += (mag - d.peak) * 0.02
	}
	if mag < d.floor {
		d.floor += (mag - d.floor) * 0.4
	} else {
		d.floor += (mag - d.floor) * 0.01
	}

	span := d.peak - d.floor
	// Hysteresis thresholds; require a minimum SNR span to call anything a tone.
	on := d.state
	if span > d.floor*0.3+1e-9 {
		onTh := d.floor + 0.55*span
		offTh := d.floor + 0.35*span
		if d.state {
			on = mag > offTh
		} else {
			on = mag > onTh
		}
		if on {
			d.lastPitch = pitch
		}
	} else {
		on = false
	}

	if on == d.state {
		d.stateHops++
		if !d.state {
			d.spaceProgress() // flush char/word as the gap grows
		}
	} else {
		if d.state {
			d.endMark(d.stateHops)
		}
		d.state = on
		d.stateHops = 1
		if on {
			// A new mark starts → the previous space is over; re-arm flushing.
			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.
func (d *Decoder) endMark(hops int) {
	h := float64(hops)
	// Reject impulse noise far shorter than a dot.
	if h < d.dotHops*0.35 {
		return
	}
	dash := h > d.dotHops*2
	if dash {
		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.7 + obs*0.3
	if d.dotHops < 3 {
		d.dotHops = 3
	}
	if d.dotHops > 60 {
		d.dotHops = 60
	}
}

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