cm-dashboard/dashboard/src/app.rs

use anyhow::Result;
use crossterm::{
    event::{self},
    execute,
    terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
};
use ratatui::{backend::CrosstermBackend, Terminal};
use std::io;
use std::time::{Duration, Instant};
use tracing::{debug, error, info, warn};

use crate::communication::{ZmqConsumer};
use crate::config::DashboardConfig;
use crate::metrics::MetricStore;
use crate::ui::{TuiApp, UiCommand};

pub struct Dashboard {
    zmq_consumer: ZmqConsumer,
    metric_store: MetricStore,
    tui_app: Option<TuiApp>,
    terminal: Option<Terminal<CrosstermBackend<io::Stdout>>>,
    headless: bool,
    initial_commands_sent: std::collections::HashSet<String>,
    config: DashboardConfig,
}

impl Dashboard {
    pub async fn new(config_path: Option<String>, headless: bool) -> Result<Self> {
        info!("Initializing dashboard");

        // Load configuration - try default path if not specified
        let config = match config_path {
            Some(path) => DashboardConfig::load_from_file(&path)?,
            None => {
                // Try default NixOS config path
                let default_path = "/etc/cm-dashboard/dashboard.toml";
                match DashboardConfig::load_from_file(default_path) {
                    Ok(config) => {
                        info!("Using default config file: {}", default_path);
                        config
                    }
                    Err(e) => {
                        error!("Configuration file is required. Use --config to specify path or ensure {} exists.", default_path);
                        error!("Failed to load default config: {}", e);
                        return Err(anyhow::anyhow!("Missing required configuration file"));
                    }
                }
            }
        };

        // Initialize ZMQ consumer
        let mut zmq_consumer = match ZmqConsumer::new(&config.zmq).await {
            Ok(consumer) => consumer,
            Err(e) => {
                error!("Failed to initialize ZMQ consumer: {}", e);
                return Err(e);
            }
        };


        // Try to connect to hosts but don't fail if none are available
        match zmq_consumer.connect_to_predefined_hosts(&config.hosts).await {
            Ok(_) => info!("Successfully connected to ZMQ hosts"),
            Err(e) => {
                warn!(
                    "Failed to connect to hosts (this is normal if no agents are running): {}",
                    e
                );
                info!("Dashboard will start anyway and connect when agents become available");
            }
        }

        // Initialize metric store
        let metric_store = MetricStore::new(10000, 24); // 10k metrics, 24h retention

        // Initialize TUI components only if not headless
        let (tui_app, terminal) = if headless {
            info!("Running in headless mode (no TUI)");
            (None, None)
        } else {
            // Initialize TUI app
            let tui_app = TuiApp::new(config.clone());

            // Setup terminal
            if let Err(e) = enable_raw_mode() {
                error!("Failed to enable raw mode: {}", e);
                error!(
                    "This usually means the dashboard is being run without a proper terminal (TTY)"
                );
                error!("Try running with --headless flag or in a proper terminal");
                return Err(e.into());
            }

            let mut stdout = io::stdout();
            if let Err(e) = execute!(stdout, EnterAlternateScreen) {
                error!("Failed to enter alternate screen: {}", e);
                let _ = disable_raw_mode();
                return Err(e.into());
            }

            let backend = CrosstermBackend::new(stdout);
            let terminal = match Terminal::new(backend) {
                Ok(term) => term,
                Err(e) => {
                    error!("Failed to create terminal: {}", e);
                    let _ = disable_raw_mode();
                    return Err(e.into());
                }
            };

            (Some(tui_app), Some(terminal))
        };

        info!("Dashboard initialization complete");

        Ok(Self {
            zmq_consumer,
            metric_store,
            tui_app,
            terminal,
            headless,
            initial_commands_sent: std::collections::HashSet::new(),
            config,
        })
    }


