mendi 0.0.2

//! Real-time fNIRS viewer for the Mendi headband.
//!
//! ```bash
//! cargo run --features tui --bin mendi-tui
//! cargo run --features tui --bin mendi-tui -- --simulate
//! ```
//!
//! Keys:
//!   q / Esc / Ctrl-C  quit
//!   p                  pause
//!   r                  resume
//!   c                  clear buffers
//!   v                  toggle smooth overlay
//!   +/-                zoom Y axis
//!   a                  auto-scale Y axis
//!   1                  optical view (IR channels)
//!   2                  raw view (all 9 optical values)

use std::collections::VecDeque;
use std::io;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use anyhow::Result;
use crossterm::{
    event::{self, Event, KeyCode, KeyModifiers},
    execute,
    terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
};
use ratatui::{
    backend::CrosstermBackend,
    layout::{Constraint, Layout, Rect},
    style::{Color, Modifier, Style},
    symbols,
    text::{Line, Span},
    widgets::{Axis, Block, Borders, Chart, Dataset, GraphType, Paragraph},
    Frame, Terminal,
};

use mendi::simulate::{SimConfig, SimulatedDevice};
use mendi::types::*;

// ── Constants ─────────────────────────────────────────────────────────────────

const WINDOW_SECS: f64 = 4.0;
const SAMPLE_HZ: f64 = 25.0;
const BUF_SIZE: usize = (WINDOW_SECS * SAMPLE_HZ) as usize;
const SMOOTH_WINDOW: usize = 5;
const SPINNER: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];

const Y_SCALES: &[f64] = &[
    500.0, 1000.0, 2000.0, 5000.0, 10000.0, 20000.0, 50000.0, 100000.0,
];
const DEFAULT_SCALE: usize = 4; // ±10000

// ── View mode ─────────────────────────────────────────────────────────────────

#[derive(Clone, Copy, PartialEq, Eq)]
enum ViewMode {
    /// 3 charts: IR left, IR right, IR pulse (the main neurofeedback signal)
    Ir,
    /// 3 charts: left (IR+red+amb), right, pulse
    Full,
}

// ── App state ─────────────────────────────────────────────────────────────────

struct App {
    // 9 optical channels: ir_l, red_l, amb_l, ir_r, red_r, amb_r, ir_p, red_p, amb_p
    bufs: [VecDeque<f64>; 9],
    // IMU
    acc: (i32, i32, i32),
    gyro: (i32, i32, i32),
    temperature: f32,
    // Battery
    battery: Option<BatteryReading>,
    // Calibration
    calibration: Option<CalibrationReading>,
    // Device
    device_info: Option<DeviceInfo>,
    diagnostics: Option<DiagnosticsReading>,
    // UI
    mode: AppMode,
    view: ViewMode,
    scale_idx: usize,
    paused: bool,
    smooth: bool,
    frame_count: u64,
    pkt_times: VecDeque<Instant>,
}

#[derive(Clone)]
enum AppMode {
    Scanning,
    Connected(String),
    Simulated,
    Disconnected,
}

impl App {
    fn new() -> Self {
        Self {
            bufs: std::array::from_fn(|_| VecDeque::with_capacity(BUF_SIZE + 4)),
            acc: (0, 0, 0),
            gyro: (0, 0, 0),
            temperature: 0.0,
            battery: None,
            calibration: None,
            device_info: None,
            diagnostics: None,
            mode: AppMode::Scanning,
            view: ViewMode::Ir,
            scale_idx: DEFAULT_SCALE,
            paused: false,
            smooth: true,
            frame_count: 0,
            pkt_times: VecDeque::with_capacity(128),
        }
    }

    fn push_frame(&mut self, f: &FrameReading) {
        if self.paused {
            return;
        }
        let vals = [
            f.ir_left as f64,
            f.red_left as f64,
            f.amb_left as f64,
            f.ir_right as f64,
            f.red_right as f64,
            f.amb_right as f64,
            f.ir_pulse as f64,
            f.red_pulse as f64,
            f.amb_pulse as f64,
        ];
        for (i, &v) in vals.iter().enumerate() {
            self.bufs[i].push_back(v);
            while self.bufs[i].len() > BUF_SIZE {
                self.bufs[i].pop_front();
            }
        }
        self.acc = (f.acc_x, f.acc_y, f.acc_z);
        self.gyro = (f.ang_x, f.ang_y, f.ang_z);
        self.temperature = f.temperature;
        self.frame_count += 1;

        let now = Instant::now();
        self.pkt_times.push_back(now);
        while self
            .pkt_times
            .front()
            .is_some_and(|t| now.duration_since(*t) > Duration::from_secs(2))
        {
            self.pkt_times.pop_front();
        }
    }

