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,
            rejected: false,
            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,
    chargen: Option<Option<ChargenState>>,
    rejected: bool,
    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();
        let saved = state.db.load_player(&self.username);
        let registration_open = state.is_registration_open();
        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 {
            let handle = session.handle();
            let mut state = self.state.lock().await;
            state.load_existing_player(self.id, saved, channel, handle);
            drop(state);

            let msg = format!(
                "{}\r\n",
                ansi::system_msg("Welcome back! Your character has been restored.")
            );
            self.send_text(session, channel, &msg);
            self.chargen = Some(None);
            self.enter_world(session, channel).await;
        } else if !registration_open {
            let msg = format!(
                "{}\r\n{}\r\n",
                ansi::error_msg("Registration is currently closed. New characters cannot be created."),
                ansi::system_msg("Contact an administrator for access. Disconnecting..."),
            );
            self.send_text(session, channel, &msg);
            self.rejected = true;
        } else {
            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,
        };

        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();

        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> {
        if self.rejected {
            session.close(channel)?;
            return Ok(());
        }

        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 {
                        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);
        });
    }
}