scroll-chat 0.1.1

//! SSH Server implementation using russh

use crate::room::{ChatRoom, RoomEvent};
use crate::tui::{render_frame, ChatTui};
use async_trait::async_trait;
use russh::server::{Auth, Handle, Handler, Msg, Session};
use russh::{Channel, ChannelId, CryptoVec, Pty};
use russh_keys::key::KeyPair;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, Mutex, RwLock};
use tracing::{debug, error, info, warn};

/// The SSH server that manages chat sessions
#[derive(Clone)]
pub struct ChatServer {
    /// Shared chat room state
    room: Arc<ChatRoom>,
    /// Data senders for each client (client_id -> sender)
    data_senders: Arc<Mutex<HashMap<usize, mpsc::Sender<Vec<u8>>>>>,
    /// Next client ID
    next_id: Arc<RwLock<usize>>,
}

impl ChatServer {
    /// Create a new chat server
    pub fn new(room: Arc<ChatRoom>) -> Self {
        Self {
            room,
            data_senders: Arc::new(Mutex::new(HashMap::new())),
            next_id: Arc::new(RwLock::new(1)),
        }
    }

    /// Run the SSH server
    pub async fn run(self, port: u16) -> anyhow::Result<()> {
        // Generate ephemeral host key (Ed25519)
        let key_pair = KeyPair::generate_ed25519().expect("Failed to generate host key");
        
        // Log the host key fingerprint
        let fingerprint = key_pair.clone_public_key()
            .map(|k| k.fingerprint())
            .unwrap_or_else(|_| "unknown".to_string());
        info!("Host key fingerprint: {}", fingerprint);

        let mut config = russh::server::Config::default();
        config.keys.push(key_pair);
        config.auth_rejection_time = Duration::from_secs(3);
        config.auth_rejection_time_initial = Some(Duration::from_secs(0));
        config.inactivity_timeout = Some(Duration::from_secs(3600)); // 1 hour timeout
        
        let config = Arc::new(config);
        let addr = format!("0.0.0.0:{}", port);
        
        info!("Starting SSH server on {}", addr);
        russh::server::run(config, &addr, self).await?;
        
        Ok(())
    }
}

impl russh::server::Server for ChatServer {
    type Handler = ChatHandler;

    fn new_client(&mut self, peer_addr: Option<std::net::SocketAddr>) -> Self::Handler {
        let client_id = {
            let mut id = futures::executor::block_on(self.next_id.write());
            let current = *id;
            *id += 1;
            current
        };

        info!("New connection from {:?} (client_id: {})", peer_addr, client_id);
        
        ChatHandler {
            client_id,
            room: self.room.clone(),
            data_senders: self.data_senders.clone(),
            session_id: None,
            username: None,
            channel_id: None,
            pty_size: (80, 24),
            authenticated: false,
        }
    }
}

