jja 0.9.1

//
// jja: swiss army knife for chess file formats
// src/obk.rs: Interface to Chessmaster book files (aka `obk')
//
// Copyright (c) 2023 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{
    collections::{HashMap, HashSet, VecDeque},
    fs::File,
    io::{Error, ErrorKind},
    path::Path,
};

use indicatif::ProgressBar;
use memmap::Mmap;
use shakmaty::{
    fen::{Epd, Fen},
    san::San,
    uci::Uci,
    Chess, Color, EnPassantMode, Position, Square,
};
use termtree::Tree;

use crate::{
    hash::{
        zobrist_hash, BookEntryHashMap, ObkMoveEntryHashMap, ZobristHashSet, ZobristHasherBuilder,
    },
    obk::{uci_from_obk_move_and_position, ObkMoveEntry, ObkNoteEntry},
    polyglot::{from_uci, is_king_on_start_square, BookEntry},
    tr,
};

/// A structure representing a Chessmaster OBK opening book with associated files and metadata.
pub struct ObkBook {
    /// The version of the OBK file.
    pub version: u8,
    /// The total number of moves in the opening book.
    pub move_count: u32,
    /// The total number of bytes of names and notes in the opening book.
    pub text_notes_count: Option<u32>,
    /// The memory-mapped data of the opening book.
    book: Mmap,
    /// A list of move entries in the book.
    pub moves: Vec<ObkMoveEntry>,
    /// A list of notes associated with moves in the book.
    pub notes: Vec<ObkNoteEntry>,
}

impl ObkBook {
    /// Reads an OBK file into the current `ObkBook` instance.
    ///
    /// This function will validate if the file is a valid Chessmaster book by
    /// checking the header bytes and verifying that the book has moves.
    ///
    /// # Arguments
    ///
    /// * `file_name` - The name of the OBK file to read.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be opened, read, or if the file is
    /// not a valid Chessmaster book or has no moves.
    pub fn open<P: AsRef<Path>>(file_name: P) -> Result<Self, Error> {
        let name = file_name.as_ref().display();
        let file = File::open(file_name.as_ref())?;
        // SAFETY: Mmap::map is unsafe because it involves file I/O which might lead to data races
        // if the underlying file is modified while the memory map is active. Here, it's safe
        // because we assume that the OBK files are not concurrently modified while they're
        // memory-mapped.
        let book = unsafe { Mmap::map(&file)? };

        // Read and validate the header bytes
        let version;
        let header = if book.starts_with(b"BOO!") {
            version = 2;
            &book[0..12]
        } else if book.starts_with(b"UGWS") {
            version = 1;
            &book[0..8]
        } else {
            return Err(Error::new(
                ErrorKind::InvalidData,
                tr!(
                    "File `{}' is not a Chessmaster book, invalid header `{}'",
                    name,
                    String::from_utf8_lossy(&book[0..4])
                ),
            ));
        };

        let move_count = u32::from_le_bytes(header[4..8].try_into().unwrap());
        if move_count == 0 {
            return Err(Error::new(
                ErrorKind::InvalidData,
                tr!("File `{}' has no moves", name),
            ));
        }

        let text_notes_count = if version == 2 {
            Some(u32::from_le_bytes(header[8..12].try_into().unwrap()))
        } else {
            None /* Version 1 has no notes */
        };

        let moves = Vec::new();
        let notes = Vec::new();

        Ok(Self {
            book,
            moves,
            notes,
            version,
            move_count,
            text_notes_count,
        })
    }

