package de.nowchess.bot.logic

import de.nowchess.api.board.PieceType
import de.nowchess.api.game.GameContext
import de.nowchess.api.move.{Move, MoveType}
import de.nowchess.bot.ai.Weights
import de.nowchess.bot.logic.{MoveOrdering, TTEntry, TTFlag, TranspositionTable}
import de.nowchess.bot.util.ZobristHash
import de.nowchess.rules.RuleSet
import de.nowchess.rules.sets.DefaultRules

import java.util.concurrent.ForkJoinPool
import java.util.concurrent.atomic.AtomicReference
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.{ExecutionContext, Future}

final class AlphaBetaSearch(
  rules: RuleSet = DefaultRules,
  tt: TranspositionTable = TranspositionTable(),
  weights: Weights,
  numThreads: Int = Runtime.getRuntime().availableProcessors()
):

  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 ASPIRATION_DELTA_MAX = 150
  private val TIME_CHECK_FREQUENCY = 1000
  private val FUTILITY_MARGIN = 100
  private val PARALLEL_DEPTH_THRESHOLD = 4

  @volatile private var timeStartMs = 0L
  @volatile private var timeLimitMs = 0L
  @volatile private var nodeCount = 0
  private val ordering = MoveOrdering.OrderingContext()
  private val threadPool = new ForkJoinPool(numThreads)
  implicit private val executionContext: ExecutionContext = ExecutionContext.fromExecutor(threadPool)

  /** 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()
    ordering.clear()
    var bestSoFar: Option[Move] = None
    var prevScore = 0
    var aspWindow = ASPIRATION_DELTA
    for depth <- 1 to maxDepth do
      val (alpha, beta) =
        if depth == 1 then (-INF, INF)
        else (prevScore - aspWindow, prevScore + aspWindow)
      val (score, move) = searchWithAspiration(context, depth, alpha, beta, aspWindow)
      prevScore = score
      move.foreach(m => bestSoFar = Some(m))
      aspWindow = ASPIRATION_DELTA
    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()
    ordering.clear()
    timeStartMs = System.currentTimeMillis()
    timeLimitMs = timeBudgetMs
    nodeCount = 0
    var bestSoFar: Option[Move] = None
    var prevScore = 0
    var depth = 1
    var aspWindow = ASPIRATION_DELTA

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

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

  private def searchWithAspiration(
    context: GameContext,
    depth: Int,
    alpha: Int,
    beta: Int,
    initialWindow: Int
  ): (Int, Option[Move]) =
    var currentAlpha = alpha
    var currentBeta = beta
    var window = initialWindow
    var attempt = 0

    while attempt < 3 && attempt < depth do
      val (score, move) = search(context, depth, 0, currentAlpha, currentBeta)
      if score > currentAlpha && score < currentBeta then
        return (score, move)
      if score <= currentAlpha then
        currentAlpha = score - window
        window = math.min(window * 2, ASPIRATION_DELTA_MAX)
      if score >= currentBeta then
        currentBeta = score + window
        window = math.min(window * 2, ASPIRATION_DELTA_MAX)
      attempt += 1

    search(context, depth, 0, -INF, INF)

  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 (weights.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
        -(weights.CHECKMATE_SCORE - ply)
      else
        weights.DRAW_SCORE
      return (score, None)

    if rules.isInsufficientMaterial(context) || rules.isFiftyMoveRule(context) then
      return (weights.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, ply, ordering)

    // Use parallel search if depth >= threshold and not root
    if depth >= PARALLEL_DEPTH_THRESHOLD && ply > 0 then
      return searchParallel(context, depth, ply, alpha, beta, ordered, hash)
    else
      return searchSequential(context, depth, ply, alpha, beta, ordered, hash)

  private def searchSequential(
    context: GameContext,
    depth: Int,
    ply: Int,
    alpha: Int,
    beta: Int,
    ordered: List[Move],
    hash: Long
  ): (Int, Option[Move]) =
    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

      // Futility pruning at frontier nodes: if static eval + margin is still below alpha, skip quiet moves
      if depth == 1 && isQuiet && moveNumber > 2 then
        val staticEval = weights.evaluate(context)
        if staticEval + FUTILITY_MARGIN < alpha then
          moveNumber += 1

      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)

      // Track history heuristic
      if isQuiet then
        val fromIdx = move.from.rank.ordinal * 8 + move.from.file.ordinal
        val toIdx = move.to.rank.ordinal * 8 + move.to.file.ordinal
        ordering.addHistory(fromIdx, toIdx, depth * depth)

      if a >= beta then
        // Record killer move
        if isQuiet then
          ordering.addKillerMove(ply, move)
        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)

  private def searchParallel(
    context: GameContext,
    depth: Int,
    ply: Int,
    alpha: Int,
    beta: Int,
    ordered: List[Move],
    hash: Long
  ): (Int, Option[Move]) =
    val results = new AtomicReference[(Int, Option[Move], Boolean)]((-INF, None, false))
    val windowRef = new AtomicReference((alpha, beta))

    val moveScores = ordered.zipWithIndex.map { case (move, moveIdx) =>
      Future {
        val isQuiet = !isCapture(move) &&
          move.moveType != MoveType.CastleKingside &&
          move.moveType != MoveType.CastleQueenside
        val child = rules.applyMove(context)(move)
        val reduction = if moveIdx > 4 && depth >= 3 && isQuiet then 1 else 0

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

        // Track history heuristic
        if isQuiet then
          val fromIdx = move.from.rank.ordinal * 8 + move.from.file.ordinal
          val toIdx = move.to.rank.ordinal * 8 + move.to.file.ordinal
          ordering.addHistory(fromIdx, toIdx, depth * depth)

        (move, score, isQuiet)
      }
    }

    var bestMove: Option[Move] = None
    var bestScore = -INF
    var cutoffFound = false

    for future <- moveScores do
      if !cutoffFound then
        val (move, score, isQuiet) = scala.concurrent.Await.result(future, scala.concurrent.duration.Duration.Inf)

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

        if bestScore >= beta then
          if isQuiet then
            ordering.addKillerMove(ply, move)
          cutoffFound = true

    // 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 = weights.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
