cm-dashboard/dashboard/src/ui/mod.rs

use anyhow::Result;
use crossterm::event::{Event, KeyCode};
use ratatui::{
    layout::{Constraint, Direction, Layout, Rect},
    style::Style,
    widgets::{Block, Paragraph},
    Frame,
};
use std::collections::HashMap;
use std::time::Instant;
use tracing::info;

pub mod theme;
pub mod widgets;

use crate::config::DashboardConfig;
use crate::metrics::MetricStore;
use cm_dashboard_shared::{Metric, Status};
use theme::{Components, Layout as ThemeLayout, Theme, Typography};
use widgets::{BackupWidget, ServicesWidget, SystemWidget, Widget};

/// Commands that can be triggered from the UI
#[derive(Debug, Clone)]
pub enum UiCommand {
    ServiceStart { hostname: String, service_name: String },
    ServiceStop { hostname: String, service_name: String },
    TriggerBackup { hostname: String },
}


/// Types of commands for status tracking
#[derive(Debug, Clone)]
pub enum CommandType {
    ServiceStart,
    ServiceStop,
    BackupTrigger,
}

/// Panel types for focus management

/// Widget states for a specific host
#[derive(Clone)]
pub struct HostWidgets {
    /// System widget state (includes CPU, Memory, NixOS info, Storage)
    pub system_widget: SystemWidget,
    /// Services widget state
    pub services_widget: ServicesWidget,
    /// Backup widget state
    pub backup_widget: BackupWidget,
    /// Scroll offsets for each panel
    pub system_scroll_offset: usize,
    pub services_scroll_offset: usize,
    pub backup_scroll_offset: usize,
    /// Last update time for this host
    pub last_update: Option<Instant>,
    /// Pending service transitions for immediate visual feedback
    pub pending_service_transitions: HashMap<String, (CommandType, String, Instant)>, // service_name -> (command_type, original_status, start_time)
}

impl HostWidgets {
    pub fn new() -> Self {
        Self {
            system_widget: SystemWidget::new(),
            services_widget: ServicesWidget::new(),
            backup_widget: BackupWidget::new(),
            system_scroll_offset: 0,
            services_scroll_offset: 0,
            backup_scroll_offset: 0,
            last_update: None,
            pending_service_transitions: HashMap::new(),
        }
    }
}


/// Main TUI application
pub struct TuiApp {
    /// Widget states per host (hostname -> HostWidgets)
    host_widgets: HashMap<String, HostWidgets>,
    /// Current active host
    current_host: Option<String>,
    /// Available hosts
    available_hosts: Vec<String>,
    /// Host index for navigation
    host_index: usize,
    /// Should quit application
    should_quit: bool,
    /// Track if user manually navigated away from localhost
    user_navigated_away: bool,
    /// Dashboard configuration
    config: DashboardConfig,
}

impl TuiApp {
    pub fn new(config: DashboardConfig) -> Self {
        Self {
            host_widgets: HashMap::new(),
            current_host: None,
            available_hosts: Vec::new(),
            host_index: 0,
            should_quit: false,
            user_navigated_away: false,
            config,
        }
    }

    /// Get or create host widgets for the given hostname
    fn get_or_create_host_widgets(&mut self, hostname: &str) -> &mut HostWidgets {
        self.host_widgets
            .entry(hostname.to_string())
            .or_insert_with(HostWidgets::new)
    }

