//! SSH shell server — interactive governed shell sessions.
//!
//! Provides an SSH server that routes commands through the namespace
//! handlers. In dev/soft mode, accepts any public key authentication.

use std::sync::Arc;

use async_trait::async_trait;
use russh::server::{Auth, Config, Handler, Msg, Server, Session};
use russh::{Channel, ChannelId};
use ssh_key::{Algorithm, PrivateKey};
use tracing::{debug, error, info, warn};

use crate::config::AgentConfig;
use crate::namespace::NamespaceRouter;
use crate::shellstream::{Namespace, ShellstreamResponse, Status};

/// The bascule-agent SSH server.
pub struct AgentSSHServer {
    config: Arc<AgentConfig>,
    russh_config: Arc<Config>,
    router: Arc<NamespaceRouter>,
}

impl AgentSSHServer {
    pub fn new(
        config: Arc<AgentConfig>,
        router: Arc<NamespaceRouter>,
    ) -> anyhow::Result<Self> {
        // Load or generate host key
        let host_key = if let Some(ref path) = config.agent.ssh.host_key_path {
            if path.exists() {
                info!(path = %path.display(), "Loading SSH host key");
                russh_keys::load_secret_key(path, None)?
            } else {
                info!("Generating ephemeral Ed25519 host key");
                let mut rng = rand::thread_rng();
                PrivateKey::random(&mut rng, Algorithm::Ed25519)?
            }
        } else {
            info!("Generating ephemeral Ed25519 host key");
            let mut rng = rand::thread_rng();
            PrivateKey::random(&mut rng, Algorithm::Ed25519)?
        };

        let russh_config = Config {
            keys: vec![host_key],
            ..Default::default()
        };

        Ok(Self {
            config,
            russh_config: Arc::new(russh_config),
            router,
        })
    }

    /// Start the SSH server. Runs until cancelled.
    pub async fn serve(mut self) -> anyhow::Result<()> {
        let addr = &self.config.agent.ssh.listen_addr;
        let listen: std::net::SocketAddr = addr
            .parse()
            .map_err(|e| anyhow::anyhow!("Invalid SSH listen address '{}': {}", addr, e))?;

        info!(addr = %listen, "SSH shell server listening");
        self.run_on_address(self.russh_config.clone(), listen).await?;
        Ok(())
    }
}

impl Server for AgentSSHServer {
    type Handler = AgentShellHandler;

    fn new_client(&mut self, peer_addr: Option<std::net::SocketAddr>) -> AgentShellHandler {
        let addr = peer_addr
            .map(|a| a.to_string())
            .unwrap_or_else(|| "unknown".to_string());

        info!(peer = %addr, "New SSH connection");

        AgentShellHandler {
            router: self.router.clone(),
            dev_mode: self.config.agent.namespaces.backend == "soft",
            peer_addr: addr,
            identity: None,
            terminal_width: 80,
            session_state: None,
        }
    }
}

/// Per-connection SSH handler.
pub struct AgentShellHandler {
    router: Arc<NamespaceRouter>,
    dev_mode: bool,
    peer_addr: String,
    identity: Option<String>,
    terminal_width: u16,
    session_state: Option<ShellSessionState>,
}

struct ShellSessionState {
    input_buffer: String,
    channel_id: ChannelId,
    session_id: [u8; 16],
}

impl AgentShellHandler {
    fn banner(&self) -> String {
        let identity = self.identity.as_deref().unwrap_or("unknown");
        format!(
            "\r\n\x1b[1;36m╔══════════════════════════════════════════╗\r\n\
             ║       bascule-agent governed shell       ║\r\n\
             ╚══════════════════════════════════════════╝\x1b[0m\r\n\
             \r\n\
             Identity: {identity}\r\n\
             Mode:     {mode}\r\n\
             \r\n\
             Type 'help' for available commands.\r\n\r\n",
            mode = if self.dev_mode { "dev (soft)" } else { "live" },
        )
    }

    fn prompt(&self) -> String {
        let identity = self.identity.as_deref().unwrap_or("?");
        format!("\x1b[1;32m{identity}\x1b[0m@\x1b[1;34magent\x1b[0m> ")
    }