/// Handler for individual SSH client connections
pub struct ChatHandler {
    client_id: usize,
    room: Arc<ChatRoom>,
    data_senders: Arc<Mutex<HashMap<usize, mpsc::Sender<Vec<u8>>>>>,
    session_id: Option<u64>,
    username: Option<String>,
    channel_id: Option<ChannelId>,
    pty_size: (u16, u16),
    authenticated: bool,
}

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

    /// Handle password authentication
    async fn auth_password(
        mut self,
        user: &str,
        password: &str,
    ) -> Result<(Self, Auth), Self::Error> {
        info!("Password auth attempt for user: {}", user);
        
        if self.room.verify_password(password) {
            self.username = Some(user.to_string());
            self.authenticated = true;
            info!("User {} authenticated successfully", user);
            Ok((self, Auth::Accept))
        } else {
            warn!("Authentication failed for user: {}", user);
            Ok((self, Auth::Reject { proceed_with_methods: None }))
        }
    }

    /// Handle channel open request
    async fn channel_open_session(
        mut self,
        channel: Channel<Msg>,
        session: Session,
    ) -> Result<(Self, bool, Session), Self::Error> {
        if !self.authenticated {
            return Ok((self, false, session));
        }

        self.channel_id = Some(channel.id());
        Ok((self, true, session))
    }

    /// Handle PTY request
    async fn pty_request(
        mut self,
        _channel_id: ChannelId,
        term: &str,
        col_width: u32,
        row_height: u32,
        _pix_width: u32,
        _pix_height: u32,
        _modes: &[(Pty, u32)],
        session: Session,
    ) -> Result<(Self, Session), Self::Error> {
        info!(
            "PTY request: term={}, size={}x{}",
            term, col_width, row_height
        );

        self.pty_size = (col_width as u16, row_height as u16);
        Ok((self, session))
    }

    /// Handle shell request - start the chat session
    async fn shell_request(
        mut self,
        channel_id: ChannelId,
        session: Session,
    ) -> Result<(Self, Session), Self::Error> {
        let username = self.username.clone().unwrap_or_else(|| "anonymous".to_string());
        let (cols, rows) = self.pty_size;
        
        // Join the room
        let session_id = self.room.join(username.clone(), cols, rows).await?;
        self.session_id = Some(session_id);
        
        info!("User {} joined room (session_id: {})", username, session_id);
        
        // Create data channel for receiving keystrokes
        let (data_tx, data_rx) = mpsc::channel::<Vec<u8>>(64);
        
        // Store the sender so data() handler can use it
        self.data_senders.lock().await.insert(self.client_id, data_tx);
        
        // Get the session handle for sending data
        let handle = session.handle();
        let room = self.room.clone();
        let client_id = self.client_id;
        let data_senders = self.data_senders.clone();
        
        // Spawn the chat session handler
        tokio::spawn(async move {
            if let Err(e) = run_chat_session(
                handle,
                channel_id,
                room,
                session_id,
                username,
                cols,
                rows,
                data_rx,
            ).await {
                error!("Chat session error: {}", e);
            }
            // Clean up sender when session ends
            data_senders.lock().await.remove(&client_id);
        });
        
        Ok((self, session))
    }

    /// Handle window resize
    async fn window_change_request(
        mut self,
        _channel_id: ChannelId,
        col_width: u32,
        row_height: u32,
        _pix_width: u32,
        _pix_height: u32,
        session: Session,
    ) -> Result<(Self, Session), Self::Error> {
        debug!("Window resize: {}x{}", col_width, row_height);
        
        self.pty_size = (col_width as u16, row_height as u16);
        
        if let Some(session_id) = self.session_id {
            self.room.update_terminal_size(
                session_id,
                col_width as u16,
                row_height as u16,
            ).await;
        }
        
        Ok((self, session))
    }

    /// Handle incoming data from the client - ROUTE TO CHAT SESSION
    async fn data(
        self,
        _channel_id: ChannelId,
        data: &[u8],
        session: Session,
    ) -> Result<(Self, Session), Self::Error> {
        // Send data to the chat session via the channel
        if let Some(sender) = self.data_senders.lock().await.get(&self.client_id) {
            let _ = sender.send(data.to_vec()).await;
        }
        Ok((self, session))
    }

    /// Handle channel close
    async fn channel_close(
        self,
        _channel_id: ChannelId,
        session: Session,
    ) -> Result<(Self, Session), Self::Error> {
        info!("Channel closed for client {}", self.client_id);
        
        if let Some(session_id) = self.session_id {
            self.room.leave(session_id).await;
        }
        
        // Clean up data sender
        self.data_senders.lock().await.remove(&self.client_id);
        
        Ok((self, session))
    }
}