    /// Update widgets with metrics from store (only for current host)
    pub fn update_metrics(&mut self, metric_store: &MetricStore) {
        
        // Check for rebuild completion by agent hash change
        
        if let Some(hostname) = self.current_host.clone() {
            // Only update widgets if we have metrics for this host
            let all_metrics = metric_store.get_metrics_for_host(&hostname);
            if !all_metrics.is_empty() {
                // Get metrics first while hostname is borrowed
                let cpu_metrics: Vec<&Metric> = all_metrics
                    .iter()
                    .filter(|m| {
                        m.name.starts_with("cpu_")
                            || m.name.contains("c_state_")
                            || m.name.starts_with("process_top_")
                    })
                    .copied()
                    .collect();
                let memory_metrics: Vec<&Metric> = all_metrics
                    .iter()
                    .filter(|m| m.name.starts_with("memory_") || m.name.starts_with("disk_tmp_"))
                    .copied()
                    .collect();
                let service_metrics: Vec<&Metric> = all_metrics
                    .iter()
                    .filter(|m| m.name.starts_with("service_"))
                    .copied()
                    .collect();
                let all_backup_metrics: Vec<&Metric> = all_metrics
                    .iter()
                    .filter(|m| m.name.starts_with("backup_"))
                    .copied()
                    .collect();

                // Clear completed transitions first
                self.clear_completed_transitions(&hostname, &service_metrics);

                // Now get host widgets and update them
                let host_widgets = self.get_or_create_host_widgets(&hostname);

                // Collect all system metrics (CPU, memory, NixOS, disk/storage)
                let mut system_metrics = cpu_metrics;
                system_metrics.extend(memory_metrics);
                
                // Add NixOS metrics - using exact matching for build display fix
                let nixos_metrics: Vec<&Metric> = all_metrics
                    .iter()
                    .filter(|m| m.name == "system_nixos_build" || m.name == "system_active_users" || m.name == "agent_version")
                    .copied()
                    .collect();
                system_metrics.extend(nixos_metrics);
                
                // Add disk/storage metrics
                let disk_metrics: Vec<&Metric> = all_metrics
                    .iter()
                    .filter(|m| m.name.starts_with("disk_"))
                    .copied()
                    .collect();
                system_metrics.extend(disk_metrics);

                host_widgets.system_widget.update_from_metrics(&system_metrics);
                host_widgets
                    .services_widget
                    .update_from_metrics(&service_metrics);
                host_widgets
                    .backup_widget
                    .update_from_metrics(&all_backup_metrics);

                host_widgets.last_update = Some(Instant::now());
            }
        }
    }

    /// Update available hosts with localhost prioritization
    pub fn update_hosts(&mut self, hosts: Vec<String>) {
        // Sort hosts alphabetically
        let mut sorted_hosts = hosts.clone();
        
        // Keep hosts that have pending transitions even if they're offline
        for (hostname, host_widgets) in &self.host_widgets {
            if !host_widgets.pending_service_transitions.is_empty() {
                if !sorted_hosts.contains(hostname) {
                    sorted_hosts.push(hostname.clone());
                }
            }
        }
        
        sorted_hosts.sort();
        self.available_hosts = sorted_hosts;
        
        // Get the current hostname (localhost) for auto-selection
        let localhost = gethostname::gethostname().to_string_lossy().to_string();
        if !self.available_hosts.is_empty() {
            if self.available_hosts.contains(&localhost) && !self.user_navigated_away {
                // Localhost is available and user hasn't navigated away - switch to it
                self.current_host = Some(localhost.clone());
                // Find the actual index of localhost in the sorted list
                self.host_index = self.available_hosts.iter().position(|h| h == &localhost).unwrap_or(0);
            } else if self.current_host.is_none() {
                // No current host - select first available (which is localhost if available)
                self.current_host = Some(self.available_hosts[0].clone());
                self.host_index = 0;
            } else if let Some(ref current) = self.current_host {
                if !self.available_hosts.contains(current) {
                    // Current host disconnected - select first available and reset navigation flag
                    self.current_host = Some(self.available_hosts[0].clone());
                    self.host_index = 0;
                    self.user_navigated_away = false; // Reset since we're forced to switch
                } else if let Some(index) = self.available_hosts.iter().position(|h| h == current) {
                    // Update index for current host
                    self.host_index = index;
                }
            }
        }
    }

