//
// jja: swiss army knife for chess file formats
// src/polyglotbook.rs: PolyGlot book file interface
//
// Copyright (c) 2023 Ali Polatel <alip@chesswob.org>
// Based in part upon PolyGlot 1.4 which is:
//     Copyright 2004-2006 Fabien Letouzey.
//     Licensed under the GNU General Public License, version 2.
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::collections::HashSet;
use std::fs::File;
use std::io::Error;
use std::path::Path;

use memmap::Mmap;

use indicatif::ProgressBar;

use shakmaty::fen::{Epd, Fen};
use shakmaty::san::San;
use shakmaty::{Chess, Color, EnPassantMode, Position};

use crate::hash::{zobrist_hash, ZobristHashSet, ZobristHasherBuilder};
use crate::polyglot::*;
use crate::tr;

use termtree::Tree;

/// `PolyGlotBook` is a struct that represents a Polyglot opening book.
pub struct PolyGlotBook {
    /// The number of entries in the Polyglot opening book.
    pub num_entries: usize,
    /// The `Mmap` handle to the Polyglot opening book file.
    /// This is `Some` if num_entries > 0, None otherwise.
    book: Option<Mmap>,
}

/// `PolyGlotIter` is an iterator to iterate through the entries.
pub struct PolyGlotBookIter<'a> {
    book: &'a PolyGlotBook,
    index: usize,
}

impl<'a> Iterator for PolyGlotBookIter<'a> {
    type Item = BookEntry;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index < self.book.num_entries {
            let entry = self.book.get(self.index);
            self.index += 1;
            entry
        } else {
            None
        }
    }
}

/// `PolyGlotBookGroupIter` is an iterator to iterate through the entries grouped by key.
pub struct PolyGlotBookGroupIter<'a> {
    book: &'a PolyGlotBook,
    index: usize,
    pending_entry: Option<Result<BookEntry, std::io::Error>>,
}

impl<'a> Iterator for PolyGlotBookGroupIter<'a> {
    type Item = Result<Vec<BookEntry>, std::io::Error>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.book.num_entries {
            return None;
        }

        let mut group_entries: Vec<BookEntry> = Vec::new();
        while self.index < self.book.num_entries {
            let entry_opt = match self.pending_entry.take() {
                Some(Ok(entry)) => Some(entry),
                Some(Err(e)) => return Some(Err(e)),
                None => self.book.get(self.index),
            };

            if let Some(entry) = entry_opt {
                if let Some(last_entry) = group_entries.last() {
                    if last_entry.key != entry.key {
                        self.pending_entry = Some(Ok(entry));
                        break;
                    }
                }
                group_entries.push(entry);
            }
            self.index += 1;
        }

        if group_entries.is_empty() {
            None
        } else {
            Some(Ok(group_entries))
        }
    }
}

impl<'a> IntoIterator for &'a mut PolyGlotBook {
    type Item = BookEntry;
    type IntoIter = PolyGlotBookIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        PolyGlotBookIter {
            book: self,
            index: 0,
        }
    }
}

impl PolyGlotBook {
    /// Opens a Polyglot opening book file.
    ///
    /// # Arguments
    ///
    /// * `file_name` - The path to the opening book file.
    ///
    /// # Returns
    ///
    /// * `Result<Self, Error>` - Returns a `Result` containing a `PolyGlotBook` instance if
    /// successful, or an error if there was an issue opening the file.
    pub fn open<P: AsRef<Path>>(file_name: P) -> Result<Self, Error> {
        let file = File::open(file_name)?;
        let num_bytes = file.metadata()?.len();
        let num_entries = (num_bytes / BOOK_ENTRY_SIZE as u64) as usize;

        let book = if num_bytes == 0 {
            None
        } else {
            // 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 BIN files are not concurrently modified while they're
            // memory-mapped.
            Some(unsafe { Mmap::map(&file)? })
        };

        Ok(Self { book, num_entries })
    }

    /// Returns an iterator which iterates over entries grouped by key.
    pub fn into_iter_grouped(&self) -> PolyGlotBookGroupIter {
        PolyGlotBookGroupIter {
            book: self,
            index: 0,
            pending_entry: None,
        }
    }

    /// Looks up the moves available in the opening book for the given chess position.
    ///
    /// # Arguments
    ///
    /// * `key: u64` - Zobrist hash of the chess position to look up.
    ///
    /// # Returns
    ///
    /// * `Option<Vec<BookEntry>>` - Returns an `Option` with a vector of `BookEntry` instances
    /// representing the possible moves if found, or `None` if no moves or book are found.
    pub fn lookup_moves(&self, key: u64) -> Option<Vec<BookEntry>> {
        // Early return if self.book is None
        self.book.as_ref()?;

        let offset = match self.find(key) {
            None => return None,
            Some(n) => n,
        };

        let mut book = Vec::new();
        let mut index = 0;

        // Read all book entries with the correct key. They're all stored
        // contiguously, so just scan through as long as the key matches.
        while offset + index < self.num_entries {
            let entry = self.get(offset + index)?;

            if entry.key != key {
                break;
            }
            book.push(entry);
            index += 1;
        }

        if book.is_empty() {
            None
        } else {
            Some(book)
        }
    }

