#import "board.typ": apply-move, at, copy-grid, file-to-col, find-pieces, is-black, is-white, on-board, piece-type, render-board, square-to-coords, starting-grid
#import "notation.typ": parse-move
#import "rules.typ": generate-moves-for, has-legal-move, is-attacked, is-in-check, pawn-moves

#let initial-state() = (
  grid: starting-grid,
  turn: "white",
  castles: "KQkq",
  ep: none,
  error: none,
  checkmate: false,
  stalemate: false,
)

#let apply-castling(state, side) = {
  let castles = state.castles
  let grid = copy-grid(state.grid)

  let row = if state.turn == "white" { 7 } else { 0 }
  let case-p = if state.turn == "white" { upper } else { lower }
  let (rook-from, rook-to, king-to) = if side == "kingside" { (7, 5, 6) } else { (0, 3, 2) }
  grid.at(row).at(4) = "."
  grid.at(row).at(rook-from) = "."
  grid.at(row).at(king-to) = case-p("K")
  grid.at(row).at(rook-to) = case-p("R")

  let lost = if state.turn == "white" { "KQ" } else { "kq" }
  for c in lost { castles = castles.replace(c, "") }

  (
    grid: grid,
    turn: if state.turn == "white" { "black" } else { "white" },
    castles: castles,
    ep: none,
    error: none,
    checkmate: false,
    stalemate: false,
  )
}

#let apply-parsed-move(state, parsed) = {
  if parsed.kind == "error" { return (..state, error: parsed.raw + ": illegal notation") }
  if parsed.kind == "castle" { return apply-castling(state, parsed.side) }

  let dest-col = parsed.dest.at(0)
  let dest-row = parsed.dest.at(1)

  // find the source pieces that can reach the destination
  let candidates = ()
  if parsed.piece == "P" {
    let pawn-file = none
    if parsed.disambig != none and parsed.disambig.type == "file" {
      pawn-file = file-to-col.at(parsed.disambig.value) + 1
    }

    for (col, row) in find-pieces(state.grid, state.turn, "P") {
      if pawn-file != none and col != pawn-file { continue }
      for m in pawn-moves(state.grid, col, row, state.turn, ep: state.ep) {
        if m.at(0) == dest-col and m.at(1) == dest-row {
          candidates.push((col, row))
          break
        }
      }
    }
  } else {
    for (col, row) in find-pieces(state.grid, state.turn, parsed.piece) {
      let moves = generate-moves-for(state.grid, col, row, state.turn, ep: state.ep)
      for m in moves {
        if m.at(0) == dest-col and m.at(1) == dest-row {
          candidates.push((col, row))
        }
      }
    }

    // apply disambiguation
    if candidates.len() > 1 and parsed.disambig != none {
      if parsed.disambig.type == "file" {
        let f = file-to-col.at(parsed.disambig.value) + 1
        candidates = candidates.filter(c => c.at(0) == f)

      } else if parsed.disambig.type == "rank" {
        let r = 9 - int(parsed.disambig.value)
        candidates = candidates.filter(c => c.at(1) == r)

      } else if parsed.disambig.type == "square" {
        let (cf, rf) = square-to-coords(parsed.disambig.value)
        candidates = candidates.filter(c => c.at(0) == cf and c.at(1) == rf)
      }
    }
  }

  // auto-promote to queen
  let promotion = parsed.promotion
  if promotion == none and parsed.piece == "P" and (dest-row == 1 or dest-row == 8) { promotion = "Q" }
  if promotion != none and "QRBN".position(upper(promotion)) == none { return (..state, error: parsed.raw + "=" + promotion + ": illegal notation") }

  // check if it's an en passant capture
  let is-ep = parsed.piece == "P" and state.ep != none and dest-col == state.ep.at(0) and dest-row == state.ep.at(1)
  let ep-captured = if is-ep { (dest-col, if state.turn == "white" { dest-row + 1 } else { dest-row - 1 }) } else { none }

  // filter to only legal moves
  let legal-candidates = ()
  for (sc, sr) in candidates {
    let sim = apply-move(state.grid, sc, sr, dest-col, dest-row, promotion: promotion, en-passant: ep-captured)
    if not is-in-check(sim, state.turn) { legal-candidates.push((sc, sr)) }
  }
  candidates = legal-candidates

  if candidates.len() == 0 { return (..state, error: parsed.raw + ": illegal move") }
  let (src-col, src-row) = candidates.at(0)

  let new-grid = apply-move(state.grid, src-col, src-row, dest-col, dest-row, promotion: promotion, en-passant: ep-captured)
  let new-turn = if state.turn == "white" { "black" } else { "white" }

  // update castling rights
  let castles = state.castles
  let lost = if parsed.piece == "K" { if state.turn == "white" { "KQ" } else { "kq" } } else {
    let flags = ""
    for (rc, rr, flag) in ((1, 8, "Q"), (8, 8, "K"), (1, 1, "q"), (8, 1, "k")) {
      if src-col == rc and src-row == rr { flags += flag }
      if dest-col == rc and dest-row == rr { flags += flag }
    }
    flags
  }
  for c in lost { castles = castles.replace(c, "") }

  // set en passant target for the next move
  let new-ep = none
  if parsed.piece == "P" and calc.abs(dest-row - src-row) == 2 { new-ep = (dest-col, int((src-row + dest-row) / 2)) }

  // check for checkmate or stalemate
  let in-check = is-in-check(new-grid, new-turn)
  let no-legal = not has-legal-move(new-grid, new-turn, ep: new-ep, castles: castles)

  (
    grid: new-grid,
    turn: new-turn,
    castles: castles,
    ep: new-ep,
    error: none,
    checkmate: in-check and no-legal,
    stalemate: not in-check and no-legal,
  )
}

#let apply-moves(moves-strs) = {
  let state = initial-state()
  let states = (state,)
  for ms in moves-strs {
    if state.at("checkmate", default: false) or state.at("stalemate", default: false) { break }
    let new-state = apply-parsed-move(state, parse-move(ms))
    states += (new-state,)
    state = new-state
  }
  states
}

#let parse-moves(body) = {
  let extract(it) = {
    if it == [ ] { " " } else if it.func() == text { it.text } else if it.func() == [].func() {
      it.children.map(extract).join()
    } else { "" }
  }
  extract(body).split().filter(s => s != "")
}

#let render(state) = {
  let banner(body) = text(fill: rgb("#0a7e8c"), weight: "bold", size: 1.2em, body)

  let out = ()
  if state.at("checkmate", default: false)      { out.push(banner[Checkmate!]) }
  else if state.at("stalemate", default: false) { out.push(banner[Stalemate!]) }
  else {
    let err = state.at("error", default: none)
    if err != none                              { out.push(text(fill: red, weight: "bold", err)) }
    else if is-in-check(state.grid, state.turn) { out.push(text(fill: yellow, weight: "bold")[Check!]) }
  }
  out.push(v(0.3em))
  out.push(render-board(state.grid))
  out.join()
}

// Usage:
//   #show: chess.with()
//   g4 d5 Bg5 Bxg4 c4 c6 cxd5 cxd5 Qb3
//
#let chess(body) = {
  set page(height: auto, width: auto, margin: (top: 0.5in, bottom: 0.5in, rest: 0.5in))
  let moves = parse-moves(body)
  let states = apply-moves(moves)
  render(states.last())
}