    /// Handle keyboard input
    pub fn handle_input(&mut self, event: Event) -> Result<Option<UiCommand>> {
        if let Event::Key(key) = event {
            match key.code {
                KeyCode::Char('q') => {
                    self.should_quit = true;
                }
                KeyCode::Left => {
                    self.navigate_host(-1);
                }
                KeyCode::Right => {
                    self.navigate_host(1);
                }
                KeyCode::Char('r') => {
                    // System rebuild command - works on any panel for current host
                    if let Some(hostname) = self.current_host.clone() {
                        // Create command that shows CM Dashboard logo and then rebuilds
                        let logo_and_rebuild = format!(
                            r"cat << 'EOF'
NixOS System Rebuild
Target: {}

EOF
ssh -tt {}@{} 'bash -ic {}'",
                            hostname,
                            self.config.ssh.rebuild_user,
                            hostname,
                            self.config.ssh.rebuild_alias
                        );
                        
                        std::process::Command::new("tmux")
                            .arg("display-popup")
                            .arg("-w")
                            .arg("80%")
                            .arg("-h") 
                            .arg("80%")
                            .arg(&logo_and_rebuild)
                            .spawn()
                            .ok(); // Ignore errors, tmux will handle them
                    }
                }
                KeyCode::Char('s') => {
                    // Service start command
                    if let (Some(service_name), Some(hostname)) = (self.get_selected_service(), self.current_host.clone()) {
                        if self.start_command(&hostname, CommandType::ServiceStart, service_name.clone()) {
                            return Ok(Some(UiCommand::ServiceStart { hostname, service_name }));
                        }
                    }
                }
                KeyCode::Char('S') => {
                    // Service stop command
                    if let (Some(service_name), Some(hostname)) = (self.get_selected_service(), self.current_host.clone()) {
                        if self.start_command(&hostname, CommandType::ServiceStop, service_name.clone()) {
                            return Ok(Some(UiCommand::ServiceStop { hostname, service_name }));
                        }
                    }
                }
                KeyCode::Char('J') => {
                    // Show service logs via journalctl in tmux popup
                    if let (Some(service_name), Some(hostname)) = (self.get_selected_service(), self.current_host.clone()) {
                        let journalctl_command = format!(
                            "ssh -tt {}@{} 'journalctl -u {}.service -f --no-pager -n 50'",
                            self.config.ssh.rebuild_user,
                            hostname,
                            service_name
                        );
                        
                        std::process::Command::new("tmux")
                            .arg("display-popup")
                            .arg("-w")
                            .arg("80%")
                            .arg("-h") 
                            .arg("80%")
                            .arg("-t")
                            .arg(format!("Logs: {}", service_name))
                            .arg(&journalctl_command)
                            .spawn()
                            .ok(); // Ignore errors, tmux will handle them
                    }
                }
                KeyCode::Char('b') => {
                    // Trigger backup
                    if let Some(hostname) = self.current_host.clone() {
                        self.start_command(&hostname, CommandType::BackupTrigger, hostname.clone());
                        return Ok(Some(UiCommand::TriggerBackup { hostname }));
                    }
                }
                KeyCode::Tab => {
                    // Tab cycles to next host
                    self.navigate_host(1);
                }
                KeyCode::Up | KeyCode::Char('k') => {
                    // Move service selection up
                    if let Some(hostname) = self.current_host.clone() {
                        let host_widgets = self.get_or_create_host_widgets(&hostname);
                        host_widgets.services_widget.select_previous();
                    }
                }
                KeyCode::Down | KeyCode::Char('j') => {
                    // Move service selection down
                    if let Some(hostname) = self.current_host.clone() {
                        let total_services = {
                            let host_widgets = self.get_or_create_host_widgets(&hostname);
                            host_widgets.services_widget.get_total_services_count()
                        };
                        let host_widgets = self.get_or_create_host_widgets(&hostname);
                        host_widgets.services_widget.select_next(total_services);
                    }
                }
                _ => {}
            }
        }
        Ok(None)
    }

    /// Navigate between hosts
    fn navigate_host(&mut self, direction: i32) {
        if self.available_hosts.is_empty() {
            return;
        }

        let len = self.available_hosts.len();
        if direction > 0 {
            self.host_index = (self.host_index + 1) % len;
        } else {
            self.host_index = if self.host_index == 0 {
                len - 1
            } else {
                self.host_index - 1
            };
        }

        self.current_host = Some(self.available_hosts[self.host_index].clone());
        
        // Check if user navigated away from localhost
        let localhost = gethostname::gethostname().to_string_lossy().to_string();
        if let Some(ref current) = self.current_host {
            if current != &localhost {
                self.user_navigated_away = true;
            } else {
                self.user_navigated_away = false; // User navigated back to localhost
            }
        }
        
        info!("Switched to host: {}", self.current_host.as_ref().unwrap());
    }





    /// Get the currently selected service name from the services widget
    fn get_selected_service(&self) -> Option<String> {
        if let Some(hostname) = &self.current_host {
            if let Some(host_widgets) = self.host_widgets.get(hostname) {
                return host_widgets.services_widget.get_selected_service();
            }
        }
        None
    }