    /// Writes all possible games contained in a Polyglot book to a PGN file.
    ///
    /// This function traverses the Polyglot 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 `look_ahead` argument determines how many number of plies to look ahead on book lookup
    /// misses. This is useful to create PGN out of books created with `--only-black` or
    /// `only-white` (currently only 0 and 1 are implemented, panics on other values).
    /// 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,
        look_ahead: usize,
        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)
        };

        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),
            &mut Vec::new(),
            fen,
            &mut 1,
            event,
            site,
            date,
            white,
            black,
            result,
            look_ahead,
            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(
        &self,
        output: &mut dyn std::io::Write,
        position: &Chess,
        position_set: &ZobristHashSet,
        move_history: &mut Vec<San>,
        fen: Option<&String>,
        round: &mut usize,
        event: &str,
        site: &str,
        date: &str,
        white: &str,
        black: &str,
        result: &str,
        look_ahead: usize,
        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();

        let hash = zobrist_hash(position);
        let entries = match look_ahead {
            0 => self.lookup_moves(hash),
            1 => {
                let mut entries = self.lookup_moves(hash).unwrap_or_default();
                for move_ in position.legal_moves() {
                    let pgmove = from_move(&move_);
                    if entries.iter().any(|entry| entry.mov == pgmove) {
                        continue; // Move already in the book!
                    }

                    let mut p = position.clone();
                    p.play_unchecked(&move_);
                    if self.lookup_moves(zobrist_hash(&p)).is_none() {
                        continue; // Checked next ply without success.
                    }

                    entries.push(BookEntry {
                        key: hash,
                        mov: pgmove,
                        weight: 1,
                        learn: 0,
                    });
                }

                if entries.is_empty() {
                    None
                } else {
                    Some(entries)
                }
            }
            n => {
                unimplemented!(
                    "{}",
                    tr!(
                        "--look-ahead={} is not implemented (use 0..=1), please report a bug if you need it!",
                        n
                    )
                );
            }
        };

        if let Some(mut entries) = entries {
            // Sort the moves by their weight in reverse order.
            entries.sort_unstable_by_key(|entry| std::cmp::Reverse(entry.weight));

            for entry in entries {
                let mov = match to_move(position, entry.mov) {
                    Some(mov) => mov,
                    None => 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_history,
                    fen,
                    round,
                    event,
                    site,
                    date,
                    white,
                    black,
                    result,
                    look_ahead,
                    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);
                }
            }
        }
    }

    /// A method that generates a tree of moves from a given position using an opening book.
    pub fn tree(&self, position: &Chess, max_ply: u16) -> Tree<String> {
        fn build_tree(
            book: &PolyGlotBook,
            position: &Chess,
            parent: &mut Tree<String>,
            ply: u16,
            max_ply: u16,
            visited_keys: &ZobristHashSet,
        ) {
            if ply >= max_ply {
                return;
            }

            let moves = book.lookup_moves(zobrist_hash(position));
            if let Some(mut book_entries) = moves {
                book_entries.sort_unstable_by_key(|mov| std::cmp::Reverse(mov.weight));
                for entry in book_entries {
                    let mov = entry.mov;
                    let m = match to_move(position, mov) {
                        Some(m) => m,
                        None => 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(
                        book,
                        &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, position, &mut root_tree, 0, max_ply, &visited_keys);

        root_tree
    }

    /// Helper function to find the index of the first book entry matching a target key.
    ///
    /// # Arguments
    ///
    /// * `target_key: u64` - The target key to find in the book.
    ///
    /// # Returns
    ///
    /// * `Option<usize>` - Returns an `Option` with the index of the first book entry with a
    /// matching key if found, or `None` if no match is found.
    pub fn find(&self, target_key: u64) -> Option<usize> {
        let mut low = 0;
        let mut high = self.num_entries;

        // Since the positions are all in sorted order, we use a binary search to find the target key.
        while low < high {
            let mid = low + (high - low) / 2;
            let entry_key = self.get_key(mid)?;

            if target_key <= entry_key {
                high = mid;
            } else {
                low = mid + 1;
            }
        }

        // If the key at the low index matches the target key, we've found a match.
        if low < self.num_entries && self.get_key(low)? == target_key {
            Some(low)
        } else {
            None
        }
    }

    /// Helper function to get a book entry from the opening book file.
    ///
    /// # Arguments
    ///
    /// * `index: usize` - The index of the book entry to get.
    ///
    /// # Returns
    ///
    /// * `Option<BookEntry>` - Returns an `Option<BookEntry>` instance. If the `index`
    /// is not smaller than the total number of entries in the book, or if the book
    /// is of zero-size, this function will return `None`.
    pub fn get(&self, index: usize) -> Option<BookEntry> {
        let book = self.book.as_ref()?;
        if index >= self.num_entries {
            return None;
        }

        let offset = index * BOOK_ENTRY_SIZE;
        let buf = &book[offset..offset + BOOK_ENTRY_SIZE];

        let key = u64::from_be_bytes([
            buf[0], buf[1], buf[2], buf[3], buf[4], buf[5], buf[6], buf[7],
        ]);
        let mov = u16::from_be_bytes([buf[8], buf[9]]);
        let weight = u16::from_be_bytes([buf[10], buf[11]]);
        let learn = u32::from_be_bytes([buf[12], buf[13], buf[14], buf[15]]);

        Some(BookEntry {
            key,
            mov,
            weight,
            learn,
        })
    }

    /// Helper function to read a book entry key from the opening book file.
    ///
    /// # Arguments
    ///
    /// * `index: usize` - The index of the book entry to read.
    ///
    /// # Returns
    ///
    /// * `Option<u64>` - Returns an `Option` with an `u64` which represents the book key.
    /// If the book is not loaded or the index is out of bounds, returns `None`.
    pub fn get_key(&self, index: usize) -> Option<u64> {
        if index >= self.num_entries {
            return None;
        }

        if let Some(book) = &self.book {
            let offset = index * BOOK_ENTRY_SIZE;
            let buf: [u8; 8] = book[offset..offset + 8].try_into().expect("slice book");
            Some(u64::from_be_bytes(buf))
        } else {
            None
        }
    }
}