"""ChessTransformer — Hugging Face Space (Gradio).

Play the transformer chess engine in the browser. Runs on CPU: the model is
only 11.7M params, so MCTS at a modest sim count returns a move in ~1-2s.

Weights are resolved in this order:
  1. ``MODEL_PATH`` env var (a checkpoint dir with model.safetensors + config)
  2. a ``model/`` dir next to this file (how ``prepare.sh`` ships the Space)
  3. the repo's ``data/models/pos2move_v2.1`` (for local dev runs)
  4. download from GitHub (LFS media) as a last resort
"""

from __future__ import annotations

import itertools
import os
import sys
from pathlib import Path

import chess
import gradio as gr
from gradio_chessboard import Chessboard

# ── make `chesstransformer` importable ──────────────────────────────────────
# On the Space the package is vendored next to this file (cwd is on sys.path).
# For local dev runs from the repo, add the repo's src/ to the path.
_repo_src = Path(__file__).resolve().parents[1] / "src"
if (_repo_src / "chesstransformer").exists():
    sys.path.insert(0, str(_repo_src))

from chesstransformer.bots import Pos2MoveV2MctsBot  # noqa: E402

START_FEN = chess.STARTING_FEN
DEFAULT_SIMS = int(os.environ.get("MCTS_SIMS", "400"))

# GitHub locations for the last-resort weight download.
_GH = "tchauffi/ChessTransformer"
_BRANCH = os.environ.get("MODEL_BRANCH", "main")
_RAW = f"https://raw.githubusercontent.com/{_GH}/{_BRANCH}/data/models/pos2move_v2.1"
_LFS = f"https://media.githubusercontent.com/media/{_GH}/{_BRANCH}/data/models/pos2move_v2.1"


def _download_weights() -> str:
    import urllib.request

    dest = Path("/tmp/pos2move_v2.1")
    dest.mkdir(parents=True, exist_ok=True)
    for fname, url in [("model_config.json", f"{_RAW}/model_config.json"),
                       ("model.safetensors", f"{_LFS}/model.safetensors")]:
        out = dest / fname
        if not out.exists():
            print(f"Downloading {fname} …")
            urllib.request.urlretrieve(url, out)
    return str(dest)


def resolve_model_dir() -> str:
    env = os.environ.get("MODEL_PATH")
    if env and (Path(env) / "model.safetensors").exists():
        return env
    for cand in (
        Path(__file__).resolve().parent / "model",
        Path("model"),
        Path(__file__).resolve().parents[1] / "data" / "models" / "pos2move_v2.1",
    ):
        if (cand / "model.safetensors").exists():
            return str(cand)
    return _download_weights()


# ── load the bot once (CPU, no compile) ─────────────────────────────────────
MODEL_DIR = resolve_model_dir()
bot = Pos2MoveV2MctsBot(
    model_dir=MODEL_DIR,
    num_simulations=DEFAULT_SIMS,
    device="cpu",
    compile=False,
    time_limit=0.0,
    # Sample the opening moves (∝ visit count) so games aren't identical every
    # time; switches to deterministic best-move play after the opening.
    move_temp=1.0,
    move_temp_plies=12,
)


def _human_color(human_color: str) -> bool:
    return chess.BLACK if human_color == "black" else chess.WHITE


def _white_pov_value(board: chess.Board) -> float:
    """Value-head eval in [-1, 1] from White's perspective (+ = White better)."""
    if board.is_game_over():
        if board.is_checkmate():
            return -1.0 if board.turn == chess.WHITE else 1.0
        return 0.0  # stalemate / draw
    v = bot.get_value(board)
    if v is None:
        return 0.0
    return v if board.turn == chess.WHITE else -v


def eval_bar_html(board: chess.Board, orientation: str) -> str:
    """A vertical eval bar (White share fills from the bottom of White's side).
    Flipped to match the board orientation so the viewer's side is at the bottom."""
    wv = _white_pov_value(board)
    white_pct = round(max(0.0, min(1.0, (wv + 1) / 2)) * 100, 1)
    white_div = f'<div style="height:{white_pct}%;background:#f5f5f5;"></div>'
    black_div = f'<div style="height:{round(100 - white_pct, 1)}%;background:#3a3a3a;"></div>'
    # orientation "white": Black on top, White on bottom; flipped when "black".
    segments = black_div + white_div if orientation == "white" else white_div + black_div
    return (
        f'<div title="Model eval (White POV): {wv:+.2f}" '
        'style="display:flex;flex-direction:column;width:26px;height:460px;'
        'border:1px solid #999;border-radius:4px;overflow:hidden;">'
        f'{segments}</div>'
    )


def _result_text(board: chess.Board, human_color: str) -> str:
    outcome = board.outcome()
    if outcome is None:
        return ""
    if outcome.winner is None:
        return f"Draw ({outcome.termination.name.lower().replace('_', ' ')}). New game?"
    who = "You win! 🎉" if outcome.winner == _human_color(human_color) else "ChessTransformer wins."
    return f"Checkmate — {who} New game?"