    fn clear(&mut self) {
        for b in &mut self.bufs {
            b.clear();
        }
        self.pkt_times.clear();
        self.frame_count = 0;
    }

    fn pkt_rate(&self) -> f64 {
        let n = self.pkt_times.len();
        if n < 2 {
            return 0.0;
        }
        let span = self.pkt_times.back().unwrap().duration_since(self.pkt_times[0]).as_secs_f64();
        if span < 1e-9 { 0.0 } else { (n as f64 - 1.0) / span }
    }

    fn y_range(&self) -> f64 {
        Y_SCALES[self.scale_idx]
    }
    fn scale_up(&mut self) {
        if self.scale_idx + 1 < Y_SCALES.len() {
            self.scale_idx += 1;
        }
    }
    fn scale_down(&mut self) {
        if self.scale_idx > 0 {
            self.scale_idx -= 1;
        }
    }
    fn auto_scale(&mut self) {
        let peak = self
            .bufs
            .iter()
            .flat_map(|b| b.iter())
            .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
        let needed = peak * 1.1;
        self.scale_idx = Y_SCALES
            .iter()
            .position(|&s| s >= needed)
            .unwrap_or(Y_SCALES.len() - 1);
    }
}

// ── Rendering ─────────────────────────────────────────────────────────────────

fn spinner_str() -> &'static str {
    let ms = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis();
    SPINNER[(ms / 100) as usize % SPINNER.len()]
}

fn sep<'a>() -> Span<'a> {
    Span::styled(" │ ", Style::default().fg(Color::DarkGray))
}

fn key(s: &str) -> Span<'_> {
    Span::styled(s, Style::default().fg(Color::Yellow).add_modifier(Modifier::BOLD))
}

fn draw(frame: &mut Frame, app: &App) {
    let root = Layout::vertical([
        Constraint::Length(3),
        Constraint::Min(0),
        Constraint::Length(3),
    ])
    .split(frame.area());

    draw_header(frame, root[0], app);
    draw_charts(frame, root[1], app);
    draw_footer(frame, root[2], app);
}

fn draw_header(frame: &mut Frame, area: Rect, app: &App) {
    let (label, color) = match &app.mode {
        AppMode::Scanning => (format!("{} Scanning…", spinner_str()), Color::Yellow),
        AppMode::Connected(name) => (format!("● {name}"), Color::Green),
        AppMode::Simulated => ("◆ Simulated".to_owned(), Color::Cyan),
        AppMode::Disconnected => (
            format!("{} Disconnected — retrying…", spinner_str()),
            Color::Red,
        ),
    };

    let bat = app
        .battery
        .as_ref()
        .map(|b| format!("Bat {:.1}V {}%", b.voltage(), b.percentage()))
        .unwrap_or_else(|| "Bat N/A".into());

    let rate = format!("{:.1} Hz", app.pkt_rate());
    let scale = format!("±{:.0}", app.y_range());
    let temp = format!("{:.1}°C", app.temperature);
    let frames = format!("{}K frm", app.frame_count / 1000);
    let view_label = match app.view {
        ViewMode::Ir => "IR",
        ViewMode::Full => "FULL",
    };

    let line = Line::from(vec![
        Span::styled(
            " MENDI fNIRS ",
            Style::default().fg(Color::White).add_modifier(Modifier::BOLD),
        ),
        sep(),
        Span::styled(label, Style::default().fg(color).add_modifier(Modifier::BOLD)),
        sep(),
        Span::styled(view_label, Style::default().fg(Color::LightYellow).add_modifier(Modifier::BOLD)),
        sep(),
        Span::styled(bat, Style::default().fg(Color::White)),
        sep(),
        Span::styled(temp, Style::default().fg(Color::White)),
        sep(),
        Span::styled(rate, Style::default().fg(Color::White)),
        sep(),
        Span::styled(scale, Style::default().fg(Color::LightBlue).add_modifier(Modifier::BOLD)),
        sep(),
        Span::styled(frames, Style::default().fg(Color::DarkGray)),
    ]);

    frame.render_widget(
        Paragraph::new(line).block(Block::default().borders(Borders::ALL)),
        area,
    );
}

