agentic-studio/network-poc/native-node/src/tui_dashboard.rs

use crossterm::{
    event::{Event, EventStream, KeyCode},
    execute,
    terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
};
use ratatui::{
    backend::CrosstermBackend,
    layout::{Constraint, Direction, Layout, Alignment},
    style::{Color, Modifier, Style},
    widgets::{Block, Borders, Paragraph, Wrap},
    Terminal,
};
use std::io;
use tokio::sync::RwLock;
use std::sync::Arc;
use futures_util::StreamExt;
use std::time::Duration;

#[derive(Clone)]
pub struct LogEntry {
    pub ty: String,
    pub msg: String,
    pub speed: Option<f64>,
    pub timestamp: String,
}

pub struct DashboardState {
    pub logs: Vec<LogEntry>,
    pub status: String,
    pub node_id: Option<u64>,
    pub sys_info: String,
    pub model_name: String,
    pub cur_task_id: Option<String>,
    pub cur_prompt: Option<String>,
    pub tasks_completed: u32,
    pub last_tokens_sec: f64,
    pub network_active_nodes: usize,
    pub network_total_tasks: u64,
    // VRAM-tila (ollama ps)
    pub vram_status: String,
    // Mallivalikko
    pub model_picker_open: bool,
    pub model_picker_items: Vec<String>,
    pub model_picker_idx: usize,
}

impl DashboardState {
    pub fn new() -> Self {
        Self {
            logs: Vec::new(),
            status: "ACTIVE".to_string(),
            node_id: None,
            sys_info: "".to_string(),
            model_name: "Yhdistetään...".to_string(),
            cur_task_id: None,
            cur_prompt: None,
            tasks_completed: 0,
            last_tokens_sec: 0.0,
            network_active_nodes: 1, // oletetaan itsemme
            network_total_tasks: 0,
            vram_status: "Haetaan...".to_string(),
            model_picker_open: false,
            model_picker_items: Vec::new(),
            model_picker_idx: 0,
        }
    }

    pub fn push_log(&mut self, ty: &str, msg: String, speed: Option<f64>) {
        let now = chrono::Local::now().format("%H:%M:%S").to_string();
        self.logs.push(LogEntry {
            timestamp: now,
            ty: ty.to_string(),
            msg,
            speed,
        });
        if self.logs.len() > 100 {
            self.logs.remove(0);
        }
    }
}

pub async fn run_dashboard(
    state: Arc<RwLock<DashboardState>>,
    cmd_tx: tokio::sync::mpsc::UnboundedSender<String>,
) -> Result<(), io::Error> {
    enable_raw_mode()?;
    let mut stdout = io::stdout();
    execute!(stdout, EnterAlternateScreen)?;
    let backend = CrosstermBackend::new(stdout);
    let mut terminal = Terminal::new(backend)?;
    terminal.clear()?;

    let mut reader = EventStream::new();
    let mut interval = tokio::time::interval(Duration::from_millis(100));

    loop {
        tokio::select! {
            _ = interval.tick() => {
                let st = state.read().await;
                terminal.draw(|f| ui(f, &st))?;
            }
            ev = reader.next() => {
                if let Some(Ok(Event::Key(key))) = ev {
                    let picker_open = state.read().await.model_picker_open;

                    if picker_open {
                        // Mallivalikko auki — navigointi
                        match key.code {
                            KeyCode::Up | KeyCode::Char('k') => {
                                let mut st = state.write().await;
                                if st.model_picker_idx > 0 { st.model_picker_idx -= 1; }
                            }
                            KeyCode::Down | KeyCode::Char('j') => {
                                let mut st = state.write().await;
                                let max = st.model_picker_items.len().saturating_sub(1);
                                if st.model_picker_idx < max { st.model_picker_idx += 1; }
                            }
                            KeyCode::Enter => {
                                let mut st = state.write().await;
                                let idx = st.model_picker_idx;
                                if let Some(model) = st.model_picker_items.get(idx).cloned() {
                                    st.model_picker_open = false;
                                    st.push_log("System", format!("Vaihdetaan malliin: {}...", model), None);
                                    let _ = cmd_tx.send(format!("change_model:{}", model));
                                }
                            }
                            KeyCode::Esc | KeyCode::Char('m') | KeyCode::Char('M') => {
                                state.write().await.model_picker_open = false;
                            }
                            _ => {}
                        }
                    } else {
                        // Normaali tila
                        match key.code {
                            KeyCode::Char('q') | KeyCode::Esc => {
                                disable_raw_mode()?;
                                execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
                                std::process::exit(0);
                            }
                            KeyCode::Char('p') | KeyCode::Char('P') => {
                                let _ = cmd_tx.send("pause".to_string());
                            }
                            KeyCode::Char('r') | KeyCode::Char('R') | KeyCode::Char('s') => {
                                let _ = cmd_tx.send("resume".to_string());
                            }
                            KeyCode::Char('m') | KeyCode::Char('M') => {
                                let _ = cmd_tx.send("fetch_models".to_string());
                            }
                            _ => {}
                        }
                    }
                }
            }
        }
    }
}

pub fn restore_terminal() {
    let _ = disable_raw_mode();
    let _ = execute!(io::stdout(), LeaveAlternateScreen);
}