    pub async fn run(&mut self) -> Result<()> {
        info!("Starting dashboard main loop");

        let mut last_metrics_check = Instant::now();
        let metrics_check_interval = Duration::from_millis(100); // Check for metrics every 100ms
        let mut last_heartbeat_check = Instant::now();
        let heartbeat_check_interval = Duration::from_secs(1); // Check for host connectivity every 1 second

        loop {
            // Handle terminal events (keyboard input) only if not headless
            if !self.headless {
                match event::poll(Duration::from_millis(50)) {
                    Ok(true) => {
                        match event::read() {
                            Ok(event) => {
                                if let Some(ref mut tui_app) = self.tui_app {
                                    // Handle input and check for commands
                                    match tui_app.handle_input(event) {
                                        Ok(Some(command)) => {
                                            // Execute the command
                                            if let Err(e) = self.execute_ui_command(command).await {
                                                error!("Failed to execute UI command: {}", e);
                                            }
                                        }
                                        Ok(None) => {
                                            // No command, check if we should quit
                                            if tui_app.should_quit() {
                                                info!("Quit requested, exiting dashboard");
                                                break;
                                            }
                                        }
                                        Err(e) => {
                                            error!("Error handling input: {}", e);
                                        }
                                    }
                                }
                            }
                            Err(e) => {
                                error!("Error reading terminal event: {}", e);
                                break;
                            }
                        }
                    }
                    Ok(false) => {} // No events available (timeout)
                    Err(e) => {
                        error!("Error polling for terminal events: {}", e);
                        break;
                    }
                }

                // Render UI immediately after handling input for responsive feedback
                if let Some(ref mut terminal) = self.terminal {
                    if let Some(ref mut tui_app) = self.tui_app {
                        if let Err(e) = terminal.draw(|frame| {
                            tui_app.render(frame, &self.metric_store);
                        }) {
                            error!("Error rendering TUI after input: {}", e);
                        }
                    }
                }
            }

            // Check for new metrics
            if last_metrics_check.elapsed() >= metrics_check_interval {
                if let Ok(Some(metric_message)) = self.zmq_consumer.receive_metrics().await {
                    debug!(
                        "Received metrics from {}: {} metrics",
                        metric_message.hostname,
                        metric_message.metrics.len()
                    );

                    // Track first contact with host (no command needed - agent sends data every 2s)
                    let is_new_host = !self
                        .initial_commands_sent
                        .contains(&metric_message.hostname);

                    if is_new_host {
                        info!(
                            "First contact with host {} - data will update automatically",
                            metric_message.hostname
                        );
                        self.initial_commands_sent
                            .insert(metric_message.hostname.clone());
                    }

                    // Update metric store
                    self.metric_store
                        .update_metrics(&metric_message.hostname, metric_message.metrics);

                    // Check for agent version mismatches across hosts
                    if let Some((current_version, outdated_hosts)) = self.metric_store.get_version_mismatches() {
                        for outdated_host in &outdated_hosts {
                            warn!("Host {} has outdated agent version (current: {})", outdated_host, current_version);
                        }
                    }

                    // Update TUI with new metrics (only if not headless)
                    if let Some(ref mut tui_app) = self.tui_app {
                        tui_app.update_metrics(&self.metric_store);
                    }
                }
                
                // Also check for command output messages
                if let Ok(Some(cmd_output)) = self.zmq_consumer.receive_command_output().await {
                    debug!(
                        "Received command output from {}: {}",
                        cmd_output.hostname,
                        cmd_output.output_line
                    );

                    // Command output (terminal popup removed - output not displayed)
                }
                
                last_metrics_check = Instant::now();
            }

            // Check for host connectivity changes (heartbeat timeouts) periodically
            if last_heartbeat_check.elapsed() >= heartbeat_check_interval {
                let timeout = Duration::from_secs(self.config.zmq.heartbeat_timeout_seconds);
                
                // Clean up metrics for offline hosts
                self.metric_store.cleanup_offline_hosts(timeout);
                
                if let Some(ref mut tui_app) = self.tui_app {
                    let connected_hosts = self.metric_store.get_connected_hosts(timeout);
                    tui_app.update_hosts(connected_hosts);
                }
                last_heartbeat_check = Instant::now();
            }

            // Render TUI (only if not headless)
            if !self.headless {
                if let Some(ref mut terminal) = self.terminal {
                    if let Some(ref mut tui_app) = self.tui_app {
                        if let Err(e) = terminal.draw(|frame| {
                            tui_app.render(frame, &self.metric_store);
                        }) {
                            error!("Error rendering TUI: {}", e);
                            break;
                        }
                    }
                }
            }

            // Small sleep to prevent excessive CPU usage
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        info!("Dashboard main loop ended");
        Ok(())
    }

    /// Execute a UI command by sending it to the appropriate agent
    async fn execute_ui_command(&self, command: UiCommand) -> Result<()> {
        match command {
            UiCommand::TriggerBackup { hostname } => {
                info!("Trigger backup requested for {}", hostname);
                // TODO: Implement backup trigger command
                info!("Backup trigger not yet implemented");
            }
        }
        Ok(())
    }

}

impl Drop for Dashboard {
    fn drop(&mut self) {
        // Restore terminal (only if not headless)
        if !self.headless {
            let _ = disable_raw_mode();
            if let Some(ref mut terminal) = self.terminal {
                let _ = execute!(terminal.backend_mut(), LeaveAlternateScreen);
                let _ = terminal.show_cursor();
            }
        }
    }
}