fn smooth_signal(data: &[(f64, f64)], window: usize) -> Vec<(f64, f64)> {
    if data.len() < 3 || window < 2 {
        return data.to_vec();
    }
    let half = window / 2;
    data.iter()
        .enumerate()
        .map(|(i, &(x, _))| {
            let start = i.saturating_sub(half);
            let end = (i + half + 1).min(data.len());
            let sum: f64 = data[start..end].iter().map(|&(_, y)| y).sum();
            (x, sum / (end - start) as f64)
        })
        .collect()
}

fn buf_to_data(buf: &VecDeque<f64>) -> Vec<(f64, f64)> {
    buf.iter()
        .enumerate()
        .map(|(i, &v)| (i as f64 / SAMPLE_HZ, v))
        .collect()
}

#[allow(clippy::too_many_arguments)]
fn draw_channel(
    frame: &mut Frame,
    area: Rect,
    title: &str,
    data: &[(f64, f64)],
    color: Color,
    dim_color: Color,
    y_min: f64,
    y_max: f64,
    smooth: bool,
) {
    let smoothed: Vec<(f64, f64)> = if smooth {
        smooth_signal(data, SMOOTH_WINDOW)
    } else {
        vec![]
    };

    let datasets: Vec<Dataset> = if smooth {
        vec![
            Dataset::default()
                .marker(symbols::Marker::Braille)
                .graph_type(GraphType::Line)
                .style(Style::default().fg(dim_color))
                .data(data),
            Dataset::default()
                .marker(symbols::Marker::Braille)
                .graph_type(GraphType::Line)
                .style(Style::default().fg(color))
                .data(&smoothed),
        ]
    } else {
        vec![Dataset::default()
            .marker(symbols::Marker::Braille)
            .graph_type(GraphType::Line)
            .style(Style::default().fg(color))
            .data(data)]
    };

    let y_mid = (y_min + y_max) / 2.0;
    let y_labels: Vec<String> = vec![
        format!("{:.0}", y_min),
        format!("{:.0}", y_mid),
        format!("{:.0}", y_max),
    ];
    let x_labels = vec![
        "0s".into(),
        format!("{:.0}s", WINDOW_SECS / 2.0),
        format!("{:.0}s", WINDOW_SECS),
    ];

    let smooth_tag = if smooth { " [S]" } else { "" };

    let chart = Chart::new(datasets)
        .block(
            Block::default()
                .title(Span::styled(
                    format!(" {title}{smooth_tag} "),
                    Style::default().fg(color).add_modifier(Modifier::BOLD),
                ))
                .borders(Borders::ALL)
                .border_style(Style::default().fg(color)),
        )
        .x_axis(
            Axis::default()
                .bounds([0.0, WINDOW_SECS])
                .labels(x_labels)
                .style(Style::default().fg(Color::DarkGray)),
        )
        .y_axis(
            Axis::default()
                .bounds([y_min, y_max])
                .labels(y_labels)
                .style(Style::default().fg(Color::DarkGray)),
        );
    frame.render_widget(chart, area);
}

fn draw_charts(frame: &mut Frame, area: Rect, app: &App) {
    match app.view {
        ViewMode::Ir => draw_ir_view(frame, area, app),
        ViewMode::Full => draw_full_view(frame, area, app),
    }
}

fn draw_ir_view(frame: &mut Frame, area: Rect, app: &App) {
    let rows = Layout::vertical([
        Constraint::Ratio(1, 3),
        Constraint::Ratio(1, 3),
        Constraint::Ratio(1, 3),
    ])
    .split(area);

    let yr = app.y_range();
    let channels = [
        (0, "IR Left", Color::Cyan, Color::Rgb(0, 60, 80)),
        (3, "IR Right", Color::Green, Color::Rgb(0, 60, 0)),
        (6, "IR Pulse", Color::Magenta, Color::Rgb(60, 0, 60)),
    ];

    for (idx, (buf_idx, name, color, dim)) in channels.iter().enumerate() {
        let data = buf_to_data(&app.bufs[*buf_idx]);
        draw_channel(frame, rows[idx], name, &data, *color, *dim, -yr, yr, app.smooth);
    }
}

