"""
ChessSLM — Play against FlameF0X/ChessSLM
Hugging Face Space  |  Gradio app

Move selection: ChessSLM generates 10–20 candidate moves via sampling,
Stockfish evaluates each and selects the best one.
"""

import re
import random
import chess
import chess.svg
import chess.pgn
import chess.engine
import gradio as gr
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# ──────────────────────────────────────────────────────────────────────────────
# Model loading (once at startup)
# ──────────────────────────────────────────────────────────────────────────────

MODEL_ID = "FlameF0X/ChessSLM-RL"

print(f"Loading {MODEL_ID}...")
device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = GPT2Tokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token
model     = GPT2LMHeadModel.from_pretrained(MODEL_ID)
model.to(device)
model.eval()
model.config.use_cache = True
print(f"✓ Model ready on {device}")

# Stockfish engine path — standard location on HuggingFace Spaces / Ubuntu
STOCKFISH_PATH = "/usr/games/stockfish"
STOCKFISH_DEPTH = 12   # evaluation depth per candidate

# ──────────────────────────────────────────────────────────────────────────────
# Chess / model logic
# ──────────────────────────────────────────────────────────────────────────────

def board_to_prompt(board: chess.Board) -> str:
    game = chess.pgn.Game()
    node = game
    for move in board.move_stack:
        node = node.add_variation(move)
    exporter = chess.pgn.StringExporter(headers=False, variations=False, comments=False)
    pgn = game.accept(exporter).strip()
    pgn = re.sub(r"\s*[\*\d][-\d/]*\s*$", "", pgn).strip()
    full_move = board.fullmove_number
    pgn += f" {full_move}." if board.turn == chess.WHITE else f" {full_move}..."
    return f"<|endoftext|>{pgn}"


def extract_move(text: str, board: chess.Board):
    text = re.sub(r"^\s*\d+\.+\s*", "", text).strip()
    for token in text.split()[:5]:
        clean = re.sub(r"[!?+#,;]+$", "", token)
        try:
            move = board.parse_san(clean)
            if move in board.legal_moves:
                return move
        except Exception:
            pass
        try:
            move = chess.Move.from_uci(clean.lower()[:4])
            if move in board.legal_moves:
                return move
        except Exception:
            pass
    return None


@torch.no_grad()
def _sample_candidate_moves(board: chess.Board, n: int = 15, max_attempts: int = 90):
    """
    Sample up to `n` distinct legal moves from ChessSLM using temperature sampling.
    Returns a list of chess.Move objects.
    """
    prompt  = board_to_prompt(board)
    inputs  = tokenizer(prompt, return_tensors="pt").to(device)

    candidates: list[chess.Move] = []
    seen: set[str] = set()
    attempts = 0

    while len(candidates) < n and attempts < max_attempts:
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=12,
            do_sample=True,
            temperature=0.85,       # higher temp → more diverse candidates
            top_k=60,
            top_p=0.95,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
        new_tokens = outputs[0][inputs.input_ids.shape[1]:]
        generated  = tokenizer.decode(new_tokens, skip_special_tokens=True)
        move       = extract_move(generated, board)
        attempts  += 1

        if move is not None and move.uci() not in seen:
            seen.add(move.uci())
            candidates.append(move)

    return candidates


def _stockfish_best(board: chess.Board, candidates: list[chess.Move]):
    """
    Use Stockfish to evaluate each candidate move and return
    (best_move, scores_dict, stockfish_available).

    scores_dict maps move-UCI → centipawn score (from the side to move).
    """
    if not candidates:
        return None, {}, False

    scores: dict[str, int] = {}

    try:
        with chess.engine.SimpleEngine.popen_uci(STOCKFISH_PATH) as engine:
            for move in candidates:
                test_board = board.copy()
                test_board.push(move)
                info  = engine.analyse(test_board, chess.engine.Limit(depth=STOCKFISH_DEPTH))
                score = info["score"].pov(board.turn)   # from the perspective of the side to move

                # Represent as centipawns (mate scores → ±9999)
                if score.is_mate():
                    cp = 9999 if score.mate() > 0 else -9999
                else:
                    cp = score.score()

                scores[move.uci()] = cp

        # Best move = highest cp from the side to move
        best_uci  = max(scores, key=lambda u: scores[u])
        best_move = chess.Move.from_uci(best_uci)
        return best_move, scores, True

    except FileNotFoundError:
        # Stockfish not installed — fall back to first candidate
        return candidates[0], {}, False
    except Exception:
        return candidates[0], {}, False


