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;

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

impl Dashboard {
    pub async fn new(config_path: Option<String>, headless: bool, raw_data: 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,
            raw_data,
            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
                                    match tui_app.handle_input(event) {
                                        Ok(_) => {
                                            // 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(agent_data)) = self.zmq_consumer.receive_agent_data().await {
                    debug!(
                        "Received agent data from {}",
                        agent_data.hostname
                    );

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

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

                    // Show raw data if requested (before processing)
                    if self.raw_data {
                        println!("RAW AGENT DATA FROM {}:", agent_data.hostname);
                        println!("{}", serde_json::to_string_pretty(&agent_data).unwrap_or_else(|e| format!("Serialization error: {}", e)));
                        println!("{}", "─".repeat(80));
                    }

                    // Store structured data directly
                    self.metric_store.store_agent_data(agent_data);

                    // 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(())
    }


}

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();
            }
        }
    }
}