    async fn process_line(
        &self,
        line: &str,
        handle: &russh::server::Handle,
        channel_id: ChannelId,
        session_id: &[u8; 16],
    ) -> bool {
        let line = line.trim();

        if line == "exit" || line == "quit" {
            let _ = handle.data(channel_id, "Goodbye.\r\n".into()).await;
            let _ = handle.close(channel_id).await;
            return true;
        }

        if line.is_empty() {
            let _ = handle.data(channel_id, self.prompt().into()).await;
            return false;
        }

        let output = match line {
            "help" => self.cmd_help(),
            "status" => self.cmd_status().await,
            "whoami" => self.cmd_whoami(),
            cmd if cmd.starts_with("audit ") => {
                self.cmd_namespace("audit", &cmd[6..], session_id).await
            }
            cmd if cmd.starts_with("crypto ") => {
                self.cmd_namespace("crypto", &cmd[7..], session_id).await
            }
            cmd if cmd.starts_with("identity ") => {
                self.cmd_namespace("identity", &cmd[9..], session_id).await
            }
            cmd if cmd.starts_with("governance ") => {
                self.cmd_namespace("governance", &cmd[11..], session_id).await
            }
            cmd if cmd.starts_with("secrets ") => {
                self.cmd_namespace("secrets", &cmd[8..], session_id).await
            }
            cmd if cmd.starts_with("attestation ") => {
                self.cmd_namespace("attestation", &cmd[12..], session_id).await
            }
            _ => format!("Unknown command: {}. Type 'help' for available commands.\r\n", line),
        };

        if !output.is_empty() {
            let _ = handle.data(channel_id, output.into()).await;
        }
        let _ = handle.data(channel_id, self.prompt().into()).await;
        false
    }

    fn cmd_help(&self) -> String {
        "\x1b[1mAvailable commands:\x1b[0m\r\n\
         \r\n\
         \x1b[33mShell:\x1b[0m\r\n\
         \x20 help           Show this help message\r\n\
         \x20 status         Show agent status\r\n\
         \x20 whoami         Show current identity\r\n\
         \x20 exit           Close the session\r\n\
         \r\n\
         \x1b[33mNamespaces:\x1b[0m\r\n\
         \x20 audit emit     Emit an audit event\r\n\
         \x20 attestation posture  Show posture level\r\n\
         \x20 crypto hash <data>   Hash data with SHA-256\r\n\
         \x20 governance gate      Check governance gate\r\n\
         \x20 identity whoami      Show identity claims\r\n\
         \x20 secrets get <path>   Look up a secret\r\n\
         \r\n"
            .to_string()
    }

    async fn cmd_status(&self) -> String {
        // Query attestation posture through namespace handler
        let payload = rmp_serde::to_vec(&serde_json::json!({})).unwrap_or_default();
        let session_id = [0u8; 16];
        let resp = self
            .router
            .handle(Namespace::Attestation, 0x01, &payload, &session_id)
            .await;

        let posture: serde_json::Value =
            rmp_serde::from_slice(&resp.payload).unwrap_or(serde_json::Value::Null);

        format!(
            "\x1b[1mAgent Status\x1b[0m\r\n\
             \x20 Mode:     {mode}\r\n\
             \x20 Posture:  {posture}\r\n\
             \x20 Peer:     {peer}\r\n\r\n",
            mode = if self.dev_mode { "soft (dev)" } else { "live" },
            posture = posture.get("level").and_then(|v| v.as_str()).unwrap_or("unknown"),
            peer = self.peer_addr,
        )
    }

    fn cmd_whoami(&self) -> String {
        format!(
            "Identity: {}\r\nMode: {}\r\nPeer: {}\r\n",
            self.identity.as_deref().unwrap_or("unknown"),
            if self.dev_mode { "dev" } else { "live" },
            self.peer_addr,
        )
    }

