//! Terminal registry - manages multiple terminal instances

use crate::pty::terminal::{Terminal, types::TerminalCommandResult};
use kodegen_mcp_schema::terminal::TerminalSnapshot;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

type TerminalMap = HashMap<(String, u32), Arc<Terminal>>;

/// Registry for managing multiple terminal instances keyed by (connection_id, terminal_id)
#[derive(Clone)]
pub struct TerminalRegistry {
    terminals: Arc<Mutex<TerminalMap>>,
}

impl TerminalRegistry {
    /// Create a new terminal registry
    pub fn new() -> Self {
        Self {
            terminals: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Find or create a terminal instance
    pub async fn find_or_create_terminal(
        &self,
        connection_id: &str,
        terminal_id: u32,
        pwd: Option<String>,
    ) -> Result<Arc<Terminal>, anyhow::Error> {
        let key = (connection_id.to_string(), terminal_id);
        let mut terminals = self.terminals.lock().await;

        if let Some(terminal) = terminals.get(&key) {
            return Ok(terminal.clone());
        }

        let mut builder = Terminal::builder()
            .terminal_id(terminal_id)
            .size(2000, 120);  // Force dimensions: 2000 rows x 120 cols

        if let Some(cwd) = pwd {
            builder = builder.cwd(cwd);
        }

        let terminal = Arc::new(builder.build().await?);

        terminals.insert(key, terminal.clone());
        Ok(terminal)
    }

    /// List all active terminals for a connection with their current states
    pub async fn list_all_terminals(
        &self,
        connection_id: &str,
    ) -> Result<TerminalCommandResult, anyhow::Error> {
        let start = std::time::Instant::now();
        let terminals = self.terminals.lock().await;

        let mut snapshots = Vec::new();
        for ((conn_id, term_id), terminal) in terminals.iter() {
            if conn_id == connection_id {
                // Get current state without blocking (use default tail of 2000)
                let state = terminal.read_current_state(*term_id, 2000).await?;
                snapshots.push(TerminalSnapshot {
                    terminal: *term_id,
                    cwd: state.cwd,
                    exit_code: state.exit_code,
                    completed: state.completed,
                });
            }
        }

        // Sort by terminal ID
        snapshots.sort_by_key(|s| s.terminal);

        let output = serde_json::to_string_pretty(&snapshots)?;

        Ok(TerminalCommandResult {
            terminal: None, // None indicates LIST response with multiple terminals
            output,
            exit_code: Some(0),
            cwd: "/".to_string(),
            duration_ms: start.elapsed().as_millis() as u64,
            completed: true,
            terminals: snapshots,
        })
    }

    /// Kill a terminal and cleanup all resources
    pub async fn kill_terminal(
        &self,
        connection_id: &str,
        terminal_id: u32,
    ) -> Result<TerminalCommandResult, anyhow::Error> {
        let start = std::time::Instant::now();
        let key = (connection_id.to_string(), terminal_id);
        let mut terminals = self.terminals.lock().await;

        if let Some(terminal_arc) = terminals.remove(&key) {
            // Drop lock before awaiting shutdown
            drop(terminals);

            // Unwrap Arc and call explicit async shutdown
            match Arc::try_unwrap(terminal_arc) {
                Ok(terminal) => terminal.shutdown().await,
                Err(arc) => {
                    log::warn!(
                        "Terminal {} still has {} references, cannot shutdown cleanly",
                        terminal_id,
                        Arc::strong_count(&arc)
                    );
                }
            }

            Ok(TerminalCommandResult {
                terminal: Some(terminal_id),
                output: format!("Terminal {} shutdown complete", terminal_id),
                exit_code: Some(0),
                cwd: "/".to_string(),
                duration_ms: start.elapsed().as_millis() as u64,
                completed: true,
                terminals: Vec::new(),
            })
        } else {
            Err(anyhow::anyhow!(
                "Terminal {} not found for connection {}",
                terminal_id,
                connection_id
            ))
        }
    }

    /// Cleanup all terminals for a connection
    ///
    /// Called when a connection drops to cleanup all associated terminal sessions.
    /// Returns the number of terminals cleaned up.
    pub async fn cleanup_connection(&self, connection_id: &str) -> usize {
        let mut terminals = self.terminals.lock().await;

        // Collect terminal IDs to remove
        let to_remove: Vec<(String, u32)> = terminals
            .keys()
            .filter(|(conn_id, _)| conn_id == connection_id)
            .cloned()
            .collect();

        let count = to_remove.len();

        // Remove and collect Arc<Terminal> instances
        let terminal_arcs: Vec<_> = to_remove
            .into_iter()
            .filter_map(|key| terminals.remove(&key))
            .collect();

        // Drop lock before awaiting shutdowns
        drop(terminals);

        // Shutdown each terminal explicitly
        for terminal_arc in terminal_arcs {
            match Arc::try_unwrap(terminal_arc) {
                Ok(terminal) => {
                    log::debug!("Shutting down terminal for connection {}", connection_id);
                    terminal.shutdown().await;
                }
                Err(arc) => {
                    log::warn!(
                        "Terminal still has {} references, cannot shutdown cleanly",
                        Arc::strong_count(&arc)
                    );
                }
            }
        }

        count
    }
}

impl Default for TerminalRegistry {
    fn default() -> Self {
        Self::new()
    }
}
