package de.nowchess.bot

import de.nowchess.api.board.PieceType
import de.nowchess.api.game.GameContext
import de.nowchess.api.move.{Move, MoveType}
import de.nowchess.rules.RuleSet
import de.nowchess.rules.sets.DefaultRules

final class AlphaBetaSearch(
  rules: RuleSet = DefaultRules,
  tt: TranspositionTable = TranspositionTable()
):

  private val INF = Int.MaxValue / 2
  private val MAX_QUIESCENCE_PLY = 64
  private val NULL_MOVE_R = 2
  private val ASPIRATION_DELTA = 50
  private val TIME_CHECK_FREQUENCY = 1000

  @volatile private var timeStartMs = 0L
  @volatile private var timeLimitMs = 0L
  @volatile private var nodeCount = 0

  /** Return the best move for the side to move, searching to maxDepth plies.
   *  Uses iterative deepening with aspiration windows. */
  def bestMove(context: GameContext, maxDepth: Int): Option[Move] =
    tt.clear()
    var bestSoFar: Option[Move] = None
    var prevScore = 0
    for depth <- 1 to maxDepth do
      val (alpha, beta) =
        if depth == 1 then (-INF, INF)
        else (prevScore - ASPIRATION_DELTA, prevScore + ASPIRATION_DELTA)
      val (score, move) = searchWithAspiration(context, depth, alpha, beta)
      prevScore = score
      move.foreach(m => bestSoFar = Some(m))
    bestSoFar

  /** Return the best move for the side to move within a time budget (ms).
   *  Uses iterative deepening, stopping when time runs out. */
  def bestMoveWithTime(context: GameContext, timeBudgetMs: Long): Option[Move] =
    tt.clear()
    timeStartMs = System.currentTimeMillis()
    timeLimitMs = timeBudgetMs
    nodeCount = 0
    var bestSoFar: Option[Move] = None
    var prevScore = 0
    var depth = 1

    while !isOutOfTime() do
      val (alpha, beta) =
        if depth == 1 then (-INF, INF)
        else (prevScore - ASPIRATION_DELTA, prevScore + ASPIRATION_DELTA)
      val (score, move) = searchWithAspiration(context, depth, alpha, beta)
      prevScore = score
      move.foreach(m => bestSoFar = Some(m))
      depth += 1
    bestSoFar

  private def isOutOfTime(): Boolean =
    System.currentTimeMillis() - timeStartMs >= timeLimitMs

  private def searchWithAspiration(
    context: GameContext,
    depth: Int,
    alpha: Int,
    beta: Int
  ): (Int, Option[Move]) =
    val (score, move) = search(context, depth, 0, alpha, beta)
    if score <= alpha then search(context, depth, 0, -INF, beta)
    else if score >= beta then search(context, depth, 0, alpha, INF)
    else (score, move)

  private def hasNonPawnMaterial(context: GameContext): Boolean =
    context.board.pieces.values.exists { piece =>
      piece.color == context.turn &&
      piece.pieceType != PieceType.Pawn &&
      piece.pieceType != PieceType.King
    }

  private def nullMoveContext(context: GameContext): GameContext =
    context.withTurn(context.turn.opposite).withEnPassantSquare(None)

  private def tryNullMove(
    context: GameContext,
    depth: Int,
    ply: Int,
    beta: Int
  ): Option[Int] =
    val nullCtx = nullMoveContext(context)
    val reductionDepth = math.max(0, depth - 1 - NULL_MOVE_R)
    val (score, _) = search(nullCtx, reductionDepth, ply + 1, -beta, -beta + 1)
    if -score >= beta then Some(beta) else None

  /** Negamax alpha-beta search returning (score, best move). */
  private def search(
    context: GameContext,
    depth: Int,
    ply: Int,
    alpha: Int,
    beta: Int
  ): (Int, Option[Move]) =
    // Periodic time check
    nodeCount += 1
    if nodeCount % TIME_CHECK_FREQUENCY == 0 && isOutOfTime() then
      return (Evaluation.evaluate(context), None)

    val hash = ZobristHash.hash(context)

    // TT probe
    tt.probe(hash).foreach { entry =>
      if entry.depth >= depth then
        entry.flag match
          case TTFlag.Exact =>
            return (entry.score, entry.bestMove)
          case TTFlag.Lower =>
            val newAlpha = math.max(alpha, entry.score)
            if newAlpha >= beta then
              return (entry.score, entry.bestMove)
          case TTFlag.Upper =>
            val newBeta = math.min(beta, entry.score)
            if alpha >= newBeta then
              return (entry.score, entry.bestMove)
    }

    // Terminal node check
    val legalMoves = rules.allLegalMoves(context)
    if legalMoves.isEmpty then
      val score = if rules.isCheckmate(context) then
        -(Evaluation.CHECKMATE_SCORE - ply)
      else
        Evaluation.DRAW_SCORE
      return (score, None)

    if rules.isInsufficientMaterial(context) || rules.isFiftyMoveRule(context) then
      return (Evaluation.DRAW_SCORE, None)

    // Leaf node: call quiescence
    if depth == 0 then
      return (quiescence(context, ply, alpha, beta), None)

    // Null move pruning
    if depth >= 3 && !rules.isCheck(context) && hasNonPawnMaterial(context) then
      tryNullMove(context, depth, ply, beta).foreach { score =>
        return (score, None)
      }

    // Get TT best move for ordering
    val ttBest = tt.probe(hash).flatMap(_.bestMove)

    // Order moves
    val ordered = MoveOrdering.sort(context, legalMoves, ttBest)

    var bestMove: Option[Move] = None
    var bestScore = -INF
    var a = alpha
    var moveNumber = 0

    for move <- ordered do
      moveNumber += 1
      val isQuiet = !isCapture(move) &&
        move.moveType != MoveType.CastleKingside &&
        move.moveType != MoveType.CastleQueenside

      val child = rules.applyMove(context)(move)
      val reduction = if moveNumber > 4 && depth >= 3 && isQuiet then 1 else 0

      val score = if reduction > 0 then
        val (reducedScore, _) = search(child, depth - 1 - reduction, ply + 1, -a - 1, -a)
        val s = -reducedScore
        if s > a then
          val (fullScore, _) = search(child, depth - 1, ply + 1, -beta, -a)
          -fullScore
        else s
      else
        val (rawScore, _) = search(child, depth - 1, ply + 1, -beta, -a)
        -rawScore

      if score > bestScore then
        bestScore = score
        bestMove = Some(move)

      a = math.max(a, score)

      if a >= beta then
        tt.store(TTEntry(hash, depth, bestScore, TTFlag.Lower, bestMove))
        return (bestScore, bestMove)

    // No cutoff: determine flag
    val flag =
      if bestScore <= alpha then TTFlag.Upper
      else TTFlag.Exact
    tt.store(TTEntry(hash, depth, bestScore, flag, bestMove))
    (bestScore, bestMove)

  /** Quiescence search: only captures until position is quiet. */
  private def quiescence(
    context: GameContext,
    ply: Int,
    alpha: Int,
    beta: Int
  ): Int =
    // Stand-pat: evaluate current position
    val standPat = Evaluation.evaluate(context)
    if standPat >= beta then
      return beta

    var a = math.max(alpha, standPat)

    // Guard against infinite quiescence
    if ply >= MAX_QUIESCENCE_PLY then
      return standPat

    // Generate only captures
    val allMoves = rules.allLegalMoves(context)
    val captures = allMoves.filter(isCapture)
    val ordered = MoveOrdering.sort(context, captures, None)

    for move <- ordered do
      val child = rules.applyMove(context)(move)
      val score = -quiescence(child, ply + 1, -beta, -a)
      if score >= beta then
        return beta
      a = math.max(a, score)
    a

  /** Predicate: is a move a capture (including promotions)? */
  private def isCapture(move: Move): Boolean = move.moveType match
    case MoveType.Normal(true) => true
    case MoveType.EnPassant => true
    case MoveType.Promotion(_) => true
    case _ => false