    /// Loads a OBK opening book file into memory.
    ///
    /// # Arguments
    ///
    /// * `progress_bar` - An optional progress bar to report progress.
    pub fn load(&mut self, progress_bar: Option<&ProgressBar>) {
        if let Some(pb) = progress_bar {
            pb.println(tr!(
                "Parsing and loading OBK opening book entries into memory."
            ));
            pb.set_message(tr!("Loading:"));
            pb.set_length(u64::from(self.move_count));
            pb.set_position(0);
        }

        let mut pos: Chess = Chess::default();
        let mut position_stack: Vec<Chess> = Vec::new();
        let mut unique_key_uci_pairs: HashSet<(u64, Uci)> = HashSet::new();

        let mut book_index: usize = if self.version == 2 {
            12
        } else {
            8 /* no notes, -4 bytes */
        };
        for _ in 0..self.move_count {
            let key = zobrist_hash(&pos);
            let move_bytes = [self.book[book_index], self.book[book_index + 1]];
            let obk_move = ObkMoveEntry::from_bytes(move_bytes, key);

            let uci = uci_from_obk_move_and_position(obk_move, &pos);
            match uci {
                Uci::Normal {
                    from,
                    to,
                    promotion: _,
                } => {
                    if from == Square::A1 && to == Square::A1 {
                        book_index += 2;
                        continue;
                    }
                }
                _ => unreachable!("{}", tr!("Unexpected null move or put move in book?")),
            };
            let mov = match uci.to_move(&pos) {
                Ok(mov) => mov,
                Err(err) => {
                    let epd = format!(
                        "{}",
                        Epd::from_position(pos.clone(), EnPassantMode::PseudoLegal)
                    );
                    eprintln!(
                        "{}",
                        tr!("Error making move uci:{} in epd:{}: {}", uci, epd, err)
                    );
                    book_index += 2;
                    continue;
                }
            };

            let prev_pos = pos.clone();
            pos.play_unchecked(&mov);

            /* Avoid duplicate UCI moves */
            let key_uci_pair = (key, uci.clone());
            if unique_key_uci_pairs.insert(key_uci_pair) {
                self.moves.push(obk_move);
            }

            if !obk_move.is_last_at_level {
                position_stack.push(prev_pos);
            }

            if obk_move.is_last_in_variation {
                if let Some(last_pos) = position_stack.pop() {
                    pos = last_pos;
                }
            }

            book_index += 2;
            if let Some(pb) = progress_bar {
                pb.inc(1);
            }
        }

        /*
        // TODO: Read the notes
        let mut note_index: usize = 0;
        while note_index < self.text_notes_count as usize {
            let note_type = self.book[book_index];
            let note_length = self.book[book_index + 1];
            let note_text = String::from_utf8_lossy(
                &self.book[book_index + 2..book_index + 2 + note_length as usize],
            )
            .to_string();

            let move_number = u32::from_le_bytes([
                self.book[book_index + 2 + note_length as usize],
                self.book[book_index + 3 + note_length as usize],
                0,
                0,
            ]);

            self.notes.push(ObkNoteEntry {
                move_number,
                note_length,
                note_type,
                note_text,
            });

            book_index += 4 + note_length as usize;
            note_index += 4 + note_length as usize;
        }
        */

        if let Some(pb) = progress_bar {
            pb.println(tr!(
                "Success loading {} OBK opening book entries into memory.",
                self.move_count
            ));
            pb.finish_with_message(tr!("Loading done."));
        }
    }

    /// Traverse the tree of possible moves from the given position, building a hashmap of book
    /// entries.
    ///
    /// The function takes a `Chess` position as input, and an optional `ProgressBar` and returns a
    /// `BookEntryHashMap`, which maps Zobrist hashes to their corresponding book entries.
    pub fn traverse_tree(
        &self,
        pos: Chess,
        progress_bar: Option<&ProgressBar>,
    ) -> BookEntryHashMap {
        let mut book: BookEntryHashMap = HashMap::with_hasher(ZobristHasherBuilder);
        let mut queue = VecDeque::new();
        let mut visited: ZobristHashSet = HashSet::with_hasher(ZobristHasherBuilder);
        let initial_zobrist = zobrist_hash(&pos);

        queue.push_back((pos, initial_zobrist));

        while let Some((pos, key)) = queue.pop_front() {
            if !visited.insert(key) {
                continue;
            }

            let entries: Vec<&ObkMoveEntry> =
                self.moves.iter().filter(|entry| entry.key == key).collect();
            if entries.is_empty() {
                continue;
            }
            let mut book_entries = Vec::new();

            for entry in entries {
                let uci = uci_from_obk_move_and_position(*entry, &pos);
                let chess_move = match uci.to_move(&pos) {
                    Ok(move_) => move_,
                    Err(_) => {
                        /* illegal uci, mmph */
                        continue;
                    }
                };
                let mov = from_uci(uci, is_king_on_start_square(&pos));
                let weight = u16::from(entry.weight);

                let book_entry = BookEntry {
                    key,
                    mov,
                    weight,
                    learn: 0,
                };

                book_entries.push(book_entry);

                let mut pos = pos.clone();
                pos.play_unchecked(&chess_move);
                let new_zobrist = zobrist_hash(&pos);
                if !book.contains_key(&new_zobrist) {
                    queue.push_back((pos, new_zobrist));
                }
            }

            let len = book_entries.len();
            book.insert(key, book_entries);

            if let Some(pb) = progress_bar {
                pb.inc(len as u64);
            }
        }

        book
    }

