awear 0.2.0

//! Full-screen ratatui TUI for AWEAR EEG devices.

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

use anyhow::Result;
use crossterm::event::{self, Event, KeyCode, KeyModifiers};
use crossterm::terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen};
use crossterm::ExecutableCommand;
use ratatui::prelude::*;
use ratatui::symbols::Marker;
use ratatui::widgets::*;
use tokio::sync::mpsc;

use awear::awear_client::AwearClient;
use awear::protocol::EEG_FREQUENCY;
use awear::types::*;

// ── Display constants ────────────────────────────────────────────────────────

const WINDOW_SECS: f64 = 2.0;
const BUF_SIZE: usize = 512; // 2s × 256 Hz
const Y_SCALES: &[f64] = &[100.0, 500.0, 1000.0, 5000.0, 10000.0, 20000.0, 32768.0];
const SMOOTH_WINDOW: usize = 9;
const SPINNER: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];

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

#[derive(Clone, Debug)]
enum AppMode {
    Scanning,
    Connecting(String),
    Connected { name: String, #[allow(dead_code)] id: String },
    Simulated,
    NoDevices,
    Disconnected,
}

struct App {
    buf: VecDeque<f64>,
    mode: AppMode,
    battery: Option<u8>,
    signal: Option<i8>,
    total_samples: u64,
    pkt_times: VecDeque<Instant>,
    scale_idx: usize,
    paused: bool,
    show_picker: bool,
    picker_entries: Vec<String>,
    picker_selected: usize,
    last_error: Option<String>,
    smooth: bool,
    spin_tick: usize,
}

impl App {
    fn new() -> Self {
        Self {
            buf: VecDeque::from(vec![0.0; BUF_SIZE]),
            mode: AppMode::Scanning,
            battery: None,
            signal: None,
            total_samples: 0,
            pkt_times: VecDeque::new(),
            scale_idx: 3,
            paused: false,
            show_picker: false,
            picker_entries: Vec::new(),
            picker_selected: 0,
            last_error: None,
            smooth: false,
            spin_tick: 0,
        }
    }

    fn push_samples(&mut self, samples: &[i16]) {
        for &s in samples {
            if self.buf.len() >= BUF_SIZE {
                self.buf.pop_front();
            }
            self.buf.push_back(s as f64);
            self.total_samples += 1;
        }
        self.pkt_times.push_back(Instant::now());
        let cutoff = Instant::now() - Duration::from_secs(2);
        while self.pkt_times.front().map_or(false, |t| *t < cutoff) {
            self.pkt_times.pop_front();
        }
    }

    fn pkt_rate(&self) -> f64 {
        if self.pkt_times.len() < 2 {
            return 0.0;
        }
        let elapsed = self.pkt_times.back().unwrap().duration_since(*self.pkt_times.front().unwrap());
        if elapsed.as_secs_f64() < 0.01 {
            return 0.0;
        }
        self.pkt_times.len() as f64 / elapsed.as_secs_f64()
    }

    fn auto_scale(&mut self) {
        let peak = self.buf.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
        let target = peak * 1.1;
        for (i, &s) in Y_SCALES.iter().enumerate() {
            if s >= target {
                self.scale_idx = i;
                return;
            }
        }
        self.scale_idx = Y_SCALES.len() - 1;
    }

    fn clear(&mut self) {
        self.buf = VecDeque::from(vec![0.0; BUF_SIZE]);
        self.total_samples = 0;
        self.pkt_times.clear();
    }
}

// ── Main ─────────────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() -> Result<()> {
    let simulate = std::env::args().any(|a| a == "--simulate");

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

    // Terminal setup — use a panic hook to restore terminal on crash
    enable_raw_mode()?;
    let mut stdout = io::stdout();
    stdout.execute(EnterAlternateScreen)?;

