#import "board.typ": at, copy-grid, find-pieces, is-black, is-white, on-board, piece-type
#import "moves.typ": bishop-dirs, generate-moves-for, king-dirs, knight-offsets, pawn-moves, rook-dirs

#let is-attacked(grid, col, row, by-color) = {
  let enemy(piece) = if by-color == "white" { is-white(piece) } else { is-black(piece) }

  // knight attacks
  for (dc, dr) in knight-offsets {
    let p = at(grid, col + dc, row + dr)
    if p != "." and piece-type(p) == "N" and enemy(p) { return true }
  }

  // king attacks
  for (dc, dr) in king-dirs {
    let p = at(grid, col + dc, row + dr)
    if p != "." and piece-type(p) == "K" and enemy(p) { return true }
  }

  // pawn attacks
  let pawn-checks = if by-color == "white" { ((-1, 1), (1, 1)) } else { ((-1, -1), (1, -1)) }
  for (dc, dr) in pawn-checks {
    let p = at(grid, col + dc, row + dr)
    if p != "." and piece-type(p) == "P" and enemy(p) { return true }
  }

  // sliding pieces
  for (dirs, types) in ((bishop-dirs, ("B", "Q")), (rook-dirs, ("R", "Q"))) {
    for (dc, dr) in dirs {
      let nc = col + dc
      let nr = row + dr
      while on-board(nc, nr) {
        let p = at(grid, nc, nr)
        if p != "." {
          if enemy(p) and types.contains(piece-type(p)) { return true }
          break // blocked by any piece
        }
        nc += dc
        nr += dr
      }
    }
  }

  false
}

#let is-in-check(grid, turn) = {
  let king-char = if turn == "white" { "K" } else { "k" }
  let enemy = if turn == "white" { "black" } else { "white" }
  for r in range(8) {
    for c in range(8) {
      if grid.at(r).at(c) == king-char {
        return is-attacked(grid, c + 1, r + 1, enemy)
      }
    }
  }
  false
}

#let can-castle(grid, row, empty-cols, check-cols, enemy) = {
  if is-attacked(grid, 5, row, enemy) { return false }
  for c in empty-cols { if at(grid, c, row) != "." { return false } }
  for c in check-cols { if is-attacked(grid, c, row, enemy) { return false } }
  true
}

#let has-legal-move(grid, turn, ep: none, castles: "") = {
  for r in range(8) {
    for c in range(8) {
      let p = grid.at(r).at(c)
      if p == "." { continue }
      if turn == "white" and not is-white(p) { continue }
      if turn == "black" and not is-black(p) { continue }

      let type = piece-type(p)

      let col = c + 1
      let row = r + 1
      let raw-moves = if type == "P" { pawn-moves(grid, col, row, turn, ep: ep) } else {
        generate-moves-for(grid, col, row, turn, ep: ep)
      }

      for m in raw-moves {
        let (dc, dr) = (m.at(0), m.at(1))
        let new-grid = copy-grid(grid)
        new-grid.at(r).at(c) = "."
        let mover = p
        let is-pawn = type == "P"
        if is-pawn and (dr == 1 or dr == 8) {
          mover = if turn == "white" { "Q" } else { "q" }
        }
        new-grid.at(dr - 1).at(dc - 1) = mover
        if is-pawn and ep != none and dc == ep.at(0) and dr == ep.at(1) {
          let captured-row = if turn == "white" { dr + 1 } else { dr - 1 }
          new-grid.at(captured-row - 1).at(dc - 1) = "."
        }
        if not is-in-check(new-grid, turn) { return true }
      }
    }
  }

  let enemy = if turn == "white" { "black" } else { "white" }
  let options = if turn == "white" {
    (("K", 8, (6, 7), (6, 7)), ("Q", 8, (2, 3, 4), (3, 4)))
  } else {
    (("k", 1, (6, 7), (6, 7)), ("q", 1, (2, 3, 4), (3, 4)))
  }

  for (flag, row, empty, check) in options {
    if castles.contains(flag) and can-castle(grid, row, empty, check, enemy) { return true }
  }

  false
}