//! SSH session handler: authentication, PTY, shell, SFTP, and input dispatch.

use std::collections::HashMap;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;

use russh::keys::{HashAlg, PublicKey};
use russh::server::{Auth, Msg, Session};
use russh::{Channel, ChannelId};
use tokio::sync::Mutex;

use crate::api::{MnwApiClient, UserInfo};
use crate::ssh::git;
use crate::ssh::sftp::SftpSession;
use crate::ssh::terminal::TerminalHandle;
use crate::staging;
use crate::tui;

/// Per-connection handler. Created by `MnwServer::new_client()`.
pub struct MnwHandler {
    api: MnwApiClient,
    peer_addr: Option<SocketAddr>,
    staging_dir: Arc<PathBuf>,
    git_user: Arc<str>,
    /// Populated after successful auth.
    user: Option<UserInfo>,
    /// Terminal dimensions (cols, rows).
    term_size: (u16, u16),
    /// TUI application handle for forwarding keypresses.
    app: Option<tui::AppHandle>,
    /// Channels stored between open and shell/subsystem request.
    channels: Arc<Mutex<HashMap<ChannelId, Channel<Msg>>>>,
    /// Active git subprocess stdin handles, keyed by channel.
    git_processes: HashMap<ChannelId, tokio::process::ChildStdin>,
}

impl MnwHandler {
    pub fn new(
        api: MnwApiClient,
        peer_addr: Option<SocketAddr>,
        staging_dir: Arc<PathBuf>,
        git_user: Arc<str>,
    ) -> Self {
        Self {
            api,
            peer_addr,
            staging_dir,
            git_user,
            user: None,
            term_size: (80, 24),
            app: None,
            channels: Arc::new(Mutex::new(HashMap::new())),
            git_processes: HashMap::new(),
        }
    }
}

impl russh::server::Handler for MnwHandler {
    type Error = anyhow::Error;

    async fn auth_publickey_offered(
        &mut self,
        _user: &str,
        key: &PublicKey,
    ) -> Result<Auth, Self::Error> {
        let fingerprint = key.fingerprint(HashAlg::Sha256).to_string();
        tracing::debug!(%fingerprint, peer = ?self.peer_addr, "key offered");

        match self.api.lookup_ssh_key(&fingerprint).await {
            Ok(Some(info)) => {
                if info.suspended {
                    tracing::warn!(user = %info.username, "suspended user attempted SSH login");
                    return Ok(Auth::Reject {
                        proceed_with_methods: None,
                        partial_success: false,
                    });
                }
                self.user = Some(info);
                Ok(Auth::Accept)
            }
            Ok(None) => {
                tracing::debug!(%fingerprint, "key not found");
                Ok(Auth::Reject {
                    proceed_with_methods: None,
                    partial_success: false,
                })
            }
            Err(e) => {
                tracing::error!(error = ?e, "SSH key lookup failed");
                Ok(Auth::Reject {
                    proceed_with_methods: None,
                    partial_success: false,
                })
            }
        }
    }

    async fn auth_publickey(
        &mut self,
        _user: &str,
        _key: &PublicKey,
    ) -> Result<Auth, Self::Error> {
        // If auth_publickey_offered accepted, the user is already stored.
        if self.user.is_some() {
            Ok(Auth::Accept)
        } else {
            Ok(Auth::Reject {
                proceed_with_methods: None,
                partial_success: false,
            })
        }
    }

    async fn channel_open_session(
        &mut self,
        channel: Channel<Msg>,
        _session: &mut Session,
    ) -> Result<bool, Self::Error> {
        let channel_id = channel.id();
        tracing::debug!(channel = %channel_id, "session channel opened");
        // Store channel for later consumption by shell_request or subsystem_request
        self.channels.lock().await.insert(channel_id, channel);
        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.term_size = (col_width as u16, row_height as u16);
        tracing::debug!(cols = col_width, rows = row_height, "PTY requested");
        Ok(())
    }

    async fn shell_request(
        &mut self,
        channel: ChannelId,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        let Some(ref user) = self.user else {
            tracing::warn!("shell_request without authenticated user");
            session.close(channel)?;
            return Ok(());
        };

        tracing::info!(user = %user.username, "launching TUI");

        let handle = session.handle();
        let terminal_handle = TerminalHandle::new(handle.clone(), channel);
        let (cols, rows) = self.term_size;

        let user_clone = user.clone();
        let staging_dir = staging::user_staging_dir(&self.staging_dir, &user.user_id);
        let app_handle = tui::launch(
            terminal_handle,
            user_clone,
            cols,
            rows,
            handle,
            channel,
            self.api.clone(),
            staging_dir,
        )?;
        self.app = Some(app_handle);

        Ok(())
    }