    let original_hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(move |panic_info| {
        let _ = disable_raw_mode();
        let _ = io::stdout().execute(LeaveAlternateScreen);
        original_hook(panic_info);
    }));

    let backend = CrosstermBackend::new(stdout);
    let mut terminal = Terminal::new(backend)?;

    // Channel for device commands (STOP, START, __DISCONNECT__)
    let (cmd_tx, cmd_rx) = mpsc::channel::<String>(32);

    // Spawn device tasks
    let app2 = app.clone();
    if simulate {
        app.lock().unwrap().mode = AppMode::Simulated;
        tokio::spawn(spawn_simulator(app2));
        drop(cmd_rx);
    } else {
        tokio::spawn(scan_and_connect(app2, cmd_rx));
    }

    // Render loop (~30 FPS)
    let tick_rate = Duration::from_millis(33);
    loop {
        // Draw
        {
            let mut a = app.lock().unwrap();
            a.spin_tick = a.spin_tick.wrapping_add(1);
            terminal.draw(|f| draw(f, &a))?;
        }

        // Input
        if event::poll(tick_rate)? {
            if let Event::Key(key) = event::read()? {
                let mut a = app.lock().unwrap();

                if key.code == KeyCode::Char('c') && key.modifiers.contains(KeyModifiers::CONTROL) {
                    break;
                }

                match key.code {
                    KeyCode::Char('q') | KeyCode::Esc => {
                        if a.show_picker {
                            a.show_picker = false;
                        } else {
                            break;
                        }
                    }
                    KeyCode::Tab => {
                        a.show_picker = !a.show_picker;
                    }
                    KeyCode::Char('+') | KeyCode::Char('=') => {
                        if a.scale_idx < Y_SCALES.len() - 1 {
                            a.scale_idx += 1;
                        }
                    }
                    KeyCode::Char('-') => {
                        if a.scale_idx > 0 {
                            a.scale_idx -= 1;
                        }
                    }
                    KeyCode::Char('a') => a.auto_scale(),
                    KeyCode::Char('v') => a.smooth = !a.smooth,
                    KeyCode::Char('p') => {
                        a.paused = true;
                        let _ = cmd_tx.try_send("STOP".to_string());
                    }
                    KeyCode::Char('r') => {
                        a.paused = false;
                        let _ = cmd_tx.try_send("START".to_string());
                    }
                    KeyCode::Char('c') => a.clear(),
                    KeyCode::Char('d') => {
                        let _ = cmd_tx.try_send("__DISCONNECT__".to_string());
                    }
                    KeyCode::Up => {
                        if a.show_picker && a.picker_selected > 0 {
                            a.picker_selected -= 1;
                        }
                    }
                    KeyCode::Down => {
                        if a.show_picker && a.picker_selected < a.picker_entries.len().saturating_sub(1) {
                            a.picker_selected += 1;
                        }
                    }
                    _ => {}
                }
            }
        }
    }

    // Cleanup — disconnect device if connected
    let _ = cmd_tx.try_send("__DISCONNECT__".to_string());
    disable_raw_mode()?;
    io::stdout().execute(LeaveAlternateScreen)?;
    Ok(())
}

// ── BLE tasks ────────────────────────────────────────────────────────────────

async fn scan_and_connect(app: Arc<Mutex<App>>, mut cmd_rx: mpsc::Receiver<String>) {
    let config = AwearClientConfig::default();
    let client = AwearClient::new(config);

    let devices = match client.scan_all().await {
        Ok(d) => d,
        Err(e) => {
            let mut a = app.lock().unwrap();
            a.mode = AppMode::NoDevices;
            a.last_error = Some(format!("Scan failed: {}", e));
            return;
        }
    };

    if devices.is_empty() {
        app.lock().unwrap().mode = AppMode::NoDevices;
        return;
    }

    {
        let mut a = app.lock().unwrap();
        a.picker_entries = devices
            .iter()
            .map(|d| format!("{} ({})", d.name, d.rssi))
            .collect();
    }

    let mut devices = devices;
    let device = devices.swap_remove(0);
    let device_name = device.name.clone();
    let device_id = device.id.clone();
    app.lock().unwrap().mode = AppMode::Connecting(device_name.clone());

    let client2 = AwearClient::new(AwearClientConfig::default());
    let (mut event_rx, handle) = match client2.connect_to(device).await {
        Ok(pair) => pair,
        Err(e) => {
            let mut a = app.lock().unwrap();
            a.mode = AppMode::Disconnected;
            a.last_error = Some(format!("Connection failed: {}", e));
            return;
        }
    };

    app.lock().unwrap().mode = AppMode::Connected {
        name: device_name,
        id: device_id,
    };

    // Wait for authentication, then start streaming
    tokio::time::sleep(Duration::from_secs(2)).await;
    if let Err(e) = handle.start().await {
        app.lock().unwrap().last_error = Some(format!("Start failed: {}", e));
    }

    // Main event + command loop
    loop {
        tokio::select! {
            event = event_rx.recv() => {
                match event {
                    Some(AwearEvent::Eeg(reading)) => {
                        let mut a = app.lock().unwrap();
                        if !a.paused {
                            a.push_samples(&reading.samples);
                        }
                    }
                    Some(AwearEvent::Battery(level)) => {
                        app.lock().unwrap().battery = Some(level);
                    }
                    Some(AwearEvent::Signal(rssi)) => {
                        app.lock().unwrap().signal = Some(rssi);
                    }
                    Some(AwearEvent::Disconnected) | None => {
                        app.lock().unwrap().mode = AppMode::Disconnected;
                        break;
                    }
                    _ => {}
                }
            }
            cmd = cmd_rx.recv() => {
                match cmd.as_deref() {
                    Some("__DISCONNECT__") => {
                        let _ = handle.stop().await;
                        let _ = handle.disconnect().await;
                        app.lock().unwrap().mode = AppMode::Disconnected;
                        break;
                    }
                    Some(c) => {
                        let _ = handle.send_command(c).await;
                    }
                    None => break,
                }
            }
        }
    }
}