    /// Should quit application
    pub fn should_quit(&self) -> bool {
        self.should_quit
    }

    /// Get current service status for state-aware command validation
    fn get_current_service_status(&self, hostname: &str, service_name: &str) -> Option<String> {
        if let Some(host_widgets) = self.host_widgets.get(hostname) {
            return host_widgets.services_widget.get_service_status(service_name);
        }
        None
    }

    /// Start command execution with immediate visual feedback 
    pub fn start_command(&mut self, hostname: &str, command_type: CommandType, target: String) -> bool {
        // Get current service status to validate command
        let current_status = self.get_current_service_status(hostname, &target);
        
        // Validate if command makes sense for current state
        let should_execute = match (&command_type, current_status.as_deref()) {
            (CommandType::ServiceStart, Some("inactive") | Some("failed") | Some("dead")) => true,
            (CommandType::ServiceStop, Some("active")) => true,
            (CommandType::ServiceStart, Some("active")) => {
                // Already running - don't execute
                false
            },
            (CommandType::ServiceStop, Some("inactive") | Some("failed") | Some("dead")) => {
                // Already stopped - don't execute  
                false
            },
            (_, None) => {
                // Unknown service state - allow command to proceed
                true
            },
            _ => true, // Default: allow other combinations
        };
        
        // ALWAYS store the pending transition for immediate visual feedback, even if we don't execute
        if let Some(host_widgets) = self.host_widgets.get_mut(hostname) {
            host_widgets.pending_service_transitions.insert(
                target.clone(), 
                (command_type, current_status.unwrap_or_else(|| "unknown".to_string()), Instant::now())
            );
        }
        
        should_execute
    }

    /// Clear pending transitions when real status updates arrive or timeout
    fn clear_completed_transitions(&mut self, hostname: &str, service_metrics: &[&Metric]) {
        if let Some(host_widgets) = self.host_widgets.get_mut(hostname) {
            let mut completed_services = Vec::new();
            
            // Check each pending transition to see if real status has changed
            for (service_name, (command_type, original_status, _start_time)) in &host_widgets.pending_service_transitions {
                
                // Look for status metric for this service
                for metric in service_metrics {
                    if metric.name == format!("service_{}_status", service_name) {
                        let new_status = metric.value.as_string();
                        
                        // Check if status has changed from original (command completed)
                        if &new_status != original_status {
                            // Verify it changed in the expected direction
                            let expected_change = match command_type {
                                CommandType::ServiceStart => &new_status == "active",
                                CommandType::ServiceStop => &new_status != "active",
                                _ => false,
                            };
                            
                            if expected_change {
                                completed_services.push(service_name.clone());
                            }
                        }
                        break;
                    }
                }
            }
            
            // Remove completed transitions
            for service_name in completed_services {
                host_widgets.pending_service_transitions.remove(&service_name);
            }
        }
    }




