mudserver/src/ssh.rs

use std::sync::atomic::{AtomicUsize, Ordering};

use russh::server::{Auth, Handle, Msg, Server, Session};
use russh::{Channel, ChannelId, CryptoVec, Pty};

use crate::ansi;
use crate::chargen::ChargenState;
use crate::commands;
use crate::game::SharedState;

pub struct MudServer {
    pub state: SharedState,
    next_id: AtomicUsize,
}

impl MudServer {
    pub fn new(state: SharedState) -> Self {
        MudServer { state, next_id: AtomicUsize::new(1) }
    }
}

impl Server for MudServer {
    type Handler = MudHandler;

    fn new_client(&mut self, addr: Option<std::net::SocketAddr>) -> MudHandler {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
        log::info!("New connection (id={id}) from {addr:?}");
        MudHandler {
            id, username: String::new(), channel: None, handle: None,
            line_buffer: String::new(), chargen: None, state: self.state.clone(),
        }
    }

    fn handle_session_error(&mut self, error: <Self::Handler as russh::server::Handler>::Error) {
        log::error!("Session error: {error:#?}");
    }
}

pub struct MudHandler {
    id: usize,
    username: String,
    channel: Option<ChannelId>,
    handle: Option<Handle>,
    line_buffer: String,
    // None = not yet determined, Some(None) = returning player, Some(Some(cg)) = in chargen
    chargen: Option<Option<ChargenState>>,
    state: SharedState,
}

impl MudHandler {
    fn send_text(&self, session: &mut Session, channel: ChannelId, text: &str) {
        let _ = session.data(channel, CryptoVec::from(text.as_bytes()));
    }

    async fn start_session(&mut self, session: &mut Session, channel: ChannelId) {
        let state = self.state.lock().await;
        let world_name = state.world.name.clone();

        // Check if returning player
        let saved = state.db.load_player(&self.username);
        drop(state);

        let welcome = format!(
            "{}\r\n{}Welcome to {}, {}!\r\n",
            ansi::CLEAR_SCREEN, ansi::welcome_banner(),
            ansi::bold(&world_name), ansi::player_name(&self.username),
        );
        self.send_text(session, channel, &welcome);

        if let Some(saved) = saved {
            // Returning player — load from DB
            let handle = session.handle();
            let mut state = self.state.lock().await;
            state.load_existing_player(self.id, saved, channel, handle);

            let msg = format!("{}\r\n", ansi::system_msg("Welcome back! Your character has been restored."));
            drop(state);
            self.send_text(session, channel, &msg);

            self.chargen = Some(None); // signal: no chargen needed
            self.enter_world(session, channel).await;
        } else {
            // New player — start chargen
            let cg = ChargenState::new();
            let state = self.state.lock().await;
            let prompt = cg.prompt_text(&state.world);
            drop(state);
            self.send_text(session, channel, &prompt);
            self.chargen = Some(Some(cg));
        }
    }

    async fn enter_world(&mut self, session: &mut Session, channel: ChannelId) {
        let state = self.state.lock().await;

        let (room_id, player_name) = match state.players.get(&self.id) {
            Some(c) => (c.player.room_id.clone(), c.player.name.clone()),
            None => return,
        };

        // Broadcast arrival
        let arrival = CryptoVec::from(
            format!("\r\n{}\r\n{}", ansi::system_msg(&format!("{player_name} has entered the world.")), ansi::prompt()).as_bytes(),
        );
        let others: Vec<_> = state.players_in_room(&room_id, self.id).iter().map(|c| (c.channel, c.handle.clone())).collect();

        // Render room
        let room_view = render_entry_room(&state, &room_id, &player_name, self.id);
        drop(state);

        self.send_text(session, channel, &room_view);

        for (ch, h) in others {
            let _ = h.data(ch, arrival.clone()).await;
        }
    }