// ── Simulator ────────────────────────────────────────────────────────────────

async fn spawn_simulator(app: Arc<Mutex<App>>) {
    let mut t: f64 = 0.0;
    let dt = 1.0 / EEG_FREQUENCY;
    let samples_per_block = 32;

    loop {
        tokio::time::sleep(Duration::from_millis(
            (1000.0 * samples_per_block as f64 / EEG_FREQUENCY) as u64,
        )).await;

        let mut samples = Vec::with_capacity(samples_per_block);
        for _ in 0..samples_per_block {
            let alpha = 2000.0 * (2.0 * std::f64::consts::PI * 10.0 * t).sin();
            let beta = 600.0 * (2.0 * std::f64::consts::PI * 22.0 * t).sin();
            let theta = 1000.0 * (2.0 * std::f64::consts::PI * 6.0 * t).sin();
            let noise = ((t * 12345.6789).sin() * 43758.5453).fract() * 800.0 - 400.0;
            let val = (alpha + beta + theta + noise).clamp(i16::MIN as f64, i16::MAX as f64) as i16;
            samples.push(val);
            t += dt;
        }

        let mut a = app.lock().unwrap();
        if !a.paused {
            a.push_samples(&samples);
        }
        // Simulate battery every ~4s
        if a.total_samples % 1024 == 0 {
            a.battery = Some(85);
            a.signal = Some(-45);
        }
    }
}

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

fn draw(f: &mut Frame, app: &App) {
    let chunks = Layout::default()
        .direction(Direction::Vertical)
        .constraints([
            Constraint::Length(3),  // header
            Constraint::Min(10),   // chart
            Constraint::Length(3), // footer
        ])
        .split(f.area());

    draw_header(f, app, chunks[0]);
    draw_chart(f, app, chunks[1]);
    draw_footer(f, app, chunks[2]);

    if app.show_picker {
        draw_picker(f, app);
    }
}

fn draw_header(f: &mut Frame, app: &App, area: Rect) {
    let spin = SPINNER[app.spin_tick / 3 % SPINNER.len()];

    let (status_text, status_color) = match &app.mode {
        AppMode::Scanning => (format!("{} Scanning...", spin), Color::Yellow),
        AppMode::Connecting(name) => (format!("{} Connecting to {}...", spin, name), Color::Yellow),
        AppMode::Connected { name, .. } => (format!("● {}", name), Color::Green),
        AppMode::Simulated => ("● Simulated".to_string(), Color::Cyan),
        AppMode::NoDevices => ("✗ No devices found".to_string(), Color::Red),
        AppMode::Disconnected => ("✗ Disconnected".to_string(), Color::Red),
    };

    let battery_text = app
        .battery
        .map(|b| format!("  🔋 {}%", b))
        .unwrap_or_default();

    let signal_text = app
        .signal
        .map(|s| format!("  📶 {} dBm", s))
        .unwrap_or_default();

    let scale = Y_SCALES[app.scale_idx];
    let rate = app.pkt_rate();

    let header = format!(
        "{}{}{}  │  ±{:.0}  │  {:.0} pkt/s  │  {} samples{}",
        status_text,
        battery_text,
        signal_text,
        scale,
        rate,
        app.total_samples,
        if app.paused { "  ⏸ PAUSED" } else { "" },
    );

    let block = Block::default()
        .borders(Borders::ALL)
        .border_style(Style::default().fg(status_color))
        .title(" AWEAR EEG ");

    let para = Paragraph::new(header)
        .style(Style::default().fg(Color::White))
        .block(block);
    f.render_widget(para, area);
}

