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::{AgentCommand, ServiceAction, ZmqCommandSender, ZmqConsumer};
use crate::config::DashboardConfig;
use crate::metrics::MetricStore;
use crate::ui::{TuiApp, UiCommand};

pub struct Dashboard {
    zmq_consumer: ZmqConsumer,
    zmq_command_sender: ZmqCommandSender,
    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);
            }
        };

        // Initialize ZMQ command sender
        let zmq_command_sender = match ZmqCommandSender::new(&config.zmq) {
            Ok(sender) => sender,
            Err(e) => {
                error!("Failed to initialize ZMQ command sender: {}", e);
                return Err(e);
            }
        };

        // Connect to predefined hosts from configuration
        let hosts = config.hosts.predefined_hosts.clone();

        // Try to connect to hosts but don't fail if none are available
        match zmq_consumer.connect_to_predefined_hosts(&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();

            // 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,
            zmq_command_sender,
            metric_store,
            tui_app,
            terminal,
            headless,
            initial_commands_sent: std::collections::HashSet::new(),
            config,
        })
    }

    /// Send a command to a specific agent
    pub async fn send_command(&mut self, hostname: &str, command: AgentCommand) -> Result<()> {
        self.zmq_command_sender
            .send_command(hostname, command)
            .await
    }

    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

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

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

                    // Check if this is the first time we've seen this host
                    let is_new_host = !self
                        .initial_commands_sent
                        .contains(&metric_message.hostname);

                    if is_new_host {
                        info!(
                            "First contact with host {}, sending initial CollectNow command",
                            metric_message.hostname
                        );

                        // Send CollectNow command for immediate refresh
                        if let Err(e) = self
                            .send_command(&metric_message.hostname, AgentCommand::CollectNow)
                            .await
                        {
                            error!(
                                "Failed to send initial CollectNow command to {}: {}",
                                metric_message.hostname, e
                            );
                        } else {
                            info!(
                                "✓ Sent initial CollectNow command to {}",
                                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);

                    // Update TUI with new hosts and metrics (only if not headless)
                    if let Some(ref mut tui_app) = self.tui_app {
                        let mut connected_hosts = self
                            .metric_store
                            .get_connected_hosts(Duration::from_secs(30));
                        
                        // Add hosts that are rebuilding but may be temporarily disconnected
                        // Use extended timeout (5 minutes) for rebuilding hosts
                        let rebuilding_hosts = self
                            .metric_store
                            .get_connected_hosts(Duration::from_secs(300));
                        
                        for host in rebuilding_hosts {
                            if !connected_hosts.contains(&host) {
                                // Check if this host is rebuilding in the UI
                                if tui_app.is_host_rebuilding(&host) {
                                    connected_hosts.push(host);
                                }
                            }
                        }
                        
                        tui_app.update_hosts(connected_hosts);
                        tui_app.update_metrics(&self.metric_store);
                    }
                }
                last_metrics_check = Instant::now();
            }

            // Render TUI (only if not headless)
            if !self.headless {
                if let (Some(ref mut terminal), Some(ref mut tui_app)) =
                    (&mut self.terminal, &mut 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::ServiceRestart { hostname, service_name } => {
                info!("Sending restart command for service {} on {}", service_name, hostname);
                let agent_command = AgentCommand::ServiceControl {
                    service_name,
                    action: ServiceAction::Restart,
                };
                self.zmq_command_sender.send_command(&hostname, agent_command).await?;
            }
            UiCommand::ServiceStart { hostname, service_name } => {
                info!("Sending start command for service {} on {}", service_name, hostname);
                let agent_command = AgentCommand::ServiceControl {
                    service_name: service_name.clone(),
                    action: ServiceAction::Start,
                };
                self.zmq_command_sender.send_command(&hostname, agent_command).await?;
            }
            UiCommand::ServiceStop { hostname, service_name } => {
                info!("Sending stop command for service {} on {}", service_name, hostname);
                let agent_command = AgentCommand::ServiceControl {
                    service_name: service_name.clone(),
                    action: ServiceAction::Stop,
                };
                self.zmq_command_sender.send_command(&hostname, agent_command).await?;
            }
            UiCommand::SystemRebuild { hostname } => {
                info!("Sending system rebuild command to {}", hostname);
                let agent_command = AgentCommand::SystemRebuild { 
                    git_url: self.config.system.nixos_config_git_url.clone(),
                    git_branch: self.config.system.nixos_config_branch.clone(),
                    working_dir: self.config.system.nixos_config_working_dir.clone(),
                    api_key_file: self.config.system.nixos_config_api_key_file.clone(),
                };
                self.zmq_command_sender.send_command(&hostname, agent_command).await?;
            }
            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();
            }
        }
    }
}