def _bot_move(board: chess.Board, sims: int):
    """Play the side to move with the bot; return its SAN and value."""
    bot.num_simulations = int(sims)
    uci, value = bot.predict(board)
    san = board.san(chess.Move.from_uci(uci))
    board.push_uci(uci)
    return san, value


def on_move(fen: str, sims: int, human_color: str):
    """Human just moved; reply with the bot's move (whichever side is to move).
    The board orientation matches the human's color, so the eval bar does too.

    A generator: first yields a "thinking" status (so the player sees the bot is
    working during the ~1-3s CPU search), then yields the bot's reply."""
    board = chess.Board(fen)
    if board.is_game_over():
        yield board.fen(), _result_text(board, human_color), eval_bar_html(board, human_color)
        return

    # Immediate feedback while the search runs (board/eval unchanged for now).
    yield gr.update(), "⏳ ChessTransformer is thinking…", gr.update()

    san, value = _bot_move(board, sims)

    if board.is_game_over():
        status = f"Bot played **{san}**.  {_result_text(board, human_color)}"
    else:
        eval_str = f"  (eval {value:+.2f})" if value is not None else ""
        status = f"Bot played **{san}**{eval_str}.  Your move."
    yield board.fen(), status, eval_bar_html(board, human_color)


# Monotonic id so every New game yields a *distinct* setup dict — otherwise a
# same-colour replay returns an identical dict and the gr.render block, seeing no
# change, wouldn't rebuild (and reset) the board.
_game_counter = itertools.count(1)


def new_game(play_as: str, sims: int):
    """Start a fresh game. If the human plays Black, the bot (White) opens.

    Returns (setup, eval_bar). The setup dict drives a gr.render block — changing
    it rebuilds the board so its orientation (chessboard.js reads it only at
    construction) actually flips to the human's side."""
    human_color = "black" if play_as == "Black" else "white"
    board = chess.Board()
    if human_color == "black":
        san, _ = _bot_move(board, sims)
        status = f"New game — you are **Black**. Bot opened with **{san}**. Your move."
    else:
        status = "New game — you are **White**. Make your move."
    setup = {"fen": board.fen(), "orientation": human_color, "color": human_color,
             "status": status, "nonce": next(_game_counter)}
    return setup, eval_bar_html(board, human_color)


with gr.Blocks(title="ChessTransformer", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# ♟️ ChessTransformer\n"
        "An **11.7M-parameter transformer trained only on human games** "
        "(no self-play, no RL), playing via AlphaZero-style MCTS. "
        "It reaches **~2100 Elo** vs Stockfish at full strength; this CPU demo "
        "runs at a lower sim count so moves come back quickly.\n\n"
        "**Pick your color and hit New game.** Drag a piece to move — the bot "
        "replies automatically.  "
        "[GitHub repo →](https://github.com/tchauffi/ChessTransformer)"
    )
    setup = gr.State({"fen": START_FEN, "orientation": "white", "color": "white",
                      "status": "You are **White**. Make your move.", "nonce": 0})

    with gr.Row():
        # Solid min_widths so the board container always has a size when the
        # gr.render block (re)mounts it — chessboard.js reloads the whole page if
        # it inits at 0 width, which is what caused the mobile flicker/reload loop.
        eval_col = gr.Column(scale=0, min_width=44)
        board_col = gr.Column(scale=3, min_width=340)
        ctrl_col = gr.Column(scale=1, min_width=200)

    with eval_col:
        eval_bar = gr.HTML(eval_bar_html(chess.Board(), "white"), visible=False)

    with ctrl_col:
        play_as = gr.Radio(["White", "Black"], value="White", label="Play as")
        sims = gr.Slider(32, 1200, value=DEFAULT_SIMS, step=8,
                         label="Engine strength (MCTS sims/move)",
                         info="Higher = stronger but slower on CPU")
        show_eval = gr.Checkbox(False, label="Show evaluation bar")
        new_btn = gr.Button("New game", variant="primary")

    # The board lives in a gr.render so switching color rebuilds it with the new
    # orientation (chessboard.js reads orientation only at construction). The
    # min_width above keeps the container sized so the remount can't init at 0.
    with board_col:
        @gr.render(inputs=setup)
        def render_board(cfg):
            board = Chessboard(value=cfg["fen"], game_mode=True, min_width=320,
                               orientation=cfg["orientation"], label="Drag to move")
            status = gr.Markdown(cfg["status"])
            color = cfg["color"]

            def handle(fen, s):  # generator → streams the "thinking" update first
                yield from on_move(fen, s, color)

            # show_progress="hidden": our "⏳ thinking…" status is the signal, so
            # Gradio's default overlay doesn't sit over the board.
            board.move(handle, inputs=[board, sims], outputs=[board, status, eval_bar],
                       show_progress="hidden")

    show_eval.change(lambda show: gr.update(visible=show), inputs=show_eval, outputs=eval_bar)
    new_btn.click(new_game, inputs=[play_as, sims], outputs=[setup, eval_bar])


if __name__ == "__main__":
    demo.launch()