fn draw_chart(f: &mut Frame, app: &App, area: Rect) {
    let scale = Y_SCALES[app.scale_idx];

    let data: Vec<(f64, f64)> = app
        .buf
        .iter()
        .enumerate()
        .map(|(i, &v)| (i as f64 / EEG_FREQUENCY, v))
        .collect();

    let mut datasets = vec![Dataset::default()
        .name("EEG")
        .marker(Marker::Braille)
        .graph_type(GraphType::Line)
        .style(Style::default().fg(Color::Cyan))
        .data(&data)];

    // Smoothed overlay
    let smoothed;
    if app.smooth && data.len() > SMOOTH_WINDOW {
        smoothed = smooth_signal(&data, SMOOTH_WINDOW);
        datasets.push(
            Dataset::default()
                .name("Smooth")
                .marker(Marker::Braille)
                .graph_type(GraphType::Line)
                .style(Style::default().fg(Color::Yellow))
                .data(&smoothed),
        );
    }

    // Stats
    let rms = if !app.buf.is_empty() {
        let sum_sq: f64 = app.buf.iter().map(|v| v * v).sum();
        (sum_sq / app.buf.len() as f64).sqrt()
    } else {
        0.0
    };
    let min_v = app.buf.iter().cloned().fold(f64::INFINITY, f64::min);
    let max_v = app.buf.iter().cloned().fold(f64::NEG_INFINITY, f64::max);

    let title = format!(
        " EEG  min={:.0}  max={:.0}  RMS={:.0}{}",
        min_v,
        max_v,
        rms,
        if app.smooth { "  [SMOOTH]" } else { "" },
    );

    // Detect clipping
    let clipping = max_v.abs() > scale * 0.95 || min_v.abs() > scale * 0.95;
    let border_color = if clipping { Color::Red } else { Color::DarkGray };

    let chart = Chart::new(datasets)
        .block(
            Block::default()
                .borders(Borders::ALL)
                .border_style(Style::default().fg(border_color))
                .title(title),
        )
        .x_axis(
            Axis::default()
                .bounds([0.0, WINDOW_SECS])
                .labels(vec!["0s".to_string(), format!("{:.0}s", WINDOW_SECS)]),
        )
        .y_axis(
            Axis::default()
                .bounds([-scale, scale])
                .labels(vec![
                    format!("-{:.0}", scale),
                    "0".to_string(),
                    format!("+{:.0}", scale),
                ]),
        );

    f.render_widget(chart, area);
}

fn draw_footer(f: &mut Frame, app: &App, area: Rect) {
    let keys = match &app.mode {
        AppMode::NoDevices => {
            "Tab: picker  s: scan  q: quit  │  Tip: check Bluetooth permissions (macOS)".to_string()
        }
        _ => {
            format!(
                "Tab: picker  +/-: scale  a: auto  v: smooth  p: pause  r: resume  c: clear  d: disconnect  q: quit",
            )
        }
    };

    let error_text = app
        .last_error
        .as_ref()
        .map(|e| format!("  ⚠ {}", e))
        .unwrap_or_default();

    let footer = Paragraph::new(format!("{}{}", keys, error_text))
        .style(Style::default().fg(Color::DarkGray))
        .block(Block::default().borders(Borders::ALL));
    f.render_widget(footer, area);
}

fn draw_picker(f: &mut Frame, app: &App) {
    let area = centered_rect(50, 60, f.area());

    // Clear background
    f.render_widget(Clear, area);

    let items: Vec<ListItem> = app
        .picker_entries
        .iter()
        .enumerate()
        .map(|(i, entry)| {
            let style = if i == app.picker_selected {
                Style::default().fg(Color::Black).bg(Color::Cyan)
            } else {
                Style::default().fg(Color::White)
            };
            ListItem::new(entry.as_str()).style(style)
        })
        .collect();

    let title = " Devices ";

    let list = List::new(items).block(
        Block::default()
            .borders(Borders::ALL)
            .border_style(Style::default().fg(Color::Cyan))
            .title(title),
    );

    f.render_widget(list, area);
}

fn centered_rect(percent_x: u16, percent_y: u16, r: Rect) -> Rect {
    let popup_layout = Layout::default()
        .direction(Direction::Vertical)
        .constraints([
            Constraint::Percentage((100 - percent_y) / 2),
            Constraint::Percentage(percent_y),
            Constraint::Percentage((100 - percent_y) / 2),
        ])
        .split(r);

    Layout::default()
        .direction(Direction::Horizontal)
        .constraints([
            Constraint::Percentage((100 - percent_x) / 2),
            Constraint::Percentage(percent_x),
            Constraint::Percentage((100 - percent_x) / 2),
        ])
        .split(popup_layout[1])[1]
}

// ── Signal processing ────────────────────────────────────────────────────────

fn smooth_signal(data: &[(f64, f64)], window: usize) -> Vec<(f64, f64)> {
    let half = window / 2;
    let mut result = Vec::with_capacity(data.len());
    for i in 0..data.len() {
        let start = i.saturating_sub(half);
        let end = (i + half + 1).min(data.len());
        let sum: f64 = data[start..end].iter().map(|(_, v)| v).sum();
        let count = (end - start) as f64;
        result.push((data[i].0, sum / count));
    }
    result
}