def get_model_move(board: chess.Board):
    """
    High-level entry point.
    Returns (chosen_move, was_model_legal, candidates, scores, stockfish_used).
    """
    n_samples  = random.randint(10, 20)
    candidates = _sample_candidate_moves(board, n=n_samples, max_attempts=n_samples * 6)

    if not candidates:
        # Model produced nothing valid; pure random fallback
        fallback = random.choice(list(board.legal_moves))
        return fallback, False, [], {}, False

    best, scores, sf_used = _stockfish_best(board, candidates)

    if best is None:
        best = candidates[0]

    return best, True, candidates, scores, sf_used


# ──────────────────────────────────────────────────────────────────────────────
# Board rendering
# ──────────────────────────────────────────────────────────────────────────────

PIECE_COLORS = {
    "square light":            "#f0d9b5",
    "square dark":             "#b58863",
    "square light lastmove":   "#cdd16e",
    "square dark lastmove":    "#aaa23a",
}


def render_board_html(board: chess.Board, last_move=None, flipped=False, size=480):
    check_square = board.king(board.turn) if board.is_check() else None
    svg = chess.svg.board(
        board,
        lastmove=last_move,
        check=check_square,
        flipped=flipped,
        size=size,
        colors=PIECE_COLORS,
    )
    return f"""
    <div style="
        display:flex; justify-content:center; align-items:center;
        padding: 16px;
        background: radial-gradient(ellipse at center, #1a1208 0%, #0d0d0d 100%);
        border-radius: 12px;
        box-shadow: 0 0 60px rgba(0,0,0,0.8), inset 0 0 30px rgba(0,0,0,0.4);
    ">
        <div style="
            border-radius: 4px;
            overflow: hidden;
            box-shadow: 0 8px 32px rgba(0,0,0,0.6), 0 0 0 3px #3d2b0e, 0 0 0 5px #6b4c1e;
        ">
            {svg}
        </div>
    </div>
    """


def get_legal_moves_san(board: chess.Board):
    moves = []
    for move in board.legal_moves:
        try:
            moves.append(board.san(move))
        except Exception:
            pass
    return sorted(moves)


def format_move_history(board: chess.Board):
    if not board.move_stack:
        return "<em style='color:#666'>No moves yet.</em>"

    temp  = chess.Board()
    lines = []
    moves = list(board.move_stack)
    i     = 0

    while i < len(moves):
        move_num  = temp.fullmove_number
        white_san = temp.san(moves[i])
        temp.push(moves[i])
        i += 1

        if i < len(moves):
            black_san = temp.san(moves[i])
            temp.push(moves[i])
            i += 1
            lines.append(
                f"<span style='color:#8a7a5a;font-size:0.8em'>{move_num}.</span> "
                f"<span style='color:#e8d5a3'>{white_san}</span> "
                f"<span style='color:#c4b48a'>{black_san}</span>"
            )
        else:
            lines.append(
                f"<span style='color:#8a7a5a;font-size:0.8em'>{move_num}.</span> "
                f"<span style='color:#e8d5a3'>{white_san}</span>"
            )

    visible = lines[-10:]
    html  = "<div style='font-family:\"Courier New\",monospace; line-height:2; font-size:0.92em;'>"
    html += "<br>".join(visible)
    html += "</div>"
    return html