/// Run the interactive chat session for a connected client
async fn run_chat_session(
    handle: Handle,
    channel_id: ChannelId,
    room: Arc<ChatRoom>,
    session_id: u64,
    username: String,
    cols: u16,
    rows: u16,
    mut data_rx: mpsc::Receiver<Vec<u8>>,
) -> anyhow::Result<()> {
    let mut tui = ChatTui::new(room.clone(), session_id, username.clone(), cols, rows);
    let mut rx = room.subscribe();
    
    // Input buffer for multi-byte sequences
    let mut input_buffer = Vec::new();
    
    // Initial render
    let frame = render_frame(&tui).await;
    handle.data(channel_id, CryptoVec::from_slice(&frame))
        .await
        .map_err(|_| anyhow::anyhow!("Failed to send initial frame"))?;
    
    loop {
        tokio::select! {
            // Handle room events
            event = rx.recv() => {
                match event {
                    Ok(RoomEvent::NewMessage(_)) | 
                    Ok(RoomEvent::UserJoined(_)) | 
                    Ok(RoomEvent::UserLeft(_)) |
                    Ok(RoomEvent::Refresh) => {
                        let frame = render_frame(&tui).await;
                        if handle.data(channel_id, CryptoVec::from_slice(&frame)).await.is_err() {
                            break;
                        }
                    }
                    Err(_) => break,
                }
            }
            
            // Handle incoming data from SSH client
            Some(data) = data_rx.recv() => {
                input_buffer.extend_from_slice(&data);
                
                // Process input
                while let Some(msg) = process_input(&mut input_buffer, &mut tui).await {
                    if msg == "__QUIT__" {
                        room.leave(session_id).await;
                        return Ok(());
                    }
                    room.send_message(session_id, msg).await;
                }
                
                // Re-render after input
                let frame = render_frame(&tui).await;
                if handle.data(channel_id, CryptoVec::from_slice(&frame)).await.is_err() {
                    break;
                }
            }
        }
    }
    
    room.leave(session_id).await;
    Ok(())
}

/// Process raw input bytes and convert to key events
async fn process_input(buffer: &mut Vec<u8>, tui: &mut ChatTui) -> Option<String> {
    if buffer.is_empty() {
        return None;
    }
    
    // Simple byte-by-byte processing
    let byte = buffer.remove(0);
    
    let key_event = match byte {
        3 => {
            // Ctrl+C
            return Some("__QUIT__".to_string());
        }
        13 => {
            // Enter
            crossterm::event::KeyEvent::new(
                crossterm::event::KeyCode::Enter,
                crossterm::event::KeyModifiers::NONE,
            )
        }
        127 | 8 => {
            // Backspace
            crossterm::event::KeyEvent::new(
                crossterm::event::KeyCode::Backspace,
                crossterm::event::KeyModifiers::NONE,
            )
        }
        27 => {
            // Escape sequence
            if buffer.len() >= 2 && buffer[0] == b'[' {
                buffer.remove(0); // Remove '['
                match buffer.remove(0) {
                    b'A' => crossterm::event::KeyEvent::new(
                        crossterm::event::KeyCode::Up,
                        crossterm::event::KeyModifiers::NONE,
                    ),
                    b'B' => crossterm::event::KeyEvent::new(
                        crossterm::event::KeyCode::Down,
                        crossterm::event::KeyModifiers::NONE,
                    ),
                    b'C' => crossterm::event::KeyEvent::new(
                        crossterm::event::KeyCode::Right,
                        crossterm::event::KeyModifiers::NONE,
                    ),
                    b'D' => crossterm::event::KeyEvent::new(
                        crossterm::event::KeyCode::Left,
                        crossterm::event::KeyModifiers::NONE,
                    ),
                    b'5' if !buffer.is_empty() && buffer[0] == b'~' => {
                        buffer.remove(0);
                        crossterm::event::KeyEvent::new(
                            crossterm::event::KeyCode::PageUp,
                            crossterm::event::KeyModifiers::NONE,
                        )
                    }
                    b'6' if !buffer.is_empty() && buffer[0] == b'~' => {
                        buffer.remove(0);
                        crossterm::event::KeyEvent::new(
                            crossterm::event::KeyCode::PageDown,
                            crossterm::event::KeyModifiers::NONE,
                        )
                    }
                    _ => return None,
                }
            } else {
                return None;
            }
        }
        c if c >= 32 && c < 127 => {
            // Printable ASCII
            crossterm::event::KeyEvent::new(
                crossterm::event::KeyCode::Char(c as char),
                crossterm::event::KeyModifiers::NONE,
            )
        }
        _ => return None,
    };
    
    tui.handle_key(key_event).await
}