use crate::{
    agent::Agent,
    board::{Board, Winner},
    piece::Piece,
};
use indicatif::{ProgressIterator, ProgressStyle};

#[derive(Clone, Debug)]
struct Move {
    /// `i` position of move
    i: usize,

    /// `j` position of move
    j: usize,

    /// [`Board`] state after move is made
    board: Board,

    /// Current winner of the match
    winner: Winner,

    /// Index of this move's parent
    parent: Option<usize>,

    /// Indices of this Move's Children
    children: Vec<usize>,

    /// Value of this move
    value: i64,

    color: Piece,
}

impl Move {
    pub const fn coords(&self) -> (usize, usize) {
        (self.i, self.j)
    }

    fn compute_self_value(&self, agent_color: Piece) -> i64 {
        let mut self_value = self.board.net_score(agent_color) as i64;

        if self.winner == Winner::Player(!agent_color) {
            // if this board results in the opponent winning, MAJORLY negatively weigh this move
            // NOTE! this branch isn't completely deleted because if so, the bot wouldn't make a move.
            // We shouldn't prune branches because we still need to always react to the opponent's moves
            self_value = i64::MIN;
        } else if self.winner == Winner::Player(agent_color) {
            // results in a win for the agent
            self_value = i64::MAX;
        }
        // TODO! handle ties... what should they be valued as? maybe `i64::MAX / 2` or 0?

        self_value
    }
}

struct FutureMoves {
    /// Arena containing all [`Move`]
    arena: Vec<Move>,

    /// Index of the [`Move`] tree's root node
    current_root: Option<usize>,

    /// Current generated depth of the Arena
    current_depth: isize,

    /// Target depth of children to generate
    max_depth: usize,

    /// Color w.r.t
    agent_color: Piece,
}

impl FutureMoves {
    pub const fn new(agent_color: Piece, max_depth: usize) -> Self {
        Self {
            arena: Vec::new(),
            current_root: None,
            current_depth: 0,
            max_depth,
            agent_color,
        }
    }

    /// Generate children for all children of `nodes`
    fn extend_layers(&mut self) {
        let mut next_nodes: Vec<usize> = (0..self.arena.len())
            .filter(|&idx| self.arena[idx].children.is_empty())
            .filter(|&idx| self.is_connected_to_root(idx)) // put here so this will not extend needlessly before prunes
            .collect();

        for _ in self.current_depth..=(self.max_depth as isize) {
            // TODO! use `i` in order to prune along-the-way
            // i.e, every 4 moves of color `self.agent_color`, do a pruning step,
            // only keeping the top [`Move`]

            next_nodes = next_nodes
                .into_iter()
                .progress_with_style(
                    ProgressStyle::with_template("Generating children: ({pos}/{len}) {per_sec}")
                        .unwrap(),
                )
                .flat_map(|node_idx| {
                    self.generate_children(
                        node_idx,
                        &self.arena[node_idx].board.clone(),
                        !self.arena[node_idx].color,
                    )
                })
                .collect();
        }
        self.current_depth = self.max_depth as isize;
    }

    /// Determines if a [`Move`] at index `idx` is connected to `self.current_root`
    /// Returns `false` if `self.current_root` is None
    fn is_connected_to_root(&self, idx: usize) -> bool {
        if let Some(root) = self.current_root {
            let mut current = Some(idx);
            while let Some(parent_idx) = current {
                if parent_idx == root {
                    return true;
                }
                current = self.arena[parent_idx].parent;
            }
        }
        false
    }

    /// Creates children for a parent (`parent`), returns an iterator it's children's indexes
    fn generate_children(
        &mut self,
        parent_idx: usize,
        board: &Board,
        color: Piece,
    ) -> impl Iterator<Item = usize> {
        let mut new: Vec<Move> =
            // use [`Board::all_positions`] here instead of [`Board::possible_moves`]
            // because we use [`Board::what_if`] later and we want to reduce calls to [`Board::propegate_from_dry`]
            Board::all_positions()
                .flat_map(|(i, j)| board.what_if(i, j, !self.arena[parent_idx].color).map(|x| (i, j, x)))
                .map(|(i, j, new_board)| Move {
                    i,
                    j,
                    board: new_board,
                    winner: new_board.game_winner(color),
                    parent: Some(parent_idx),
                    children: Vec::new(),
                    value: 0,
                    color: !self.arena[parent_idx].color
                }).collect();

        // keep the TOP_K children `self.agent_color`-color moves
        const TOP_K_CHILDREN: usize = 1;

        // we want to keep only the best move of the agent
        if color == self.agent_color && new.len() > TOP_K_CHILDREN {
            // TODO! Move this to `extend_layers` so we can prune based on recursive [`Move`] value
            // negative, because we want the max value to be at the first index
            new.sort_by_key(|x| -x.compute_self_value(self.agent_color));
            new.drain(TOP_K_CHILDREN..);
        }

        let start_idx = self.arena.len();
        self.arena.extend(new);
        let new_indices = start_idx..self.arena.len();

        self.arena[parent_idx].children.extend(new_indices.clone());

        new_indices
    }

