//! waterfall.rs — Dual-CSS HF Mode
//! Real-time scrolling spectrogram rendered as an egui texture.
//!
//! WSJT-X / JTDX style: the waterfall shows the full configured audio band
//! (display.waterfall_freq_min..max), independent of the modem passband, so you
//! always see the whole spectrum rather than a zoom on the modem signal. The
//! history buffer scrolls top-down; each new spectrum row is inserted at the top
//! after rolling the buffer down (by `scroll_speed` rows).
//!
//! The spectrum length is variable (FFT size is runtime-adjustable), so rows are
//! stored at whatever resolution the audio worker produced and mapped to display
//! columns by frequency at paint time.

use egui::{Color32, ColorImage, FontId, Pos2, Rect, Sense, Stroke, TextureHandle, TextureOptions};

use crate::config::AppConfig;

pub const HISTORY: usize = 300; // time rows kept in scrolling buffer
pub const DB_RANGE: f32 = 70.0; // dB window displayed

/// Colormap palettes selectable from the Display settings.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum ColorScheme {
    Default, // black → navy → cyan → yellow → white
    Hot,     // black → red → yellow → white
    Gray,    // black → white
}

impl ColorScheme {
    pub fn from_key(s: &str) -> Self {
        match s {
            "hot" => ColorScheme::Hot,
            "gray" | "grey" => ColorScheme::Gray,
            _ => ColorScheme::Default,
        }
    }

    fn key(&self) -> &'static str {
        match self {
            ColorScheme::Default => "default",
            ColorScheme::Hot => "hot",
            ColorScheme::Gray => "gray",
        }
    }
}

/// Build a 256-entry colormap for the given scheme.
fn build_colormap(scheme: ColorScheme) -> Vec<[u8; 3]> {
    let n = 256;
    let mut cmap = vec![[0u8; 3]; n];
    for (i, slot) in cmap.iter_mut().enumerate() {
        let t = i as f32 / (n as f32 - 1.0);
        *slot = match scheme {
            ColorScheme::Default => default_ramp(t),
            ColorScheme::Hot => hot_ramp(t),
            ColorScheme::Gray => {
                let v = (t * 255.0) as u8;
                [v, v, v]
            }
        };
    }
    cmap
}

/// Default ramp: black → navy → cyan → yellow → white (verbatim from waterfall.py).
fn default_ramp(t: f32) -> [u8; 3] {
    let (r, g, b);
    if t < 0.12 {
        r = 0.0;
        g = 0.0;
        b = t / 0.12 * 45.0;
    } else if t < 0.35 {
        let p = (t - 0.12) / 0.23;
        r = 0.0;
        g = p * 35.0;
        b = 45.0 + p * 115.0;
    } else if t < 0.62 {
        let p = (t - 0.35) / 0.27;
        r = p * 90.0;
        g = 35.0 + p * 195.0;
        b = 160.0 + p * 40.0;
    } else if t < 0.82 {
        let p = (t - 0.62) / 0.20;
        r = 90.0 + p * 165.0;
        g = 255.0;
        b = 200.0 - p * 200.0;
    } else {
        let p = (t - 0.82) / 0.18;
        r = 255.0;
        g = 255.0 - p * 30.0;
        b = p * 220.0;
    }
    [r as u8, g as u8, b as u8]
}