    async fn subsystem_request(
        &mut self,
        channel_id: ChannelId,
        name: &str,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        if name != "sftp" {
            tracing::debug!(subsystem = name, "unsupported subsystem requested");
            session.close(channel_id)?;
            return Ok(());
        }

        let Some(ref user) = self.user else {
            tracing::warn!("subsystem_request without authenticated user");
            session.close(channel_id)?;
            return Ok(());
        };

        // Take the stored channel for this ID
        let channel = self
            .channels
            .lock()
            .await
            .remove(&channel_id);

        let Some(channel) = channel else {
            tracing::error!("no stored channel for SFTP subsystem");
            session.close(channel_id)?;
            return Ok(());
        };

        let user_staging = staging::user_staging_dir(&self.staging_dir, &user.user_id);
        let sftp_session = SftpSession::new(
            user.user_id.clone(),
            user.creator_tier.clone(),
            user_staging,
        );

        tracing::info!(user = %user.username, "starting SFTP session");

        let stream = channel.into_stream();
        tokio::spawn(async move {
            russh_sftp::server::run(stream, sftp_session).await;
        });

        Ok(())
    }

    async fn exec_request(
        &mut self,
        channel: ChannelId,
        data: &[u8],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        let command_line = String::from_utf8_lossy(data);
        let handle = session.handle();

        let Some(ref user) = self.user else {
            let _ = handle.close(channel).await;
            return Ok(());
        };

        // Git operations: bidirectional streaming via subprocess proxy
        if let Some((operation, raw_path)) = git::parse_git_command(&command_line)
            && let Some((owner, repo_name)) = git::parse_repo_path(raw_path)
        {
            tracing::info!(
                user = %user.username,
                %operation,
                %owner,
                %repo_name,
                "git operation"
            );

            match self
                .api
                .git_authorize(&user.user_id, operation, owner, repo_name)
                .await
            {
                Ok(auth) => {
                    match git::spawn_git_process(
                        &self.git_user,
                        operation,
                        &auth.repo_path,
                        channel,
                        handle.clone(),
                    )
                    .await
                    {
                        Ok(stdin) => {
                            self.git_processes.insert(channel, stdin);
                            session.channel_success(channel)?;
                        }
                        Err(e) => {
                            tracing::error!(error = ?e, "failed to spawn git process");
                            let msg = bytes::Bytes::from(
                                "fatal: internal error\r\n".to_string(),
                            );
                            let _ = handle.data(channel, msg).await;
                            let _ = handle.exit_status_request(channel, 1).await;
                            let _ = handle.eof(channel).await;
                            let _ = handle.close(channel).await;
                        }
                    }
                }
                Err(e) => {
                    let msg =
                        bytes::Bytes::from(format!("fatal: {e}\r\n"));
                    let _ = handle.extended_data(channel, 1, msg).await;
                    let _ = handle.exit_status_request(channel, 1).await;
                    let _ = handle.eof(channel).await;
                    let _ = handle.close(channel).await;
                }
            }
            return Ok(());
        }

        // Check if this looks like a legacy SCP transfer
        if command_line.starts_with("scp ") {
            let msg: &[u8] =
                b"Use scp (not scp -O) or sftp to upload files to cli.makenot.work\r\n";
            let bytes = bytes::Bytes::copy_from_slice(msg);
            tokio::spawn(async move {
                let _ = handle.data(channel, bytes).await;
                let _ = handle.close(channel).await;
            });
            return Ok(());
        }

        // Execute the command
        let user = user.clone();
        let api = self.api.clone();
        let cmd = command_line.to_string();
        tracing::info!(user = %user.username, command = %cmd, "exec command");

        tokio::spawn(async move {
            let output = crate::commands::execute(&cmd, &user, &api).await;
            let bytes = bytes::Bytes::from(output);
            let _ = handle.data(channel, bytes).await;
            let _ = handle.exit_status_request(channel, 0).await;
            let _ = handle.eof(channel).await;
            let _ = handle.close(channel).await;
        });

        Ok(())
    }

    async fn data(
        &mut self,
        channel: ChannelId,
        data: &[u8],
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        if let Some(stdin) = self.git_processes.get_mut(&channel) {
            use tokio::io::AsyncWriteExt;
            let _ = stdin.write_all(data).await;
        } else if let Some(ref app) = self.app {
            app.send_input(data).await;
        }
        Ok(())
    }

    async fn channel_eof(
        &mut self,
        channel: ChannelId,
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        if let Some(stdin) = self.git_processes.remove(&channel) {
            drop(stdin); // Closes pipe → subprocess sees EOF
        }
        Ok(())
    }

    async fn window_change_request(
        &mut self,
        _channel: ChannelId,
        col_width: u32,
        row_height: u32,
        _pix_width: u32,
        _pix_height: u32,
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        self.term_size = (col_width as u16, row_height as u16);
        if let Some(ref app) = self.app {
            app.send_resize(col_width as u16, row_height as u16).await;
        }
        Ok(())
    }

    async fn auth_succeeded(&mut self, _session: &mut Session) -> Result<(), Self::Error> {
        if let Some(ref user) = self.user {
            tracing::info!(user = %user.username, peer = ?self.peer_addr, "authenticated");
        }
        Ok(())
    }
}