zeal/src/interface/uci.rs

use std::{
    collections::HashMap,
    io::{BufRead, Write},
    str::SplitWhitespace,
};

use anyhow::{anyhow, bail};

use crate::{
    board::{
        board::{Board, Color},
        game::Game,
    },
    movegen::r#move::Move,
    search::{
        iterative_deepening, time::TimeInfo, transposition_table::TranspositionTable, INC,
        MAX_DEPTH, TIME,
    },
};

#[derive(PartialEq, Eq, Debug)]
pub enum Command {
    Uci,
    IsReady,
    UciNewGame,
    Position,
    Go,
    Stop,
    Quit,
    SetOption,
}

fn parse_command(parts: &mut SplitWhitespace) -> anyhow::Result<Command> {
    match parts.next() {
        Some("uci") => Ok(Command::Uci),
        Some("isready") => Ok(Command::IsReady),
        Some("ucinewgame") => Ok(Command::UciNewGame),
        Some("position") => Ok(Command::Position),
        Some("go") => Ok(Command::Go),
        Some("stop") => Ok(Command::Stop),
        Some("quit") => Ok(Command::Quit),
        Some("setoption") => Ok(Command::SetOption),
        _ => bail!("Unrecognised command"),
    }
}

pub enum Response {
    UciOk,
    ReadyOk,
    BestMove(String),
    Info(String),
}

pub fn write_response(handle_out: &mut impl Write, response: &Response) -> anyhow::Result<()> {
    writeln!(handle_out, "{response}").map_err(|e| anyhow!(e))?;
    handle_out.flush().map_err(|e| anyhow!(e))?;
    Ok(())
}

use std::fmt;

impl fmt::Display for Response {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UciOk => write!(f, "id name zeal\nid author stefiosif\nuciok"),
            Self::ReadyOk => write!(f, "readyok"),
            Self::BestMove(best_move) => write!(f, "bestmove {best_move}"),
            Self::Info(info) => write!(f, "{info}"),
        }
    }
}

#[derive(Debug)]
struct UciParameters {
    depth: Option<u8>,
    game: Option<Game>,
    wtime: Option<u128>,
    btime: Option<u128>,
    winc: Option<u128>,
    binc: Option<u128>,
    options: HashMap<String, String>,
}

impl UciParameters {
    fn new() -> Self {
        Self {
            depth: None,
            game: None,
            wtime: None,
            btime: None,
            winc: None,
            binc: None,
            options: HashMap::new(),
        }
    }

    fn add_move_time(&mut self, movetime: u128) {
        self.wtime = Some(movetime);
        self.btime = Some(movetime);
    }

    fn add_depth(&mut self, depth: u8) {
        self.depth = Some(depth);
    }

    fn add_wtime(&mut self, wtime: u128) {
        self.wtime = Some(wtime);
    }

    fn add_btime(&mut self, btime: u128) {
        self.btime = Some(btime);
    }

    fn add_winc(&mut self, winc: u128) {
        self.winc = Some(winc);
    }

    fn add_binc(&mut self, binc: u128) {
        self.binc = Some(binc);
    }

    fn add_game(&mut self, game: Game) {
        self.game = Some(game);
    }

    fn add_option(&mut self, name: String, value: String) {
        self.options.insert(name, value);
    }
}

impl Default for UciParameters {
    fn default() -> Self {
        Self::new()
    }
}

pub fn uci_position(position: &mut SplitWhitespace) -> anyhow::Result<Game> {
    let state = position
        .next()
        .ok_or_else(|| anyhow!("Expected startpos or fen"))?;
    let mut game = match state {
        "startpos" => Game::new(),
        "fen" => {
            let fen_parts: Vec<&str> = position.take_while(|&part| part != "moves").collect();
            let fen = fen_parts.join(" ");
            Game::from_fen(&fen).map_err(|e| anyhow!("Failed to parse FEN: {fen}, {e}"))?
        }
        _ => bail!("Expected startpos or fen"),
    };

    if Some("moves") != position.next() {
        return Ok(game);
    }

    for mv_str in position {
        let mv = Move::parse_from_str(&game, mv_str)
            .map_err(|e| anyhow!("Failed to parse move: {e}"))?;
        game.make_move(&mv);
    }

    Ok(game)
}

pub fn uci_go(go_iter: &mut SplitWhitespace, game: &mut Game) -> anyhow::Result<Move> {
    let mut params = UciParameters::new();
    while let Some(subcommand) = go_iter.next() {
        match subcommand {
            "wtime" => params.add_wtime(parse_next(go_iter, "wtime")?),
            "btime" => params.add_btime(parse_next(go_iter, "btime")?),
            "winc" => params.add_winc(parse_next(go_iter, "winc")?),
            "binc" => params.add_binc(parse_next(go_iter, "binc")?),
            "depth" => params.add_depth(parse_next(go_iter, "depth")?),
            "movetime" => params.add_move_time(parse_next(go_iter, "movetime")?),
            _ => (),
        }
    }

    let (time, inc) = match game.current_player() {
        Color::White => (params.wtime.unwrap_or(TIME), params.winc.unwrap_or(INC)),
        Color::Black => (params.btime.unwrap_or(TIME), params.binc.unwrap_or(INC)),
    };

    iterative_deepening::iterative_deepening(
        game,
        params.depth.unwrap_or(MAX_DEPTH),
        &TimeInfo::new(std::time::Instant::now(), time, inc),
    )?
    .ok_or_else(|| anyhow!("No stored best move found"))
}