    /// Writes all possible games contained in a Chessmater book to a PGN file.
    ///
    /// This function traverses the Chessmaster book, which is a type of opening book, and writes
    /// all possible games to the output file in PGN format. A game is considered "possible" if it
    /// follows a path of moves in the book from the given starting position to a position with no
    /// more book moves. Each game is written as a separate round, and the rounds are numbered
    /// consecutively starting from 1.
    ///
    /// The `output` argument is a mutable reference to a `Write` trait object where the generated PGN will be written.
    /// The `event`, `site`, `date`, `white`, `black`, and `result` arguments are used to fill in the corresponding PGN tags for each game.
    /// The `max_ply` argument determines the limit of variation depth in plies.
    /// The `progress_bar` is an optional reference to a progress bar to report progress.
    ///
    /// # Errors
    ///
    /// This function will panic if writing to the output file fails.
    ///
    /// # Panics
    ///
    /// Panics if the disk is full or the file isn't writable.
    #[allow(clippy::too_many_arguments)]
    pub fn write_pgn(
        &self,
        output: &mut dyn std::io::Write,
        position: &Chess,
        event: &str,
        site: &str,
        date: &str,
        white: &str,
        black: &str,
        result: &str,
        max_ply: usize,
        progress_bar: Option<&ProgressBar>,
    ) {
        let fen_header: String;
        let fen = if *position == Chess::default() {
            None
        } else {
            fen_header = Fen::from_position(position.clone(), EnPassantMode::Legal).to_string();
            Some(&fen_header)
        };

        // Prepare move lookup table for faster lookups.
        // TODO: Do we want to store them by default in a HashMap, rather than Vec?
        let mut move_map: ObkMoveEntryHashMap = HashMap::with_hasher(ZobristHasherBuilder);
        for entry in &self.moves {
            move_map.entry(entry.key).or_default().push(*entry);
        }

        if let Some(progress_bar) = progress_bar {
            progress_bar.set_message(tr!("Writing:"));
            progress_bar.set_length(0);
            progress_bar.set_position(0);
        }
        Self::_write_pgn(
            output,
            position,
            &HashSet::with_hasher(ZobristHasherBuilder),
            &move_map,
            &mut Vec::new(),
            fen,
            &mut 1,
            event,
            site,
            date,
            white,
            black,
            result,
            max_ply,
            position.turn(),
            progress_bar,
        );
        if let Some(progress_bar) = progress_bar {
            progress_bar.set_message(tr!("Writing done."));
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn _write_pgn(
        output: &mut dyn std::io::Write,
        position: &Chess,
        position_set: &ZobristHashSet,
        move_map: &ObkMoveEntryHashMap,
        move_history: &mut Vec<San>,
        fen: Option<&String>,
        round: &mut usize,
        event: &str,
        site: &str,
        date: &str,
        white: &str,
        black: &str,
        result: &str,
        max_ply: usize,
        initial_color: Color,
        progress_bar: Option<&ProgressBar>,
    ) {
        // Return if the maximum ply is reached
        if move_history.len() >= max_ply {
            return;
        }

        // Each recursive call gets a localized copy of visited positions, preventing global skips.
        // TODO: This is a relatively memory-intensive operation but does the right thing.
        let mut position_set = position_set.clone();

        if let Some(entries) = move_map.get(&zobrist_hash(position)) {
            let mut entries = entries.clone();
            entries.sort_unstable_by_key(|mov| std::cmp::Reverse(mov.weight));

            for entry in entries {
                let mov = match Uci::from(entry).to_move(position) {
                    Ok(mov) => mov,
                    Err(_) => continue, // TODO: warn about illegal move?
                };
                let san = San::from_move(position, &mov);
                move_history.push(san);
                let mut new_position = position.clone();
                new_position.play_unchecked(&mov);

                // If the new position has been seen before, skip it to avoid infinite recursion.
                let hash = zobrist_hash(&new_position);
                if !position_set.insert(hash) {
                    // Insert returned false, the set already contained this value.
                    move_history.pop();
                    continue;
                }

                // Recursively generate all games starting from the new position.
                Self::_write_pgn(
                    output,
                    &new_position,
                    &position_set,
                    move_map,
                    move_history,
                    fen,
                    round,
                    event,
                    site,
                    date,
                    white,
                    black,
                    result,
                    max_ply,
                    initial_color,
                    progress_bar,
                );

                // Undo the move and remove it from the move history.
                move_history.pop();
            }
        } else {
            // This is a leaf node.
            if !move_history.is_empty() {
                let opening = move_history
                    .iter()
                    .enumerate()
                    .map(|(i, san)| {
                        let move_number = i / 2 + 1;
                        let move_text = san.to_string();
                        match (initial_color, i, i % 2) {
                            (Color::White, _, 0) => format!("{}. {} ", move_number, move_text),
                            (Color::Black, 0, 0) => format!("{}... {} ", move_number, move_text),
                            (Color::Black, _, 1) => format!("{}. {} ", move_number + 1, move_text),
                            _ => format!("{} ", move_text),
                        }
                    })
                    .collect::<String>();

                let fen_header = if let Some(fen) = fen {
                    format!("[FEN \"{}\"]\n[Setup \"1\"]\n", fen)
                } else {
                    String::new()
                };

                writeln!(
                    output,
                    "[Event \"{}\"]\n\
                    [Site \"{}\"]\n\
                    [Date \"{}\"]\n\
                    [Round \"{}\"]\n\
                    [White \"{}\"]\n\
                    [Black \"{}\"]\n\
                    [Result \"{}\"]\n{}\
                    [Annotator \"{} v{}\"]",
                    event,
                    site,
                    date,
                    round,
                    white,
                    black,
                    result,
                    fen_header,
                    crate::built_info::PKG_NAME,
                    crate::built_info::PKG_VERSION,
                )
                .expect("write output PGN");

                writeln!(output, "\n{} {}\n", opening.trim(), result).expect("write output PGN");
                *round += 1;
                if let Some(progress_bar) = progress_bar {
                    progress_bar.inc(1);
                }
            }
        }
    }

    /// Generate a tree of possible moves from the given position up to the specified maximum ply.
    ///
    /// The function takes a `shakmaty::Chess` position and a maximum ply as input, and returns a
    /// `Tree` containing `String`s representing moves.
    pub fn tree(&self, position: &Chess, max_ply: u16) -> Tree<String> {
        fn build_tree(
            moves: &Vec<ObkMoveEntry>,
            position: &Chess,
            parent: &mut Tree<String>,
            ply: u16,
            max_ply: u16,
            visited_keys: &ZobristHashSet,
        ) {
            if ply >= max_ply {
                return;
            }

            let key = zobrist_hash(position);
            let mut book_entries: Vec<&ObkMoveEntry> =
                moves.iter().filter(|entry| entry.key == key).collect();
            if book_entries.is_empty() {
                return;
            }
            book_entries.sort_unstable_by_key(|mov| std::cmp::Reverse(mov.weight));

            for entry in book_entries {
                let m = match Uci::from(*entry).to_move(position) {
                    Ok(m) => m,
                    Err(_) => continue,
                };
                let mut new_position = position.clone();
                new_position.play_unchecked(&m);

                let key = zobrist_hash(&new_position);
                if visited_keys.contains(&key) {
                    continue;
                }

                let mut new_visited_keys = visited_keys.clone(); // Clone visited_keys
                new_visited_keys.insert(key);

                let mut new_tree = Tree::new(San::from_move(position, &m).to_string());
                build_tree(
                    moves,
                    &new_position,
                    &mut new_tree,
                    ply + 1,
                    max_ply,
                    &new_visited_keys,
                );

                parent.push(new_tree);
            }
        }

        let epd = format!(
            "{}",
            Epd::from_position(position.clone(), EnPassantMode::PseudoLegal)
        );
        let mut root_tree = Tree::new(epd);

        let key = zobrist_hash(position);
        let mut visited_keys: ZobristHashSet = HashSet::with_hasher(ZobristHasherBuilder);
        visited_keys.insert(key);

        build_tree(
            &self.moves,
            position,
            &mut root_tree,
            0,
            max_ply,
            &visited_keys,
        );

        root_tree
    }
}