regy 0.1.1

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64};

use crate::domain::errors::{AgentError, AgentResult, ErrorCode};
use crate::domain::policy::CommandPolicy;
use crate::domain::resize::TerminalSize;
use crate::infrastructure::pty_session::PtySession;

#[derive(Clone)]
pub(crate) struct SessionOutput {
    pub session_id: String,
    pub session: Arc<PtySession>,
    pub next_seq: Arc<AtomicU64>,
    pub active: Arc<AtomicBool>,
    pub exit_queued: Arc<AtomicBool>,
}

pub(crate) struct SessionSpawn {
    pub session_id: String,
    pub command: Vec<String>,
    pub cwd: Option<String>,
}

struct SessionEntry {
    session: Arc<PtySession>,
    next_seq: Arc<AtomicU64>,
    active: Arc<AtomicBool>,
    exit_queued: Arc<AtomicBool>,
}

pub struct SessionManager {
    sessions: HashMap<String, SessionEntry>,
    policy: CommandPolicy,
    default_command: Vec<String>,
}

#[allow(dead_code)]
impl SessionManager {
    pub fn new(policy: CommandPolicy, default_command: Vec<String>) -> Self {
        Self {
            sessions: HashMap::new(),
            policy,
            default_command,
        }
    }

    /// Returns a clone of the current command policy. Used to share the same
    /// allowlist between the session manager and other read-only
    /// infrastructure (e.g. `DirectoryBrowser`).
    pub fn policy_snapshot(&self) -> CommandPolicy {
        self.policy.clone()
    }

    pub async fn create_session(
        &mut self,
        session_id: String,
        command: Option<Vec<String>>,
        cwd: Option<String>,
        size: TerminalSize,
    ) -> AgentResult<u32> {
        let spawn = self.prepare_session(session_id, command, cwd)?;
        let session =
            PtySession::spawn(spawn.session_id.clone(), spawn.command, spawn.cwd, size).await?;
        let pid = session.pid();
        self.insert_session(spawn.session_id, session)?;
        Ok(pid)
    }

    pub(crate) fn prepare_session(
        &self,
        session_id: String,
        command: Option<Vec<String>>,
        cwd: Option<String>,
    ) -> AgentResult<SessionSpawn> {
        if self.sessions.contains_key(&session_id) {
            return Err(session_exists());
        }
        let command = command.unwrap_or_else(|| self.default_command.clone());
        let validated = self.policy.validate(&command, cwd.as_deref())?;
        Ok(SessionSpawn {
            session_id,
            command: validated.command,
            cwd: validated.cwd,
        })
    }

    pub(crate) fn insert_session(
        &mut self,
        session_id: String,
        session: PtySession,
    ) -> AgentResult<SessionOutput> {
        if self.sessions.contains_key(&session_id) {
            return Err(session_exists());
        }
        let session = Arc::new(session);
        let next_seq = Arc::new(AtomicU64::new(0));
        let active = Arc::new(AtomicBool::new(true));
        let exit_queued = Arc::new(AtomicBool::new(false));
        self.sessions.insert(
            session_id.clone(),
            SessionEntry {
                session: session.clone(),
                next_seq: next_seq.clone(),
                active: active.clone(),
                exit_queued: exit_queued.clone(),
            },
        );
        Ok(SessionOutput {
            session_id,
            session,
            next_seq,
            active,
            exit_queued,
        })
    }

    pub(crate) fn session_output(&self, session_id: &str) -> AgentResult<SessionOutput> {
        self.sessions
            .get(session_id)
            .map(|entry| SessionOutput {
                session_id: session_id.to_string(),
                session: entry.session.clone(),
                next_seq: entry.next_seq.clone(),
                active: entry.active.clone(),
                exit_queued: entry.exit_queued.clone(),
            })
            .ok_or_else(session_missing)
    }

    pub(crate) fn output_sessions(&self) -> Vec<SessionOutput> {
        self.sessions
            .iter()
            .map(|(session_id, entry)| SessionOutput {
                session_id: session_id.clone(),
                session: entry.session.clone(),
                next_seq: entry.next_seq.clone(),
                active: entry.active.clone(),
                exit_queued: entry.exit_queued.clone(),
            })
            .collect()
    }

    pub(crate) fn remove_session(&mut self, session_id: &str) -> AgentResult<SessionOutput> {
        self.sessions
            .remove(session_id)
            .map(|entry| SessionOutput {
                session_id: session_id.to_string(),
                session: entry.session,
                next_seq: entry.next_seq,
                active: entry.active,
                exit_queued: entry.exit_queued,
            })
            .ok_or_else(session_missing)
    }

    pub(crate) fn restore_session(&mut self, output: SessionOutput) {
        self.sessions.insert(
            output.session_id,
            SessionEntry {
                session: output.session,
                next_seq: output.next_seq,
                active: output.active,
                exit_queued: output.exit_queued,
            },
        );
    }

    pub async fn write_input(&self, session_id: &str, bytes: &[u8]) -> AgentResult<()> {
        self.session_output(session_id)?
            .session
            .write_input(bytes)
            .await
    }

    pub async fn resize_session(&self, session_id: &str, size: TerminalSize) -> AgentResult<()> {
        self.session_output(session_id)?.session.resize(size).await
    }

    pub async fn kill_session(&mut self, session_id: &str) -> AgentResult<()> {
        let output = self.remove_session(session_id)?;
        if let Err(err) = output.session.kill().await {
            self.restore_session(output);
            return Err(err);
        }
        Ok(())
    }

    pub async fn shutdown(&mut self) {
        for (session_id, entry) in std::mem::take(&mut self.sessions) {
            if entry.session.kill().await.is_err() {
                self.sessions.insert(session_id, entry);
            }
        }
    }
}

fn session_missing() -> AgentError {
    AgentError::new(ErrorCode::SessionNotFound, "session not found")
}

fn session_exists() -> AgentError {
    AgentError::new(ErrorCode::SessionAlreadyExists, "session already exists")
}