    /// Render the dashboard (real btop-style multi-panel layout)
    pub fn render(&mut self, frame: &mut Frame, metric_store: &MetricStore) {
        let size = frame.size();

        // Clear background to true black like btop
        frame.render_widget(
            Block::default().style(Style::default().bg(Theme::background())),
            size,
        );

        // Create real btop-style layout: multi-panel with borders
        // Three-section layout: title bar, main content, statusbar
        let main_chunks = Layout::default()
            .direction(Direction::Vertical)
            .constraints([
                Constraint::Length(1), // Title bar
                Constraint::Min(0),    // Main content area
                Constraint::Length(1), // Statusbar
            ])
            .split(size);

        // New layout: left panels | right services (100% height)
        let content_chunks = ratatui::layout::Layout::default()
            .direction(Direction::Horizontal)
            .constraints([
                Constraint::Percentage(ThemeLayout::LEFT_PANEL_WIDTH), // Left side: system, backup
                Constraint::Percentage(ThemeLayout::RIGHT_PANEL_WIDTH), // Right side: services (100% height)
            ])
            .split(main_chunks[1]); // main_chunks[1] is now the content area (between title and statusbar)

        // Check if backup panel should be shown
        let show_backup = if let Some(hostname) = self.current_host.clone() {
            let host_widgets = self.get_or_create_host_widgets(&hostname);
            host_widgets.backup_widget.has_data()
        } else {
            false
        };

        // Left side: dynamic layout based on backup data availability
        let left_chunks = if show_backup {
            // Show both system and backup panels
            ratatui::layout::Layout::default()
                .direction(Direction::Vertical)
                .constraints([
                    Constraint::Percentage(ThemeLayout::SYSTEM_PANEL_HEIGHT), // System section
                    Constraint::Percentage(ThemeLayout::BACKUP_PANEL_HEIGHT), // Backup section
                ])
                .split(content_chunks[0])
        } else {
            // Show only system panel (full height)
            ratatui::layout::Layout::default()
                .direction(Direction::Vertical)
                .constraints([Constraint::Percentage(100)]) // System section takes full height
                .split(content_chunks[0])
        };

        // Render title bar
        self.render_btop_title(frame, main_chunks[0], metric_store);

        // Render new panel layout
        self.render_system_panel(frame, left_chunks[0], metric_store);
        if show_backup && left_chunks.len() > 1 {
            self.render_backup_panel(frame, left_chunks[1]);
        }

        // Render services widget for current host
        if let Some(hostname) = self.current_host.clone() {
            let is_focused = true; // Always show service selection
            let (scroll_offset, pending_transitions) = {
                let host_widgets = self.get_or_create_host_widgets(&hostname);
                (host_widgets.services_scroll_offset, host_widgets.pending_service_transitions.clone())
            };
            let host_widgets = self.get_or_create_host_widgets(&hostname);
            host_widgets
                .services_widget
                .render_with_transitions(frame, content_chunks[1], is_focused, scroll_offset, &pending_transitions); // Services takes full right side
        }

        // Render statusbar at the bottom
        self.render_statusbar(frame, main_chunks[2]); // main_chunks[2] is the statusbar area

    }

    /// Render btop-style minimal title with host status colors
    fn render_btop_title(&self, frame: &mut Frame, area: Rect, metric_store: &MetricStore) {
        use ratatui::style::Modifier;
        use ratatui::text::{Line, Span};
        use theme::StatusIcons;

        if self.available_hosts.is_empty() {
            let title_text = "cm-dashboard • no hosts discovered";
            let title = Paragraph::new(title_text)
                .style(Style::default().fg(Theme::background()).bg(Theme::status_color(Status::Unknown)));
            frame.render_widget(title, area);
            return;
        }

        // Calculate worst-case status across all hosts
        let mut worst_status = Status::Ok;
        for host in &self.available_hosts {
            let host_status = self.calculate_host_status(host, metric_store);
            worst_status = Status::aggregate(&[worst_status, host_status]);
        }

        // Use the worst status color as background
        let background_color = Theme::status_color(worst_status);

        // Split the title bar into left and right sections
        let chunks = Layout::default()
            .direction(Direction::Horizontal)
            .constraints([Constraint::Length(15), Constraint::Min(0)])
            .split(area);

        // Left side: "cm-dashboard" text
        let left_span = Span::styled(
            " cm-dashboard", 
            Style::default().fg(Theme::background()).bg(background_color).add_modifier(Modifier::BOLD)
        );
        let left_title = Paragraph::new(Line::from(vec![left_span]))
            .style(Style::default().bg(background_color));
        frame.render_widget(left_title, chunks[0]);

        // Right side: hosts with status indicators
        let mut host_spans = Vec::new();
        
        for (i, host) in self.available_hosts.iter().enumerate() {
            if i > 0 {
                host_spans.push(Span::styled(
                    " ", 
                    Style::default().fg(Theme::background()).bg(background_color)
                ));
            }

            // Always show normal status icon based on metrics (no command status at host level)
            let host_status = self.calculate_host_status(host, metric_store);
            let status_icon = StatusIcons::get_icon(host_status);

            // Add status icon with background color as foreground against status background
            host_spans.push(Span::styled(
                format!("{} ", status_icon),
                Style::default().fg(Theme::background()).bg(background_color),
            ));

            if Some(host) == self.current_host.as_ref() {
                // Selected host in bold background color against status background
                host_spans.push(Span::styled(
                    host.clone(),
                    Style::default()
                        .fg(Theme::background())
                        .bg(background_color)
                        .add_modifier(Modifier::BOLD),
                ));
            } else {
                // Other hosts in normal background color against status background
                host_spans.push(Span::styled(
                    host.clone(),
                    Style::default().fg(Theme::background()).bg(background_color),
                ));
            }
        }

        // Add right padding
        host_spans.push(Span::styled(
            " ", 
            Style::default().fg(Theme::background()).bg(background_color)
        ));

        let host_line = Line::from(host_spans);
        let host_title = Paragraph::new(vec![host_line])
            .style(Style::default().bg(background_color))
            .alignment(ratatui::layout::Alignment::Right);
        frame.render_widget(host_title, chunks[1]);
    }