fn draw_full_view(frame: &mut Frame, area: Rect, app: &App) {
    let rows = Layout::vertical([
        Constraint::Ratio(1, 3),
        Constraint::Ratio(1, 3),
        Constraint::Ratio(1, 3),
    ])
    .split(area);

    let yr = app.y_range();
    let groups: [(&str, [usize; 3], [Color; 3]); 3] = [
        ("Left", [0, 1, 2], [Color::Cyan, Color::Red, Color::Yellow]),
        ("Right", [3, 4, 5], [Color::Green, Color::LightRed, Color::LightYellow]),
        ("Pulse", [6, 7, 8], [Color::Magenta, Color::LightRed, Color::Yellow]),
    ];

    for (row_idx, (label, buf_idxs, colors)) in groups.iter().enumerate() {
        // Pre-build data vecs so they outlive the Dataset borrows
        let data_vecs: Vec<Vec<(f64, f64)>> = buf_idxs
            .iter()
            .map(|bi| {
                app.bufs[*bi]
                    .iter()
                    .enumerate()
                    .map(|(i, &v)| (i as f64 / SAMPLE_HZ, v.clamp(-yr, yr)))
                    .collect()
            })
            .collect();

        let datasets: Vec<Dataset> = data_vecs
            .iter()
            .zip(colors.iter())
            .map(|(data, col)| {
                Dataset::default()
                    .marker(symbols::Marker::Braille)
                    .graph_type(GraphType::Line)
                    .style(Style::default().fg(*col))
                    .data(data)
            })
            .collect();

        let y_labels: Vec<String> = vec![
            format!("{:.0}", -yr),
            "0".into(),
            format!("{:.0}", yr),
        ];

        let chart = Chart::new(datasets)
            .block(
                Block::default()
                    .title(Span::styled(
                        format!(" {label}: IR(c) Red(r) Amb(y) "),
                        Style::default()
                            .fg(colors[0])
                            .add_modifier(Modifier::BOLD),
                    ))
                    .borders(Borders::ALL),
            )
            .x_axis(
                Axis::default()
                    .bounds([0.0, WINDOW_SECS])
                    .labels(vec!["0s".into(), format!("{:.0}s", WINDOW_SECS)])
                    .style(Style::default().fg(Color::DarkGray)),
            )
            .y_axis(
                Axis::default()
                    .bounds([-yr, yr])
                    .labels(y_labels)
                    .style(Style::default().fg(Color::DarkGray)),
            );
        frame.render_widget(chart, rows[row_idx]);
    }
}

fn draw_footer(frame: &mut Frame, area: Rect, app: &App) {
    let pause_span = if app.paused {
        Span::styled("  ⏸ PAUSED", Style::default().fg(Color::Red).add_modifier(Modifier::BOLD))
    } else {
        Span::raw("")
    };

    let (ax, ay, az) = app.acc;
    let (gx, gy, gz) = app.gyro;

    let keys = Line::from(vec![
        Span::raw(" "),
        key("[1]"), Span::raw("IR  "),
        key("[2]"), Span::raw("Full  "),
        key("[+/-]"), Span::raw("Scale  "),
        key("[a]"), Span::raw("Auto  "),
        key("[v]"), Span::raw(if app.smooth { "Raw  " } else { "Smooth  " }),
        key("[p]"), Span::raw("Pause  "),
        key("[r]"), Span::raw("Resume  "),
        key("[c]"), Span::raw("Clear  "),
        key("[q]"), Span::raw("Quit"),
        pause_span,
    ]);

    let imu = Line::from(vec![
        Span::raw(" "),
        Span::styled("Acc ", Style::default().fg(Color::DarkGray)),
        Span::styled(
            format!("({ax:+6},{ay:+6},{az:+6})"),
            Style::default().fg(Color::Cyan),
        ),
        Span::raw("  "),
        Span::styled("Gyro ", Style::default().fg(Color::DarkGray)),
        Span::styled(
            format!("({gx:+6},{gy:+6},{gz:+6})"),
            Style::default().fg(Color::Magenta),
        ),
        Span::raw("  "),
        if let Some(c) = &app.calibration {
            Span::styled(
                format!("Cal L={:.1} R={:.1} P={:.1}", c.offset_left, c.offset_right, c.offset_pulse),
                Style::default().fg(Color::DarkGray),
            )
        } else {
            Span::raw("")
        },
    ]);

    frame.render_widget(
        Paragraph::new(vec![keys, imu]).block(Block::default().borders(Borders::ALL)),
        area,
    );
}

// ── Spawn event consumer ──────────────────────────────────────────────────────