/// Classic "hot" ramp: black → red → yellow → white.
fn hot_ramp(t: f32) -> [u8; 3] {
    let r = (t / 0.4).clamp(0.0, 1.0);
    let g = ((t - 0.3) / 0.4).clamp(0.0, 1.0);
    let b = ((t - 0.7) / 0.3).clamp(0.0, 1.0);
    [(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8]
}

pub struct Waterfall {
    history: Vec<Vec<f32>>, // HISTORY rows, each a full dBFS spectrum (variable len)
    colormap: Vec<[u8; 3]>,
    scheme: ColorScheme,
    texture: Option<TextureHandle>,
    spec_len: usize, // most recent spectrum length (bins), for axis mapping
}

impl Default for Waterfall {
    fn default() -> Self {
        Self::new()
    }
}

impl Waterfall {
    pub fn new() -> Self {
        Self {
            history: vec![Vec::new(); HISTORY],
            colormap: build_colormap(ColorScheme::Default),
            scheme: ColorScheme::Default,
            texture: None,
            spec_len: 0,
        }
    }

    /// Insert a new spectrum row at the top, scrolling the rest down by one.
    pub fn push_spectrum(&mut self, spectrum_db: &[f32]) {
        self.spec_len = spectrum_db.len();
        self.history.pop();
        self.history.insert(0, spectrum_db.to_vec());
    }

    /// Clear history. No live caller since config changes stopped blanking
    /// the display (old rows just scroll off); kept for future use.
    #[allow(dead_code)]
    pub fn clear(&mut self) {
        for row in self.history.iter_mut() {
            row.clear();
        }
    }

    fn ensure_colormap(&mut self, scheme: ColorScheme) {
        if scheme != self.scheme {
            self.scheme = scheme;
            self.colormap = build_colormap(scheme);
        }
    }

    /// Map a frequency (Hz) to a spectrum bin index for the current spec length.
    /// Bins span 0..fs/2 across `spec_len` entries (spec_len = fft_size/2).
    fn freq_to_bin(&self, freq: u32, fs: u32) -> usize {
        if fs == 0 || self.spec_len == 0 {
            return 0;
        }
        // bin = freq / (fs/2) * spec_len
        let frac = freq as f32 / (fs as f32 / 2.0);
        ((frac * self.spec_len as f32) as usize).min(self.spec_len.saturating_sub(1))
    }

    /// Draw the waterfall, filling the given UI region.
    pub fn ui(&mut self, ui: &mut egui::Ui, cfg: &AppConfig) {
        self.ensure_colormap(ColorScheme::from_key(&cfg.display.color_scheme));

        let (rect, _response) = ui.allocate_exact_size(ui.available_size(), Sense::hover());
        let w = rect.width().max(1.0) as usize;
        let h = rect.height().max(1.0) as usize;
        if w == 0 || h == 0 {
            return;
        }

        let fs = cfg.audio.sample_rate;
        // WSJT-X-style: display the whole configured band, clamped to Nyquist.
        // Keep f_lo strictly below Nyquist so the clamp below can't invert
        // (std clamp panics when min > max).
        let nyquist = (fs / 2).max(2);
        let f_lo = cfg.display.waterfall_freq_min.min(nyquist - 1);
        let f_hi = cfg.display.waterfall_freq_max.clamp(f_lo + 1, nyquist);

        let b_lo = self.freq_to_bin(f_lo, fs);
        let b_hi = self.freq_to_bin(f_hi, fs).max(b_lo + 1);
        let n_band = b_hi - b_lo;

        let n_rows = HISTORY.min(h);
        let floor = cfg.display.waterfall_floor;
        let d_range = DB_RANGE * cfg.display.waterfall_gain;
        let d_range = if d_range.abs() < 1e-6 { 1e-6 } else { d_range };

        // Map each output column to a band-relative bin (handles n_band ≷ w).
        let col_idx: Vec<usize> = (0..w)
            .map(|x| {
                if w == 1 {
                    0
                } else {
                    (x as f32 / (w as f32 - 1.0) * (n_band as f32 - 1.0)).round() as usize
                }
            })
            .collect();

        let mut pixels = vec![Color32::BLACK; w * n_rows];
        for row in 0..n_rows {
            let data_row = &self.history[row];
            for x in 0..w {
                let bin = b_lo + col_idx[x];
                let db = data_row.get(bin).copied().unwrap_or(floor);
                let idx = (((db - floor) / d_range) * 255.0).clamp(0.0, 255.0) as usize;
                let [r, g, b] = self.colormap[idx.min(255)];
                pixels[row * w + x] = Color32::from_rgb(r, g, b);
            }
        }

        let image = ColorImage {
            size: [w, n_rows],
            pixels,
            source_size: egui::Vec2::new(w as f32, n_rows as f32),
        };

        // Texture id includes the scheme so a palette switch forces a refresh.
        let tex_name = format!("waterfall_{}", self.scheme.key());
        let texture = self.texture.get_or_insert_with(|| {
            ui.ctx()
                .load_texture(&tex_name, ColorImage::example(), TextureOptions::LINEAR)
        });
        texture.set(image, TextureOptions::LINEAR);

        let painter = ui.painter_at(rect);
        painter.image(
            texture.id(),
            rect,
            Rect::from_min_max(Pos2::ZERO, Pos2::new(1.0, 1.0)),
            Color32::WHITE,
        );

        self.draw_freq_axis(&painter, rect, f_lo, f_hi);
    }

    fn draw_freq_axis(&self, painter: &egui::Painter, rect: Rect, f_lo: u32, f_hi: u32) {
        let f_lo = f_lo as i32;
        let f_hi = f_hi as i32;
        let f_span = (f_hi - f_lo).max(1) as f32;
        let w = rect.width();

        let tick_step = if f_span <= 1000.0 {
            100
        } else if f_span <= 3000.0 {
            500
        } else {
            1000
        };
        let first_tick = ((f_lo / tick_step) + 1) * tick_step;

        let font = FontId::monospace(9.0);
        let tick_color = Color32::from_rgba_unmultiplied(0, 212, 255, 110);
        let line_color = Color32::from_rgba_unmultiplied(0, 212, 255, 80);

        let mut freq = first_tick;
        while freq < f_hi {
            let x = rect.left() + (freq - f_lo) as f32 / f_span * w;
            if x > rect.left() && x < rect.right() - 1.0 {
                painter.line_segment(
                    [
                        Pos2::new(x, rect.bottom() - 18.0),
                        Pos2::new(x, rect.bottom() - 10.0),
                    ],
                    Stroke::new(1.0, line_color),
                );
                painter.text(
                    Pos2::new(x - 14.0, rect.bottom() - 12.0),
                    egui::Align2::LEFT_TOP,
                    format!("{freq}"),
                    font.clone(),
                    tick_color,
                );
            }
            freq += tick_step;
        }

        // Band-edge labels.
        let edge = Color32::from_rgba_unmultiplied(0, 212, 255, 65);
        painter.text(
            Pos2::new(rect.left() + 2.0, rect.bottom() - 12.0),
            egui::Align2::LEFT_TOP,
            format!("{f_lo}"),
            font.clone(),
            edge,
        );
        painter.text(
            Pos2::new(rect.right() - 32.0, rect.bottom() - 12.0),
            egui::Align2::LEFT_TOP,
            format!("{f_hi}"),
            font,
            edge,
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn default_colormap_endpoints() {
        let cmap = build_colormap(ColorScheme::Default);
        assert_eq!(cmap.len(), 256);
        assert_eq!(cmap[0], [0, 0, 0]);
        let [r, g, b] = cmap[255];
        assert!(r > 250 && g > 200 && b > 200, "got {:?}", cmap[255]);
    }

    #[test]
    fn hot_and_gray_endpoints() {
        let hot = build_colormap(ColorScheme::Hot);
        assert_eq!(hot[0], [0, 0, 0]);
        assert_eq!(hot[255], [255, 255, 255]);
        let gray = build_colormap(ColorScheme::Gray);
        assert_eq!(gray[0], [0, 0, 0]);
        assert_eq!(gray[255], [255, 255, 255]);
        // Gray is monochrome: r == g == b at every step.
        assert!(gray.iter().all(|[r, g, b]| r == g && g == b));
    }

    #[test]
    fn scheme_from_str() {
        assert!(matches!(ColorScheme::from_key("hot"), ColorScheme::Hot));
        assert!(matches!(ColorScheme::from_key("gray"), ColorScheme::Gray));
        assert!(matches!(ColorScheme::from_key("grey"), ColorScheme::Gray));
        assert!(matches!(
            ColorScheme::from_key("default"),
            ColorScheme::Default
        ));
        assert!(matches!(
            ColorScheme::from_key("nonsense"),
            ColorScheme::Default
        ));
    }

    #[test]
    fn push_keeps_history_bounded() {
        let mut wf = Waterfall::new();
        for _ in 0..(HISTORY + 50) {
            wf.push_spectrum(&vec![-50.0; 1024]);
        }
        assert_eq!(wf.history.len(), HISTORY);
        assert!((wf.history[0][10] - (-50.0)).abs() < 1e-6);
        assert_eq!(wf.spec_len, 1024);
    }

    #[test]
    fn band_clamp_cannot_invert() {
        // Regression: waterfall_freq_min ≥ Nyquist used to make the f_hi
        // clamp panic (min > max). Mirror the ui() band math here.
        let fs: u32 = 48000;
        let nyquist = (fs / 2).max(2);
        for freq_min in [0u32, 23_999, 24_000, 50_000] {
            let f_lo = freq_min.min(nyquist - 1);
            let f_hi = 3500u32.clamp(f_lo + 1, nyquist); // must not panic
            assert!(f_hi > f_lo);
        }
    }

    #[test]
    fn freq_to_bin_maps_nyquist() {
        let mut wf = Waterfall::new();
        wf.push_spectrum(&vec![0.0; 1024]); // fft_size 2048 → 1024 bins
        let fs = 48000;
        // 0 Hz → bin 0; near Nyquist (24 kHz) → last bin.
        assert_eq!(wf.freq_to_bin(0, fs), 0);
        assert_eq!(wf.freq_to_bin(24000, fs), 1023);
        // 12 kHz is half of Nyquist → ~bin 512.
        let mid = wf.freq_to_bin(12000, fs);
        assert!((mid as i64 - 512).abs() <= 1, "got {mid}");
    }
}