    async fn cmd_namespace(
        &self,
        ns_name: &str,
        args: &str,
        session_id: &[u8; 16],
    ) -> String {
        let ns = match ns_name {
            "audit" => Namespace::Audit,
            "crypto" => Namespace::Crypto,
            "identity" => Namespace::Identity,
            "governance" => Namespace::Governance,
            "secrets" => Namespace::Secrets,
            "attestation" => Namespace::Attestation,
            _ => return format!("Unknown namespace: {ns_name}\r\n"),
        };

        let (function_id, payload) = match parse_namespace_command(ns_name, args) {
            Ok(v) => v,
            Err(e) => return format!("Error: {e}\r\n"),
        };

        let resp = self.router.handle(ns, function_id, &payload, session_id).await;

        let status = Status::from_u16(resp.status);
        let payload_parsed: serde_json::Value =
            rmp_serde::from_slice(&resp.payload).unwrap_or(serde_json::Value::Null);

        let status_str = match status {
            Some(Status::Ok) => "\x1b[32mOK\x1b[0m",
            Some(Status::Error) => "\x1b[31mERROR\x1b[0m",
            Some(Status::Denied) => "\x1b[31mDENIED\x1b[0m",
            Some(Status::CeremonyRequired) => "\x1b[33mCEREMONY_REQUIRED\x1b[0m",
            None => "UNKNOWN",
        };

        format!(
            "[{status_str}] {}\r\n",
            serde_json::to_string_pretty(&payload_parsed).unwrap_or_else(|_| "{}".to_string())
        )
        .replace('\n', "\r\n")
    }
}

/// Parse a namespace subcommand into (function_id, msgpack payload).
fn parse_namespace_command(ns: &str, args: &str) -> anyhow::Result<(u16, Vec<u8>)> {
    let parts: Vec<&str> = args.splitn(2, ' ').collect();
    let subcmd = parts.first().copied().unwrap_or("");
    let rest = parts.get(1).copied().unwrap_or("");

    match (ns, subcmd) {
        ("audit", "emit") => {
            let payload = serde_json::json!({
                "event_type": "shell.command",
                "subject": "ssh-user",
                "resource": "shell",
                "action": "emit",
                "outcome": "success",
            });
            Ok((0x01, rmp_serde::to_vec(&payload)?))
        }
        ("audit", "anchor") => {
            Ok((0x02, rmp_serde::to_vec(&serde_json::json!({"event_ids": []}))?))
        }
        ("crypto", "hash") => {
            let data = if rest.is_empty() { "test" } else { rest };
            Ok((0x03, rmp_serde::to_vec(&serde_json::json!({"data": data}))?))
        }
        ("identity", "whoami") | ("identity", "authenticate") => {
            Ok((0x01, rmp_serde::to_vec(&serde_json::json!({"token": "ssh-session"}))?))
        }
        ("governance", "gate") => {
            Ok((0x01, rmp_serde::to_vec(&serde_json::json!({
                "subject": "ssh-user",
                "resource": "shell",
                "action": "access",
            }))?))
        }
        ("secrets", "get") => {
            let path = if rest.is_empty() { "test/secret" } else { rest };
            Ok((0x01, rmp_serde::to_vec(&serde_json::json!({"path": path}))?))
        }
        ("attestation", "posture") => {
            Ok((0x01, rmp_serde::to_vec(&serde_json::json!({}))?))
        }
        ("attestation", "sat") => {
            Ok((0x02, rmp_serde::to_vec(&serde_json::json!({}))?))
        }
        _ => anyhow::bail!("Unknown {ns} subcommand: {subcmd}"),
    }
}

#[async_trait]
impl Handler for AgentShellHandler {
    type Error = anyhow::Error;

    async fn auth_publickey(
        &mut self,
        user: &str,
        public_key: &ssh_key::PublicKey,
    ) -> Result<Auth, Self::Error> {
        let fingerprint = public_key.fingerprint(ssh_key::HashAlg::Sha256).to_string();

        if self.dev_mode {
            // Dev mode: accept any key
            info!(user = %user, fingerprint = %fingerprint, "SSH auth accepted (dev mode)");
            self.identity = Some(user.to_string());
            Ok(Auth::Accept)
        } else {
            // TODO: validate against authorized keys store
            warn!(fingerprint = %fingerprint, "SSH auth rejected: live mode not implemented");
            Ok(Auth::Reject {
                proceed_with_methods: None,
            })
        }
    }

