use crate::position::Position;
use crate::types::{Move, Color, PieceType};
use crate::movegen::generate_moves;
use crate::evaluate::evaluate;
use crate::tt::{TranspositionTable, TT_EXACT, TT_ALPHA, TT_BETA};
use crate::zobrist::Zobrist;
use crate::timeman::TimeManager;
use crate::movepick::{score_moves, pick_next_move};

pub struct Searcher {
    pub tt: TranspositionTable,
    pub zobrist: Zobrist,
    pub killers: [[Move; 64]; 2],
    pub nodes: u64,
}

impl Searcher {
    pub fn new() -> Self {
        Self {
            tt: TranspositionTable::new(16),
            zobrist: Zobrist::new(),
            killers: [[Move::NONE; 64]; 2],
            nodes: 0,
        }
    }

    pub fn search_best(&mut self, pos: &mut Position, tm: &mut TimeManager, max_depth: i32) -> Move {
        let mut best_move = Move::NONE;
        self.nodes = 0;

        for depth in 1..=max_depth {
            let mut current_best = Move::NONE;
            let score = self.alpha_beta(pos, tm, -30000, 30000, depth, 0, &mut current_best, false);
            
            if !tm.tick() { break; } 
            
            best_move = current_best;
            println!("info depth {} score cp {} nodes {}", depth, score, self.nodes);
        }
        best_move
    }

    fn alpha_beta(&mut self, pos: &mut Position, tm: &mut TimeManager, mut alpha: i32, beta: i32, depth: i32, ply: usize, best_move: &mut Move, is_null: bool) -> i32 {
        if !tm.tick() { return 0; }
        self.nodes += 1;

        if depth <= 0 { return self.quiescence(pos, alpha, beta); }

        let hash = self.zobrist.hash(pos);
        let mut tt_move = Move::NONE;
        
        if let Some(entry) = self.tt.probe(hash) {
            tt_move = entry.best_move;
            if entry.depth >= depth as u8 && ply > 0 {
                if entry.flag == TT_EXACT { return entry.score; }
                if entry.flag == TT_ALPHA && entry.score <= alpha { return alpha; }
                if entry.flag == TT_BETA && entry.score >= beta { return beta; }
            }
        }

        let us = pos.side_to_move;
        let mut king_bb = pos.pieces[crate::types::PieceType::King as usize] & pos.colors[us as usize];
        let in_check = if king_bb != 0 {
            let king_sq = crate::bitboard::pop_lsb(&mut king_bb);
            pos.is_square_attacked(crate::types::Square::from_u8(king_sq), us)
        } else {
            false
        };

        // Null Move Pruning (NMP)
        // If we are not in check, not already a null move, and have depth to spare
        if !in_check && depth >= 3 && !is_null && ply > 0 {
            let eval = evaluate(pos);
            // If the position is so good that even doing NOTHING maintains an advantage
            if eval >= beta {
                pos.side_to_move = if pos.side_to_move == Color::White { Color::Black } else { Color::White };
                let null_score = -self.alpha_beta(pos, tm, -beta, -beta + 1, depth - 3, ply + 1, &mut Move::NONE, true);
                pos.side_to_move = if pos.side_to_move == Color::White { Color::Black } else { Color::White }; 
                
                if null_score >= beta {
                    return beta; // Cut off the branch entirely
                }
            }
        }

        let mut moves = Vec::new();
        generate_moves(pos, &mut moves);
        
        let mut move_list: Vec<(Move, i32)> = moves.into_iter().map(|m| (m, 0)).collect();
        score_moves(pos, &mut move_list, tt_move, self.killers[0][ply], self.killers[1][ply]);

        let mut legal_moves = 0;
        let mut max_score = -30000;
        let mut local_best = Move::NONE;
        let original_alpha = alpha;
        let mut b_search_pv = true;

        for i in 0..move_list.len() {
            let m = pick_next_move(&mut move_list, i);
            if !pos.do_move(m) { continue; }
            legal_moves += 1;
            
            let is_capture = pos.get_piece_at(m.to()) != crate::types::PieceType::None;
            let is_promo = m.promotion() != crate::types::PieceType::None;
            
            let mut score;
            
            // Late Move Reductions (LMR)
            // If a move is late in the list and doesn't look critical, search it less
            let mut reduction = 0;
            if depth >= 3 && legal_moves > 4 && !in_check && !is_capture && !is_promo {
                reduction = 1;
                if legal_moves > 10 { reduction = 2; }
            }

            if b_search_pv {
                score = -self.alpha_beta(pos, tm, -beta, -alpha, depth - 1, ply + 1, &mut Move::NONE, false);
            } else {
                score = -self.alpha_beta(pos, tm, -alpha - 1, -alpha, depth - 1 - reduction, ply + 1, &mut Move::NONE, false);
                if score > alpha && reduction > 0 {
                    // Re-search without reduction if it actually turned out to be a good move
                    score = -self.alpha_beta(pos, tm, -alpha - 1, -alpha, depth - 1, ply + 1, &mut Move::NONE, false);
                }
                if score > alpha && score < beta {
                    // Full window search if it beat alpha
                    score = -self.alpha_beta(pos, tm, -beta, -alpha, depth - 1, ply + 1, &mut Move::NONE, false);
                }
            }
            
            pos.undo_move(m);

            if score > max_score {
                max_score = score;
                local_best = m;
            }
            if score > alpha { alpha = score; b_search_pv = false; }
            if alpha >= beta {
                if !is_capture {
                    self.killers[1][ply] = self.killers[0][ply];
                    self.killers[0][ply] = m;
                }
                break;
            }
        }

        if legal_moves == 0 {
            if in_check { return -30000 + ply as i32; } // Checkmate
            return 0; // Stalemate
        }

        let flag = if max_score <= original_alpha { TT_ALPHA } else if max_score >= beta { TT_BETA } else { TT_EXACT };
        self.tt.store(hash, depth as u8, flag, max_score, local_best);
        
        *best_move = local_best;
        max_score
    }

    fn quiescence(&mut self, pos: &mut Position, mut alpha: i32, beta: i32) -> i32 {
        self.nodes += 1;
        let stand_pat = evaluate(pos);
        
        if stand_pat >= beta { return beta; }
        if alpha < stand_pat { alpha = stand_pat; }

        let mut moves = Vec::new();
        generate_moves(pos, &mut moves);
        
        let mut move_list: Vec<(Move, i32)> = moves.into_iter()
            .filter(|m| pos.get_piece_at(m.to()) != crate::types::PieceType::None || m.promotion() != crate::types::PieceType::None)
            .map(|m| (m, 0)).collect();
            
        score_moves(pos, &mut move_list, Move::NONE, Move::NONE, Move::NONE);

        for i in 0..move_list.len() {
            let m = pick_next_move(&mut move_list, i);
            if !pos.do_move(m) { continue; }
            
            let score = -self.quiescence(pos, -beta, -alpha);
            pos.undo_move(m);

            if score >= beta { return beta; }
            if score > alpha { alpha = score; }
        }
        alpha
    }
}