fn ui(f: &mut ratatui::Frame, st: &DashboardState) {
    let chunks = Layout::default()
        .direction(Direction::Vertical)
        .constraints([
            Constraint::Length(3),  // Header
            Constraint::Min(0),     // Body
            Constraint::Length(3),  // Footer / Status
        ].as_ref())
        .split(f.area());

    // --- Header ---
    let header_text = match st.node_id {
        Some(id) => format!(" Kipinä Agentic Node #{} ", id),
        None => " Kipinä Agentic Node (Yhdistää...) ".to_string(),
    };
    let header = Paragraph::new(header_text)
        .style(Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD))
        .alignment(Alignment::Center)
        .block(Block::default().borders(Borders::ALL).style(Style::default().fg(Color::DarkGray)));
    f.render_widget(header, chunks[0]);

    // --- Body ---
    let body_chunks = Layout::default()
        .direction(Direction::Vertical)
        .constraints([
            Constraint::Length(8), // Yläosan info ja tehtävä
            Constraint::Min(0),    // Lokit / Chat alas
        ].as_ref())
        .split(chunks[1]);

    let top_panels = Layout::default()
        .direction(Direction::Horizontal)
        .constraints([
            Constraint::Percentage(40), // Vasen paneeli (Info)
            Constraint::Percentage(60), // Oikea paneeli (Tehtävä)
        ].as_ref())
        .split(body_chunks[0]);

    // Vasen paneeli: Laitteisto, Malli & Verkosto
    let info_text = format!(
        "🚀 Malli: {}\n🎮 VRAM: {}\n💻 Järjestelmä: {}\n📊 Tehdyt: {} | Nopeus: {} t/s\n🌐 Verkosto: {} solmua | {} tehtävää",
        st.model_name, st.vram_status, st.sys_info, st.tasks_completed, st.last_tokens_sec, st.network_active_nodes, st.network_total_tasks
    );
    let left_panel = Paragraph::new(info_text)
        .block(Block::default().title(" Laitteisto ja AI ").borders(Borders::ALL))
        .style(Style::default().fg(Color::White))
        .wrap(Wrap { trim: true });
    f.render_widget(left_panel, top_panels[0]);

    // Oikea paneeli: Käynnissä oleva tehtävä
    let task_title = match &st.cur_task_id {
        Some(id) => format!(" Työn alla: {} ", id),
        None => " Vapaana ".to_string(),
    };
    let task_content = st.cur_prompt.clone().unwrap_or_else(|| "Odotetaan tehtäviä Hubilta...".to_string());
    
    let task_style = if st.cur_task_id.is_some() {
        Style::default().fg(Color::Magenta)
    } else {
        Style::default().fg(Color::DarkGray)
    };
    
    let task_panel = Paragraph::new(task_content)
        .wrap(Wrap { trim: true })
        .block(Block::default().title(task_title).borders(Borders::ALL).style(task_style));
    f.render_widget(task_panel, top_panels[1]);

    // Alaosan paneeli: Tapahtumaloki koko leveydeltä
    let area_height = body_chunks[1].height.saturating_sub(2) as usize;
    let skip_count = if st.logs.len() > area_height { st.logs.len() - area_height } else { 0 };
    
    let visible_logs: Vec<ratatui::text::Line> = st.logs.iter().skip(skip_count).map(|log| {
        let ty_color = match log.ty.as_str() {
            "System" => Color::Yellow,
            "Network" => Color::Blue,
            "Task" => Color::Magenta,
            "Ping" => Color::DarkGray,
            _ => Color::White,
        };
        
        let speed_str = if let Some(s) = log.speed {
            format!(" | {:.1} tok/s", s)
        } else {
            "".to_string()
        };
        
        ratatui::text::Line::from(vec![
            ratatui::text::Span::styled(&log.timestamp, Style::default().fg(Color::DarkGray)),
            ratatui::text::Span::raw(" "),
            ratatui::text::Span::styled(format!("{: <8}", log.ty), Style::default().fg(ty_color).add_modifier(Modifier::BOLD)),
            ratatui::text::Span::raw(" | "),
            ratatui::text::Span::styled(log.msg.clone(), Style::default().fg(Color::White)),
            ratatui::text::Span::styled(speed_str, Style::default().fg(ty_color)),
        ])
    }).collect();

    let logs_panel = Paragraph::new(visible_logs)
        .block(Block::default().title(" Tapahtumaloki ").borders(Borders::ALL).style(Style::default().fg(Color::Cyan)));
    f.render_widget(logs_panel, body_chunks[1]);

    // --- Footer / Status ---
    let status_color = if st.status == "ACTIVE" { Color::Green } else { Color::Yellow };
    let status_text = format!(" Tila: {} | [P] Pause [R] Työhön [M] Malli [Q] Sulje ", st.status);
    let footer = Paragraph::new(status_text)
        .style(Style::default().fg(status_color).add_modifier(Modifier::BOLD))
        .alignment(Alignment::Center)
        .block(Block::default().borders(Borders::ALL));
    f.render_widget(footer, chunks[2]);

    // --- Mallivalikko-overlay ---
    if st.model_picker_open && !st.model_picker_items.is_empty() {
        let area = f.area();
        let popup_h = (st.model_picker_items.len() as u16 + 4).min(area.height - 4);
        let popup_w = 50.min(area.width - 4);
        let popup = ratatui::layout::Rect::new(
            (area.width - popup_w) / 2,
            (area.height - popup_h) / 2,
            popup_w,
            popup_h,
        );

        // Tausta
        f.render_widget(ratatui::widgets::Clear, popup);

        let items: Vec<ratatui::text::Line> = st.model_picker_items.iter().enumerate().map(|(i, name)| {
            if i == st.model_picker_idx {
                ratatui::text::Line::from(format!(" ▸ {} ", name))
                    .style(Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD))
            } else {
                ratatui::text::Line::from(format!("   {} ", name))
                    .style(Style::default().fg(Color::White))
            }
        }).collect();

        let picker = Paragraph::new(items)
            .block(Block::default()
                .title(" Vaihda malli [↑↓] Enter=valitse Esc=peruuta ")
                .borders(Borders::ALL)
                .style(Style::default().fg(Color::Cyan)));
        f.render_widget(picker, popup);
    }
}