    /// Calculate overall status for a host based on its metrics
    fn calculate_host_status(&self, hostname: &str, metric_store: &MetricStore) -> Status {
        let metrics = metric_store.get_metrics_for_host(hostname);

        if metrics.is_empty() {
            return Status::Unknown;
        }

        // First check if we have the aggregated host status summary from the agent
        if let Some(host_summary_metric) = metric_store.get_metric(hostname, "host_status_summary") {
            return host_summary_metric.status;
        }

        // Fallback to old aggregation logic with proper Pending handling
        let mut has_critical = false;
        let mut has_warning = false;
        let mut has_pending = false;
        let mut ok_count = 0;

        for metric in &metrics {
            match metric.status {
                Status::Critical => has_critical = true,
                Status::Warning => has_warning = true,
                Status::Pending => has_pending = true,
                Status::Ok => ok_count += 1,
                Status::Unknown => {} // Ignore unknown for aggregation
            }
        }

        // Priority order: Critical > Warning > Pending > Ok > Unknown
        if has_critical {
            Status::Critical
        } else if has_warning {
            Status::Warning
        } else if has_pending {
            Status::Pending
        } else if ok_count > 0 {
            Status::Ok
        } else {
            Status::Unknown
        }
    }

    /// Render dynamic statusbar with context-aware shortcuts
    fn render_statusbar(&self, frame: &mut Frame, area: Rect) {
        let shortcuts = self.get_context_shortcuts();
        let statusbar_text = shortcuts.join(" • ");
        
        let statusbar = Paragraph::new(statusbar_text)
            .style(Typography::secondary())
            .alignment(ratatui::layout::Alignment::Center);
            
        frame.render_widget(statusbar, area);
    }

    /// Get context-aware shortcuts based on focused panel
    fn get_context_shortcuts(&self) -> Vec<String> {
        let mut shortcuts = Vec::new();
        
        // Global shortcuts
        shortcuts.push("Tab: Host".to_string());
        shortcuts.push("↑↓/jk: Select".to_string());
        shortcuts.push("r: Rebuild".to_string());
        shortcuts.push("s/S: Start/Stop".to_string());
        shortcuts.push("J: Logs".to_string());
        
        // Always show quit
        shortcuts.push("q: Quit".to_string());
        
        shortcuts
    }

    fn render_system_panel(&mut self, frame: &mut Frame, area: Rect, _metric_store: &MetricStore) {
        let system_block = Components::widget_block("system");
        let inner_area = system_block.inner(area);
        frame.render_widget(system_block, area);
        // Get current host widgets, create if none exist
        if let Some(hostname) = self.current_host.clone() {
            let scroll_offset = {
                let host_widgets = self.get_or_create_host_widgets(&hostname);
                host_widgets.system_scroll_offset
            };
            let host_widgets = self.get_or_create_host_widgets(&hostname);
            host_widgets.system_widget.render_with_scroll(frame, inner_area, scroll_offset, &hostname);
        }
    }

    fn render_backup_panel(&mut self, frame: &mut Frame, area: Rect) {
        let backup_block = Components::widget_block("backup");
        let inner_area = backup_block.inner(area);
        frame.render_widget(backup_block, area);

        // Get current host widgets for backup widget
        if let Some(hostname) = self.current_host.clone() {
            let scroll_offset = {
                let host_widgets = self.get_or_create_host_widgets(&hostname);
                host_widgets.backup_scroll_offset
            };
            let host_widgets = self.get_or_create_host_widgets(&hostname);
            host_widgets.backup_widget.render_with_scroll(frame, inner_area, scroll_offset);
        }
    }


}