    async fn auth_password(
        &mut self,
        user: &str,
        _password: &str,
    ) -> Result<Auth, Self::Error> {
        if self.dev_mode {
            info!(user = %user, "SSH password auth accepted (dev mode)");
            self.identity = Some(user.to_string());
            Ok(Auth::Accept)
        } else {
            Ok(Auth::Reject {
                proceed_with_methods: None,
            })
        }
    }

    async fn channel_open_session(
        &mut self,
        _channel: Channel<Msg>,
        _session: &mut Session,
    ) -> Result<bool, Self::Error> {
        Ok(true)
    }

    async fn pty_request(
        &mut self,
        _channel: ChannelId,
        _term: &str,
        col_width: u32,
        _row_height: u32,
        _pix_width: u32,
        _pix_height: u32,
        _modes: &[(russh::Pty, u32)],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        self.terminal_width = col_width.min(200) as u16;
        session.request_success();
        Ok(())
    }

    async fn shell_request(
        &mut self,
        channel: ChannelId,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        session.request_success();

        let session_id: [u8; 16] = uuid::Uuid::new_v4().into_bytes();

        let banner = self.banner();
        let prompt = self.prompt();
        session.data(channel, format!("{banner}{prompt}").into())?;

        self.session_state = Some(ShellSessionState {
            input_buffer: String::new(),
            channel_id: channel,
            session_id,
        });

        Ok(())
    }

    async fn data(
        &mut self,
        channel: ChannelId,
        data: &[u8],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        // Take session state temporarily to avoid borrow conflicts
        let mut state = match self.session_state.take() {
            Some(s) => s,
            None => return Ok(()),
        };

        for &byte in data {
            match byte {
                // Enter
                b'\r' | b'\n' => {
                    session.data(channel, "\r\n".into())?;
                    let line = state.input_buffer.clone();
                    state.input_buffer.clear();

                    let handle = session.handle();
                    let sid = state.session_id;
                    if self.process_line(&line, &handle, channel, &sid).await {
                        // Session ended — don't put state back
                        return Ok(());
                    }
                }
                // Backspace / DEL
                0x7f | 0x08 => {
                    if !state.input_buffer.is_empty() {
                        state.input_buffer.pop();
                        session.data(channel, "\x08 \x08".into())?;
                    }
                }
                // Ctrl-C
                0x03 => {
                    state.input_buffer.clear();
                    session.data(channel, "^C\r\n".into())?;
                    let prompt = self.prompt();
                    session.data(channel, prompt.into())?;
                }
                // Ctrl-D (EOF)
                0x04 => {
                    session.data(channel, "\r\nGoodbye.\r\n".into())?;
                    session.close(channel)?;
                    return Ok(());
                }
                // Printable characters
                b if b >= 0x20 && b < 0x7f => {
                    state.input_buffer.push(b as char);
                    session.data(channel, vec![b].into())?;
                }
                // Ignore other control characters
                _ => {}
            }
        }

        // Put session state back
        self.session_state = Some(state);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_namespace_commands() {
        let (fn_id, _) = parse_namespace_command("audit", "emit").unwrap();
        assert_eq!(fn_id, 0x01);

        let (fn_id, _) = parse_namespace_command("crypto", "hash test data").unwrap();
        assert_eq!(fn_id, 0x03);

        let (fn_id, _) = parse_namespace_command("governance", "gate").unwrap();
        assert_eq!(fn_id, 0x01);

        let (fn_id, _) = parse_namespace_command("attestation", "posture").unwrap();
        assert_eq!(fn_id, 0x01);

        let (fn_id, _) = parse_namespace_command("secrets", "get db/password").unwrap();
        assert_eq!(fn_id, 0x01);

        assert!(parse_namespace_command("audit", "nonexistent").is_err());
    }

    #[test]
    fn test_parse_crypto_hash_payload() {
        let (_, payload) = parse_namespace_command("crypto", "hash hello world").unwrap();
        let parsed: serde_json::Value = rmp_serde::from_slice(&payload).unwrap();
        assert_eq!(parsed["data"], "hello world");
    }

    #[test]
    fn test_parse_secrets_get_path() {
        let (_, payload) = parse_namespace_command("secrets", "get db/password").unwrap();
        let parsed: serde_json::Value = rmp_serde::from_slice(&payload).unwrap();
        assert_eq!(parsed["path"], "db/password");
    }
}