zeal/src/interface/uci.rs

use std::{
    io::{BufRead, Write},
    str::SplitWhitespace,
    time::Instant,
};

use anyhow::{anyhow, bail};

use crate::{
    board::{
        board::{Board, Color},
        game::Game,
    },
    movegen::r#move::Move,
    search::{
        iterative_deepening, transposition_table::TranspositionTable, MAX_DEPTH, MAX_TT_SIZE,
        REMAINING_TIME_DEFAULT,
    },
};

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

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("quit") => Ok(Command::Quit),
        _ => bail!("Unrecognised command"),
    }
}

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

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 {
    movetime: Option<usize>,
    depth: Option<u8>,
    game: Option<Game>,
    wtime: Option<u128>,
    btime: Option<u128>,
}

impl UciParameters {
    const fn new() -> Self {
        Self {
            movetime: None,
            depth: None,
            game: None,
            wtime: None,
            btime: None,
        }
    }

    fn add_movetime(&mut self, movetime: usize) {
        self.movetime = 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_game(&mut self, game: Game) {
        self.game = Some(game);
    }
}

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")?),
            "depth" => params.add_depth(parse_next(go_iter, "depth")?),
            "movetime" => params.add_movetime(parse_next(go_iter, "movetime")?),
            _ => (),
        }
    }

    let time = Instant::now();
    let remaining_time = match game.current_player() {
        Color::White => params.wtime.unwrap_or(REMAINING_TIME_DEFAULT),
        Color::Black => params.btime.unwrap_or(REMAINING_TIME_DEFAULT),
    };

    iterative_deepening::iterative_deepening(game, MAX_DEPTH, remaining_time)?.ok_or_else(|| {
        anyhow!(
            "No stored best move found. Time: {}",
            remaining_time - time.elapsed().as_millis()
        )
    })
}

fn parse_next<T: std::str::FromStr>(go_iter: &mut SplitWhitespace, val: &str) -> anyhow::Result<T> {
    go_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)?;
        let response = match command {
            Command::Uci => Response::UciOk,
            Command::IsReady => Response::ReadyOk,
            Command::UciNewGame => {
                params.game.as_mut().map_or_else(
                    || Response::Info("Failed to unwrap from UciParameter".to_string()),
                    |game| {
                        game.tt = TranspositionTable::new(MAX_TT_SIZE);
                        game.board = Board::startpos();
                        Response::Info("Initialized new game".to_string())
                    },
                );
                Response::Info("Clear cache".to_string())
            }
            Command::Position => {
                params.add_game(uci_position(&mut parts)?);
                Response::Info("Initialized position".to_string())
            }
            Command::Go => params.game.as_mut().map_or_else(
                || Response::Info("Failed to unwrap from UciParameter".to_string()),
                |game| match uci_go(&mut parts, game) {
                    Ok(best_move) => Response::BestMove(best_move.parse_into_str()),
                    Err(e) => Response::Info(e.to_string()),
                },
            ),
            // TODO: Command::Stop => (),
            Command::Quit => break,
        };

        write_response(&mut output, &response)?;
        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\
            Clear cache\n\
            Initialized position\n\
            bestmove e3f2\n";
        let actual_response = String::from_utf8(output).expect("Invalid UTF-8 in output");

        assert_eq!(expected_response, actual_response);

        Ok(())
    }

    #[test]
    fn test_cute_chess_bug() -> anyhow::Result<()> {
        init_attacks();
        let commands = "uci\n\
            ucinewgame\n\
            position startpos moves b1c3 g7g6 d2d4 f8g7 e2e4 f7f5 e4f5 e7e6 f5g6 h7g6 g1f3 b8c6 f1c4 d7d6 c1e3 e6e5 d4d5 c6e7 c3b5 c7c6 d5c6 b7c6 b5d6 d8d6 d1d6 c8d7 f3e5 d7e6 d6e6 g8f6 e5g6 f6g8 e6f7 e8d7 f7g7 g8f6 g7e7 d7c8 g6h8 c8b8 e7d8 b8b7 d8f6 b7c7 e3a7 a8a7 h8g6 c7b6 g6e5 a7c7 f6c6 b6a7
            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\
            Clear cache\n\
            Initialized position\n\
            bestmove c6c7\n";
        let actual_response = String::from_utf8(output).expect("Invalid UTF-8 in output");

        assert_eq!(expected_response, actual_response);

        Ok(())
    }
}