    async fn finish_chargen(&mut self, race_id: String, class_id: String, session: &mut Session, channel: ChannelId) {
        let handle = session.handle();
        let mut state = self.state.lock().await;

        let race_name = state.world.races.iter().find(|r| r.id == race_id).map(|r| r.name.clone()).unwrap_or_default();
        let class_name = state.world.classes.iter().find(|c| c.id == class_id).map(|c| c.name.clone()).unwrap_or_default();

        state.create_new_player(self.id, self.username.clone(), race_id, class_id, channel, handle);
        state.save_player_to_db(self.id);

        drop(state);

        let msg = format!("\r\n{}\r\n\r\n", ansi::system_msg(&format!("Character created: {} the {} {}", self.username, race_name, class_name)));
        self.send_text(session, channel, &msg);

        self.chargen = Some(None);
        self.enter_world(session, channel).await;
    }

    async fn handle_disconnect(&self) {
        let mut state = self.state.lock().await;
        if let Some(conn) = state.remove_player(self.id) {
            let departure = CryptoVec::from(
                format!("\r\n{}\r\n{}", ansi::system_msg(&format!("{} has left the world.", conn.player.name)), ansi::prompt()).as_bytes(),
            );
            let others: Vec<_> = state.players_in_room(&conn.player.room_id, self.id).iter().map(|c| (c.channel, c.handle.clone())).collect();
            drop(state);
            for (ch, h) in others { let _ = h.data(ch, departure.clone()).await; }
            log::info!("{} disconnected (id={})", conn.player.name, self.id);
        }
    }
}

fn render_entry_room(state: &crate::game::GameState, room_id: &str, player_name: &str, player_id: usize) -> String {
    let room = match state.world.get_room(room_id) { Some(r) => r, None => return String::new() };
    let mut out = String::new();
    out.push_str(&format!("{} {}\r\n", ansi::room_name(&room.name), ansi::system_msg(&format!("[{}]", room.region))));
    out.push_str(&format!("  {}\r\n", room.description));

    let npc_strs: Vec<String> = room.npcs.iter().filter_map(|id| {
        let npc = state.world.get_npc(id)?;
        let alive = state.npc_instances.get(id).map(|i| i.alive).unwrap_or(true);
        if !alive { return None; }
        let att = state.npc_attitude_toward(id, player_name);
        let color = match att {
            crate::world::Attitude::Friendly => ansi::GREEN,
            crate::world::Attitude::Neutral | crate::world::Attitude::Wary => ansi::YELLOW,
            _ => ansi::RED,
        };
        Some(ansi::color(color, &npc.name))
    }).collect();
    if !npc_strs.is_empty() {
        out.push_str(&format!("\r\n{}{}\r\n", ansi::color(ansi::DIM, "Present: "), npc_strs.join(", ")));
    }

    let others = state.players_in_room(room_id, player_id);
    if !others.is_empty() {
        let names: Vec<String> = others.iter().map(|c| ansi::player_name(&c.player.name)).collect();
        out.push_str(&format!("{}{}\r\n", ansi::color(ansi::GREEN, "Players here: "), names.join(", ")));
    }

    if !room.exits.is_empty() {
        let mut dirs: Vec<&String> = room.exits.keys().collect();
        dirs.sort();
        let dir_strs: Vec<String> = dirs.iter().map(|d| ansi::direction(d)).collect();
        out.push_str(&format!("{} {}\r\n", ansi::color(ansi::DIM, "Exits:"), dir_strs.join(", ")));
    }
    out.push_str(&ansi::prompt());
    out
}

impl russh::server::Handler for MudHandler {
    type Error = russh::Error;

    async fn auth_password(&mut self, user: &str, _password: &str) -> Result<Auth, Self::Error> {
        self.username = user.to_string();
        log::info!("Auth accepted for '{}' (id={})", user, self.id);
        Ok(Auth::Accept)
    }