def format_candidates_html(board: chess.Board, candidates: list, scores: dict,
                            chosen_move, sf_used: bool) -> str:
    """
    Render a small table showing all ChessSLM candidates with their
    Stockfish evaluation scores and which one was chosen.
    """
    if not candidates:
        return ""

    rows = []
    for move in candidates:
        san     = board.san(move) if move in board.legal_moves else move.uci()
        cp      = scores.get(move.uci())
        is_best = (chosen_move is not None and move.uci() == chosen_move.uci())

        if cp is None:
            score_str = "<span style='color:#666'>—</span>"
        elif cp >= 9999:
            score_str = "<span style='color:#6cf06c'>+M</span>"
        elif cp <= -9999:
            score_str = "<span style='color:#f06c6c'>−M</span>"
        elif cp > 0:
            score_str = f"<span style='color:#a8d88a'>+{cp/100:.2f}</span>"
        elif cp < 0:
            score_str = f"<span style='color:#d88a8a'>{cp/100:.2f}</span>"
        else:
            score_str = "<span style='color:#aaa'>0.00</span>"

        star = "⭐ " if is_best else "&nbsp;&nbsp;&nbsp;"
        row  = (
            f"<tr style='background:{'#1e1a08' if is_best else 'transparent'};'>"
            f"<td style='padding:2px 6px;color:#c8a96e'>{star}</td>"
            f"<td style='padding:2px 8px;font-weight:{'700' if is_best else '400'};color:#e8d5a3'>{san}</td>"
            f"<td style='padding:2px 8px;text-align:right'>{score_str}</td>"
            f"</tr>"
        )
        rows.append(row)

    sf_label = (
        "<span style='color:#8a6;font-size:0.8em'>✔ Stockfish</span>"
        if sf_used else
        "<span style='color:#a64;font-size:0.8em'>⚠ Stockfish N/A — first candidate used</span>"
    )

    return (
        f"<div style='margin-top:6px'>"
        f"<span style='font-family:Cinzel,serif;font-size:0.75em;color:#8a7a5a;"
        f"letter-spacing:0.08em;text-transform:uppercase'>ChessSLM Candidates</span>"
        f"&nbsp;{sf_label}"
        f"<table style='width:100%;border-collapse:collapse;margin-top:4px;font-family:\"Courier New\",monospace;font-size:0.88em'>"
        + "".join(rows)
        + "</table></div>"
    )


def game_status(board: chess.Board, player_color: str):
    if board.is_checkmate():
        winner = "Black" if board.turn == chess.WHITE else "White"
        if (winner == "White") == (player_color == "white"):
            return "♟ Checkmate — You win! 🎉", "win"
        else:
            return "♟ Checkmate — ChessSLM wins!", "loss"
    if board.is_stalemate():
        return "½ Stalemate — Draw", "draw"
    if board.is_insufficient_material():
        return "½ Insufficient material — Draw", "draw"
    if board.is_seventyfive_moves():
        return "½ 75-move rule — Draw", "draw"
    if board.is_fivefold_repetition():
        return "½ Fivefold repetition — Draw", "draw"
    if board.is_check():
        return "⚠ Check!", "check"
    whose = "Your turn" if (board.turn == chess.WHITE) == (player_color == "white") else "ChessSLM is thinking..."
    return whose, "playing"


# ──────────────────────────────────────────────────────────────────────────────
# Gradio callbacks
# ──────────────────────────────────────────────────────────────────────────────

def new_game(player_color_choice: str):
    board        = chess.Board()
    player_color = "white" if player_color_choice == "⬜ White (move first)" else "black"
    flipped      = (player_color == "black")
    last_move    = None
    log_lines    = []
    candidates_html = ""

    if player_color == "black":
        move, legal, candidates, scores, sf_used = get_model_move(board)
        san       = board.san(move)

        # Build candidate table before pushing
        cand_html = format_candidates_html(board, candidates, scores, move, sf_used)

        board.push(move)
        last_move = move
        log_lines.append(f"ChessSLM opens with <strong>{san}</strong>")
        candidates_html = cand_html

    legal_moves = get_legal_moves_san(board)
    status_text, _ = game_status(board, player_color)
    board_html      = render_board_html(board, last_move=last_move, flipped=flipped)
    history_html    = format_move_history(board)

    if log_lines:
        log_html = (
            "<br>".join(f"<span style='color:#c8a96e'>{l}</span>" for l in log_lines)
            + candidates_html
        )
    else:
        log_html = "<em style='color:#666'>Game started.</em>"

    state = {
        "fen":           board.fen(),
        "move_stack":    [m.uci() for m in board.move_stack],
        "player_color":  player_color,
        "last_move_uci": last_move.uci() if last_move else None,
        "game_over":     False,
    }

    return (
        board_html,
        gr.Dropdown(choices=legal_moves, value=None, interactive=True, label="Your move"),
        status_text,
        history_html,
        log_html,
        state,
    )