fn uci_option(option_iter: &mut SplitWhitespace, game: &mut Game) -> anyhow::Result<()> {
    let mut params = UciParameters::new();

    //TODO: Implement the rest of the options and refactor
    while let Some(subcommand) = option_iter.next() {
        if subcommand == "name" {
            let name: String = parse_next(option_iter, "name")?;
            option_iter.next();
            let value: String = parse_next(option_iter, "value")?;
            params.add_option(name, value);
        }
    }

    if let Some(name) = params.options.get("name") {
        if name == "Hash" {
            let value = params.options.get("value").expect("Expected key value");
            let size = value
                .parse::<usize>()
                .map_err(|_| anyhow!("Invalid value for Hash"))?;
            game.tt = TranspositionTable::new_with_mb_size(size);
        }
    }

    Ok(())
}

fn parse_next<T: std::str::FromStr>(iter: &mut SplitWhitespace, val: &str) -> anyhow::Result<T> {
    iter.next()
        .ok_or_else(|| anyhow!("Expected {val}"))
        .and_then(|v| v.parse::<T>().map_err(|_| anyhow!("Invalid {val}")))
}

pub fn uci_loop<R: BufRead, W: Write>(input: R, mut output: W) -> anyhow::Result<()> {
    let mut params = UciParameters::new();

    for line in input.lines() {
        let line_str = line.unwrap_or_else(|_| "quit".to_string());
        let mut parts = line_str.split_whitespace();
        let command = parse_command(&mut parts)?;
        match command {
            Command::Uci => write_response(&mut output, &Response::UciOk)?,
            Command::IsReady => write_response(&mut output, &Response::ReadyOk)?,
            Command::UciNewGame => {
                if let Some(game) = params.game.as_mut() {
                    game.tt = TranspositionTable::new();
                    game.board = Board::startpos();
                } else {
                    let game = Game::new();
                    params.add_game(game);
                }
            }
            Command::Position => params.add_game(uci_position(&mut parts)?),
            Command::Go => {
                if let Some(game) = params.game.as_mut() {
                    match uci_go(&mut parts, game) {
                        Ok(best_move) => write_response(
                            &mut output,
                            &Response::BestMove(best_move.parse_into_str()),
                        ),
                        Err(e) => write_response(&mut output, &Response::Info(e.to_string())),
                    }?;
                } else {
                    write_response(
                        &mut output,
                        &Response::Info("Failed to unwrap from UciParameter".to_string()),
                    )?;
                };
            }
            Command::Stop => break,
            Command::Quit => break,
            Command::SetOption => {
                if let Some(game) = params.game.as_mut() {
                    uci_option(&mut parts, game)?;
                } else {
                    let game = Game::new();
                    params.add_game(game);
                };
            }
        };

        output.flush().map_err(|e| anyhow!(e))?
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use crate::{
        board::{fen::from_fen, square::Square},
        interface::uci::{parse_command, Command},
        movegen::{attack_generator::init_attacks, r#move::Move},
    };

    use super::uci_go;
    use super::uci_loop;
    use super::uci_position;

    const FEN: [&str; 2] = [
        "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1",
        "r1bqkbnr/pppppppp/2n4B/8/3P4/8/PPP1PPPP/RN1QKBNR b KQkq - 2 2",
    ];

    #[test]
    fn test_uci_position() -> anyhow::Result<()> {
        let command_position = "position startpos";
        let mut parts = command_position.split_whitespace();
        let command = parse_command(&mut parts)?;

        assert_eq!(command, Command::Position);
        assert_eq!(from_fen(FEN[0]).unwrap(), uci_position(&mut parts)?);

        let command_position_moves = "position startpos moves d2d4 b8c6 c1h6";
        let mut parts = command_position_moves.split_whitespace();
        let command = parse_command(&mut parts)?;

        assert_eq!(command, Command::Position);
        assert_eq!(from_fen(FEN[1]).unwrap(), uci_position(&mut parts)?);

        Ok(())
    }

    const FEN_MATE_IN_1: &str = "8/8/8/8/8/4q1k1/8/5K2 b - - 0 1";

    #[test]
    fn test_uci_go() -> anyhow::Result<()> {
        init_attacks();
        let mut game = from_fen(FEN_MATE_IN_1).unwrap();
        let command_go = "go depth 2";
        let mut parts = command_go.split_whitespace();
        let response = uci_go(&mut parts, &mut game)?;

        assert_eq!(Move::new(Square::E3, Square::F2), response);

        Ok(())
    }

    #[test]
    fn test_uci_loop() -> anyhow::Result<()> {
        init_attacks();
        let commands = "uci\n\
            ucinewgame\n\
            position fen 8/8/8/8/8/4q1k1/8/5K2 b - - 0 1\n\
            go\n\
            quit";
        let input = Cursor::new(commands);
        let mut output: Vec<_> = vec![];

        uci_loop(input, &mut output)?;

        let expected_response = "id name zeal\n\
            id author stefiosif\n\
            uciok\n\
            bestmove e3f2\n";
        let actual_response = String::from_utf8(output).expect("Invalid UTF-8 in output");

        assert_eq!(expected_response, actual_response);

        Ok(())
    }
}