    async fn auth_publickey(&mut self, user: &str, _key: &russh::keys::ssh_key::PublicKey) -> Result<Auth, Self::Error> {
        self.username = user.to_string();
        Ok(Auth::Accept)
    }

    async fn auth_none(&mut self, user: &str) -> Result<Auth, Self::Error> {
        self.username = user.to_string();
        Ok(Auth::Accept)
    }

    async fn channel_open_session(&mut self, channel: Channel<Msg>, session: &mut Session) -> Result<bool, Self::Error> {
        self.channel = Some(channel.id());
        self.handle = Some(session.handle());
        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: &[(Pty, u32)], session: &mut Session) -> Result<(), Self::Error> {
        session.channel_success(channel)?;
        Ok(())
    }

    async fn shell_request(&mut self, channel: ChannelId, session: &mut Session) -> Result<(), Self::Error> {
        session.channel_success(channel)?;
        self.start_session(session, channel).await;
        Ok(())
    }

    async fn data(&mut self, channel: ChannelId, data: &[u8], session: &mut Session) -> Result<(), Self::Error> {
        for &byte in data {
            match byte {
                3 | 4 => {
                    self.handle_disconnect().await;
                    session.close(channel)?;
                    return Ok(());
                }
                8 | 127 => {
                    if !self.line_buffer.is_empty() {
                        self.line_buffer.pop();
                        session.data(channel, CryptoVec::from(&b"\x08 \x08"[..]))?;
                    }
                }
                b'\r' | b'\n' => {
                    if byte == b'\n' && self.line_buffer.is_empty() { continue; }
                    session.data(channel, CryptoVec::from(&b"\r\n"[..]))?;
                    let line = std::mem::take(&mut self.line_buffer);

                    // Handle chargen flow
                    let mut chargen_done = None;
                    let mut chargen_active = false;
                    if let Some(ref mut chargen_opt) = self.chargen {
                        if let Some(ref mut cg) = chargen_opt {
                            chargen_active = true;
                            let result = {
                                let state = self.state.lock().await;
                                cg.handle_input(&line, &state.world)
                            };
                            let msg_text = match result { Ok(msg) | Err(msg) => msg };
                            let _ = session.data(channel, CryptoVec::from(msg_text.as_bytes()));
                            if cg.is_done() {
                                chargen_done = cg.result();
                            }
                        }
                    }
                    if let Some((race_id, class_id)) = chargen_done {
                        self.chargen = None;
                        self.finish_chargen(race_id, class_id, session, channel).await;
                        continue;
                    }
                    if chargen_active {
                        // Still in chargen, show next prompt
                        if let Some(Some(ref cg)) = self.chargen {
                            let state = self.state.lock().await;
                            let prompt = cg.prompt_text(&state.world);
                            drop(state);
                            let _ = session.data(channel, CryptoVec::from(prompt.as_bytes()));
                        }
                        continue;
                    }
                    if self.chargen.is_none() {
                        continue;
                    }

                    let keep_going = commands::execute(&line, self.id, &self.state, session, channel).await?;
                    if !keep_going {
                        self.handle_disconnect().await;
                        session.close(channel)?;
                        return Ok(());
                    }
                }
                27 => {}
                b if b < 32 => {}
                _ => {
                    self.line_buffer.push(byte as char);
                    session.data(channel, CryptoVec::from(&[byte][..]))?;
                }
            }
        }
        Ok(())
    }

    async fn channel_eof(&mut self, _channel: ChannelId, _session: &mut Session) -> Result<(), Self::Error> {
        self.handle_disconnect().await;
        Ok(())
    }

    async fn channel_close(&mut self, _channel: ChannelId, _session: &mut Session) -> Result<(), Self::Error> {
        self.handle_disconnect().await;
        Ok(())
    }
}

impl Drop for MudHandler {
    fn drop(&mut self) {
        let state = self.state.clone();
        let id = self.id;
        tokio::spawn(async move {
            let mut state = state.lock().await;
            state.remove_player(id);
        });
    }
}