    /// Given an index from `self.arena`, what depth is it at? 1-indexed (ROOT IS AT INDEX 1)
    fn depth_of(&self, node_idx: usize) -> usize {
        let mut depth = 0;
        let mut current = Some(node_idx);
        while let Some(parent_idx) = current {
            depth += 1;
            current = self.arena[parent_idx].parent;
        }

        depth
    }

    fn compute_values(&mut self) {
        // PERF! doing this filtering here previously visited moves reduces time
        // spent in `depth_of` from 27.79% -> 2.9%
        let mut visited = vec![false; self.arena.len()];

        for depth in (0..=self.current_depth).rev() {
            for (idx, was_visited) in visited.iter_mut().enumerate() {
                if *was_visited {
                    continue;
                } else {
                    *was_visited = true;
                }

                if self.depth_of(idx) != depth as usize {
                    continue;
                }

                let self_value = self.arena[idx].compute_self_value(self.agent_color)
                    / (self.current_depth - depth + 1) as i64;

                let children_value = self.arena[idx]
                    .children
                    .iter()
                    .map(|&child| self.arena[child].value)
                    .sum::<i64>()
                    .checked_div(self.arena[idx].children.len() as i64)
                    .unwrap_or(0);

                self.arena[idx].value = self_value + children_value;
            }
        }
    }

    pub fn best_move(&self) -> Option<(usize, usize)> {
        self.current_root
            .and_then(|x| {
                self.arena[x]
                    .children
                    .iter()
                    .max_by_key(|&&idx| self.arena[idx].value)
            })
            .inspect(|&&x| {
                assert_eq!(
                    self.arena[x].color, self.agent_color,
                    "selected move color should be the same as the color of the agent"
                );
            })
            .map(|&x| self.arena[x].coords())
    }

    /// Updates `FutureMoves` based on the current state of the board
    /// The board is supposed to be after the opposing move
    pub fn update(&mut self, board: &Board) {
        let curr_board = self
            .arena
            .iter()
            .enumerate()
            .find(|(_, m)| {
                &m.board == board && (m.parent == self.current_root) && self.current_root.is_some()
            })
            .map(|(idx, _)| idx);

        if let Some(curr_board_idx) = curr_board {
            self.update_root_idx(curr_board_idx);
        } else {
            println!("Generating root of FutureMoves");
            self.arena.clear();
            self.arena.push(Move {
                i: 0,
                j: 0,
                board: *board,
                winner: Winner::None,
                parent: None,
                children: Vec::new(),
                value: 0,
                color: !self.agent_color,
            });
            self.update_root_idx(0);
        }
    }

    pub fn update_root_coord(&mut self, i: usize, j: usize) -> bool {
        self.arena
            .iter()
            .enumerate()
            .find_map(|(idx, node)| {
                (node.parent == self.current_root && node.i == i && node.j == j).then_some(idx)
            })
            .inspect(|&root| self.update_root_idx(root))
            .is_some()
    }

    fn update_root_idx(&mut self, idx: usize) {
        self.current_root = Some(idx);
        self.current_depth -= self.depth_of(idx) as isize - 1;
        self.prune_unrelated();
        self.extend_layers();
        self.compute_values();
    }

    fn prune_unrelated(&mut self) {
        let Some(root) = self.current_root else {
            return;
        };

        // make sure `root` doesn't reference another node
        self.arena[root].parent = None;

        let mut retain = vec![false; self.arena.len()];

        // stack is going to be AT MAXIMUM, the size of the array,
        // so lets just pre-allocate it
        let mut stack: Vec<usize> = Vec::with_capacity(self.arena.len());
        stack.push(root);

        // traverse children of the current root
        while let Some(idx) = stack.pop() {
            retain[idx] = true;
            stack.extend(self.arena[idx].children.iter());
        }

        let mut index_map = vec![None; self.arena.len()];

        self.arena = retain
            .into_iter()
            .enumerate() // old_idx
            .zip(self.arena.drain(..))
            .flat_map(|((old_idx, keep), node)| keep.then_some((old_idx, node))) // filter out unrelated nodes
            .enumerate() // new_idx
            .map(|(new_idx, (old_idx, mut node))| {
                index_map[old_idx] = Some(new_idx);

                if let Some(parent) = node.parent.as_mut() {
                    if let Some(new_parent) = index_map[*parent] {
                        *parent = new_parent;
                    }
                }

                node.children.retain_mut(|c| {
                    if let Some(new_c) = index_map[*c] {
                        *c = new_c;
                        true
                    } else {
                        false
                    }
                });

                node
            })
            .collect();

        self.current_root = index_map[root];
    }
}

pub struct ComplexAgent {
    color: Piece,
    future_moves: FutureMoves,
}

impl ComplexAgent {
    pub const fn new(color: Piece) -> Self {
        const MAX_DEPTH: usize = 15;
        Self {
            color,
            future_moves: FutureMoves::new(color, MAX_DEPTH),
        }
    }
}

impl Agent for ComplexAgent {
    fn next_move(&mut self, board: &Board) -> Option<(usize, usize)> {
        self.future_moves.update(board);

        println!("# of moves stored: {}", self.future_moves.arena.len());

        self.future_moves.best_move().inspect(|&(i, j)| {
            self.future_moves.update_root_coord(i, j);
        })
    }

    fn name(&self) -> &'static str {
        "Complex Agent"
    }

    fn color(&self) -> Piece {
        self.color
    }
}