def make_player_move(move_san: str, state: dict):
    if not state or state.get("game_over"):
        return gr.update(), gr.update(), "Game is over. Start a new game.", gr.update(), gr.update(), state
    if not move_san:
        return gr.update(), gr.update(), "Please select a move first.", gr.update(), gr.update(), state

    board = chess.Board()
    for uci in state["move_stack"]:
        board.push(chess.Move.from_uci(uci))

    player_color = state["player_color"]
    flipped      = (player_color == "black")
    log_lines    = []
    candidates_html = ""

    try:
        player_move = board.parse_san(move_san)
    except Exception:
        return gr.update(), gr.update(), f"Invalid move: {move_san}", gr.update(), gr.update(), state

    board.push(player_move)
    log_lines.append(f"You played <strong>{move_san}</strong>")
    last_move = player_move

    status_text, status_key = game_status(board, player_color)
    game_over = status_key in ("win", "loss", "draw")

    if not game_over:
        move, legal, candidates, scores, sf_used = get_model_move(board)
        model_san = board.san(move)

        # Build candidate table before pushing
        candidates_html = format_candidates_html(board, candidates, scores, move, sf_used)

        board.push(move)
        last_move = move

        flag = "" if legal else " <em style='color:#888'>(random fallback)</em>"
        log_lines.append(f"ChessSLM plays <strong>{model_san}</strong>{flag}")

        status_text, status_key = game_status(board, player_color)
        game_over = status_key in ("win", "loss", "draw")

    state = {
        "fen":           board.fen(),
        "move_stack":    [m.uci() for m in board.move_stack],
        "player_color":  player_color,
        "last_move_uci": last_move.uci() if last_move else None,
        "game_over":     game_over,
    }

    legal_moves  = [] if game_over else get_legal_moves_san(board)
    board_html   = render_board_html(board, last_move=last_move, flipped=flipped)
    history_html = format_move_history(board)
    log_html     = (
        "<br>".join(f"<span style='color:#c8a96e'>{l}</span>" for l in log_lines)
        + candidates_html
    )

    return (
        board_html,
        gr.Dropdown(choices=legal_moves, value=None, interactive=not game_over, label="Your move"),
        status_text,
        history_html,
        log_html,
        state,
    )


# ──────────────────────────────────────────────────────────────────────────────
# CSS
# ──────────────────────────────────────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Cinzel:wght@400;600;700&family=Crimson+Text:ital,wght@0,400;0,600;1,400&display=swap');

body, .gradio-container {
    background: #0d0d0d !important;
    color: #e8d5a3 !important;
}

.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
    font-family: 'Crimson Text', Georgia, serif !important;
}

h1, h2, h3 {
    font-family: 'Cinzel', serif !important;
    letter-spacing: 0.08em;
}

#title-block {
    text-align: center;
    padding: 2rem 0 1rem;
    border-bottom: 1px solid #3d2b0e;
    margin-bottom: 1.5rem;
}

.panel {
    background: #141008 !important;
    border: 1px solid #3d2b0e !important;
    border-radius: 8px !important;
    padding: 1rem !important;
}

#status-bar {
    text-align: center;
    font-family: 'Cinzel', serif;
    font-size: 1.1em;
    letter-spacing: 0.05em;
    padding: 0.6rem 1rem;
    border-radius: 6px;
    background: #1a1208;
    border: 1px solid #4a3520;
    color: #f0c060;
}

button.primary {
    background: linear-gradient(135deg, #8b6914 0%, #c4922a 50%, #8b6914 100%) !important;
    border: 1px solid #d4a843 !important;
                color: #fff8e8 !important;
    font-family: 'Cinzel', serif !important;
    letter-spacing: 0.06em !important;
    font-size: 0.9em !important;
    border-radius: 4px !important;
    transition: all 0.2s ease !important;
}
button.primary:hover {
    background: linear-gradient(135deg, #a07820 0%, #d4a843 50%, #a07820 100%) !important;
    box-shadow: 0 0 16px rgba(212,168,67,0.4) !important;
}

button.secondary {
    background: #1e1810 !important;
    border: 1px solid #5a4020 !important;
    color: #c8a96e !important;
    font-family: 'Cinzel', serif !important;
    letter-spacing: 0.04em !important;
    border-radius: 4px !important;
}

select, .gr-dropdown select {
    background: #1a1208 !important;
    border: 1px solid #5a4020 !important;
    color: #e8d5a3 !important;
    font-family: 'Crimson Text', serif !important;
    font-size: 1em !important;
}

#move-log {
    background: #0f0c06 !important;
    border: 1px solid #3d2b0e !important;
    border-radius: 6px;
    padding: 0.8rem 1rem;
    font-family: 'Crimson Text', serif;
    font-size: 0.95em;
    line-height: 1.8;
    min-height: 80px;
    color: #c8a96e;
}

#history-panel {
    background: #0f0c06 !important;
    border: 1px solid #3d2b0e !important;
    border-radius: 6px;
    padding: 0.8rem 1rem;
    min-height: 200px;
    max-height: 320px;
    overflow-y: auto;
}