fn spawn_event_task(mut rx: tokio::sync::mpsc::Receiver<MendiEvent>, app: Arc<Mutex<App>>) {
    tokio::spawn(async move {
        while let Some(ev) = rx.recv().await {
            let mut s = app.lock().unwrap();
            match ev {
                MendiEvent::Connected(info) => {
                    s.mode = AppMode::Connected(info.to_string());
                    s.device_info = Some(info);
                }
                MendiEvent::Disconnected => {
                    s.mode = AppMode::Disconnected;
                    break;
                }
                MendiEvent::Frame(f) => {
                    s.push_frame(&f);
                }
                MendiEvent::Battery(b) => {
                    s.battery = Some(b);
                }
                MendiEvent::Calibration(c) => {
                    s.calibration = Some(c);
                }
                MendiEvent::Diagnostics(d) => {
                    s.diagnostics = Some(d);
                }
                MendiEvent::SensorRead(_) => {}
            }
        }
    });
}

// ── Entry point ───────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() -> Result<()> {
    // Log to file
    {
        use std::fs::File;
        if let Ok(file) = File::create("mendi-tui.log") {
            env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
                .target(env_logger::Target::Pipe(Box::new(file)))
                .init();
        }
    }

    let simulate = std::env::args().any(|a| a == "--simulate");

    let app = Arc::new(Mutex::new(App::new()));

    // ── Start data source ─────────────────────────────────────────────────────
    let sim_handle = if simulate {
        let config = SimConfig {
            frame_rate_hz: SAMPLE_HZ,
            ..Default::default()
        };
        let (rx, handle) = SimulatedDevice::start(config);
        {
            let mut s = app.lock().unwrap();
            s.mode = AppMode::Simulated;
        }
        spawn_event_task(rx, Arc::clone(&app));
        Some(handle)
    } else {
        // Real BLE connection
        let (tx_conn, mut rx_conn) = tokio::sync::mpsc::channel::<tokio::sync::mpsc::Receiver<MendiEvent>>(1);

        tokio::spawn(async move {
            use mendi::mendi_client::{MendiClient, MendiClientConfig};
            loop {
                let client = MendiClient::new(MendiClientConfig::default());
                match client.connect().await {
                    Ok((rx, _handle)) => {
                        let _ = tx_conn.send(rx).await;
                        // Keep handle alive by leaking it (it's Send but not easily shared)
                        // The disconnect event will come through the channel
                        std::mem::forget(_handle);
                        break;
                    }
                    Err(e) => {
                        log::warn!("Connection failed: {e}");
                        tokio::time::sleep(Duration::from_secs(3)).await;
                    }
                }
            }
        });

        // Wait for connection in background, forward events when ready
        let app_clone2 = Arc::clone(&app);
        tokio::spawn(async move {
            if let Some(rx) = rx_conn.recv().await {
                spawn_event_task(rx, app_clone2);
            }
        });

        None
    };

    // ── Terminal setup ────────────────────────────────────────────────────────
    enable_raw_mode()?;
    let mut stdout = io::stdout();
    execute!(stdout, EnterAlternateScreen)?;
    let mut terminal = Terminal::new(CrosstermBackend::new(stdout))?;
    let tick = Duration::from_millis(33);

    // ── Main loop ─────────────────────────────────────────────────────────────
    loop {
        {
            let s = app.lock().unwrap();
            terminal.draw(|f| draw(f, &s))?;
        }

        if !event::poll(tick)? {
            continue;
        }
        let Event::Key(key) = event::read()? else {
            continue;
        };

        let ctrl_c = key.modifiers.contains(KeyModifiers::CONTROL)
            && key.code == KeyCode::Char('c');

        match key.code {
            KeyCode::Char('q') | KeyCode::Esc => break,
            _ if ctrl_c => break,

            KeyCode::Char('1') => app.lock().unwrap().view = ViewMode::Ir,
            KeyCode::Char('2') => app.lock().unwrap().view = ViewMode::Full,

            KeyCode::Char('+') | KeyCode::Char('=') => app.lock().unwrap().scale_up(),
            KeyCode::Char('-') => app.lock().unwrap().scale_down(),
            KeyCode::Char('a') => app.lock().unwrap().auto_scale(),

            KeyCode::Char('v') => {
                let mut s = app.lock().unwrap();
                s.smooth = !s.smooth;
            }
            KeyCode::Char('p') => app.lock().unwrap().paused = true,
            KeyCode::Char('r') => app.lock().unwrap().paused = false,
            KeyCode::Char('c') if !key.modifiers.contains(KeyModifiers::CONTROL) => {
                app.lock().unwrap().clear();
            }
            _ => {}
        }
    }

    // ── Teardown ──────────────────────────────────────────────────────────────
    if let Some(h) = sim_handle {
        h.disconnect();
    }
    disable_raw_mode()?;
    execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
    terminal.show_cursor()?;
    Ok(())
}