.gr-radio label {
    color: #e8d5a3 !important;
    font-family: 'Crimson Text', serif !important;
}

label span {
    color: #a08050 !important;
    font-family: 'Cinzel', serif !important;
    font-size: 0.8em !important;
    letter-spacing: 0.06em !important;
    text-transform: uppercase !important;
}

::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: #0d0d0d; }
::-webkit-scrollbar-thumb { background: #5a4020; border-radius: 3px; }
"""

# ──────────────────────────────────────────────────────────────────────────────
# Layout
# ──────────────────────────────────────────────────────────────────────────────

with gr.Blocks(css=CSS, title="ChessSLM — Play vs AI") as demo:

    state = gr.State({})

    gr.HTML("""
    <div id="title-block">
        <h1 style="
            font-family:'Cinzel',serif;
            font-size:2.4em;
            font-weight:700;
            color:#e8c96e;
            text-shadow: 0 0 30px rgba(232,180,80,0.4);
            margin:0 0 0.3rem;
            letter-spacing:0.12em;
        ">♛ ChessSLM</h1>
        <p style="
            font-family:'Crimson Text',serif;
            color:#8a7a5a;
            font-size:1.1em;
            font-style:italic;
            margin:0;
        ">GPT-2 generates 10–20 candidate moves · Stockfish picks the best</p>
    </div>
    """)

    with gr.Row():

        with gr.Column(scale=3):
            board_display = gr.HTML(
                value=render_board_html(chess.Board()),
                label="Board"
            )
            status_display = gr.HTML(
                value="<div id='status-bar'>Choose your colour and press New Game</div>"
            )

        with gr.Column(scale=2):

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>NEW GAME</h3>")

            color_choice = gr.Radio(
                choices=["⬜ White (move first)", "⬛ Black (move second)"],
                value="⬜ White (move first)",
                label="Play as",
            )

            new_game_btn = gr.Button("♟  New Game", variant="primary", size="lg")

            gr.HTML("<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>")

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>YOUR MOVE</h3>")

            move_dropdown = gr.Dropdown(
                choices=[],
                value=None,
                label="Select move (SAN notation)",
                interactive=False,
            )

            move_btn = gr.Button("▶  Make Move", variant="secondary")

            gr.HTML("<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>")

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>MOVE LOG</h3>")

            log_display = gr.HTML(
                value="<div id='move-log'><em style='color:#555'>Start a new game to begin.</em></div>",
            )

            gr.HTML("<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>")

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>GAME HISTORY</h3>")

            history_display = gr.HTML(
                value="<div id='history-panel'><em style='color:#555'>No moves yet.</em></div>",
            )

    gr.HTML("""
    <div style="
        text-align:center; margin-top:2rem; padding-top:1rem;
        border-top:1px solid #2a1e0a;
        font-family:'Crimson Text',serif; font-size:0.85em; color:#5a4a30;
    ">
        Model: <a href="https://huggingface.co/FlameF0X/ChessSLM" target="_blank"
            style="color:#8a6a30; text-decoration:none;">FlameF0X/ChessSLM</a>
        &nbsp;·&nbsp; GPT-2 samples 10–20 candidates (temp=0.85) · Stockfish depth=12 selects the best
    </div>
    """)

    new_game_btn.click(
        fn=new_game,
        inputs=[color_choice],
        outputs=[board_display, move_dropdown, status_display, history_display, log_display, state],
    )

    move_btn.click(
        fn=make_player_move,
        inputs=[move_dropdown, state],
        outputs=[board_display, move_dropdown, status_display, history_display, log_display, state],
    )

    move_dropdown.select(
        fn=make_player_move,
        inputs=[move_dropdown, state],
        outputs=[board_display, move_dropdown, status_display, history_display, log_display, state],
    )


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