"""
ChessQween — Play against zhoudoe23/CheesQween
Hugging Face Space  |  Gradio app
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import re
import random
import chess
import chess.svg
import chess.pgn
import gradio as gr
import torch
import spaces

torch.cuda.is_available = lambda: False

from transformers import GPT2LMHeadModel, GPT2Tokenizer
import outlines
from typing import Literal
from transformers import AutoTokenizer, AutoModelForCausalLM


# ──────────────────────────────────────────────────────────────────────────────
# Available models
# ──────────────────────────────────────────────────────────────────────────────

AVAILABLE_MODELS = {
    "ChessQween": "zhoudoe23/ChessQween-124m",
    "ChessQween1.5-nano": "zhoudoe23/ChessQween1.5-nano",
    "ChessQween2-tiny": "zhoudoe23/ChessQween2-tiny",
    "ChessQween3-base": "zhoudoe23/ChessQween3-base",

}

MODEL_DESCRIPTIONS = {
    "ChessQween": "Trained on 200,000 best games on Lichess, at least 1900 elo",
    "ChessQween1.5-nano": "A smaller version, but trained on 400,000 master games, at least 2300 elo",
    "ChessQween2-tiny": "New model type with new tokenizer",
    "ChessQween3-base": "Trained with Stockfish's best moves, using FEN board",
}
MODEL_TYPE_MAP = {
    "ChessQween": 1,
    "ChessQween1.5-nano": 1,
    "ChessQween2-tiny": 2,
    "ChessQween3-base": 3,
}


device = torch.device("cpu")

# Lazy cache: {model_id: (tokenizer, model)}
_model_cache: dict = {}

@spaces.GPU(duration=30)
def foo(bar):
    return bar

current_model_type = 1
def load_model(model_key: str):
    """Load (or retrieve from cache) tokenizer + model for the given key."""
    model_id = AVAILABLE_MODELS[model_key]
    if model_id not in _model_cache:
        print(f"Loading {model_id} …")
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        tokenizer.pad_token = tokenizer.eos_token
        
        raw_model = AutoModelForCausalLM.from_pretrained(model_id)
        raw_model.to(device)
        raw_model.eval()
        raw_model.config.use_cache = True
        
        # ─────────── 新增：用 Outlines 包装模型 ───────────
        # outlines 会封装原始模型和 tokenizer，接管 token 级别的掩码生成
        # outlines_model = outlines.from_transformers(raw_model, tokenizer)
        
        # 将 outlines_model 存入缓存（tokenizer 依然保留，方便后续拼接 prompt 选用）
        # _model_cache[model_id] = (tokenizer, raw_model, outlines_model)
        _model_cache[model_id] = (tokenizer, raw_model)

        print(f"✓ {model_id} ready on {device}")
    return _model_cache[model_id]

# ──────────────────────────────────────────────────────────────────────────────
# Chess / model logic
# ──────────────────────────────────────────────────────────────────────────────
def get_history_uci(board: chess.Board) -> str:
    rounds = []
    current_round = []

    for move in board.move_stack:
        current_round.append(move.uci())
        if len(current_round) == 2:
            rounds.append(" ".join(current_round))
            current_round = []

    if current_round:  # 当前轮只有白方走了一步
        rounds.append(" ".join(current_round))

    if not rounds:
        return ""

    history_str = " | ".join(rounds)

    # 关键逻辑：如果是偶数步（轮到白方走棋），且不是开局，末尾必须补上 " |"
    if len(board.move_stack) % 2 == 0 and len(board.move_stack) > 0:
        history_str += " |"

    return " " + history_str

def board_to_prompt(board: chess.Board, model_type) -> str:
    if model_type == 1:
        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().replace("\n"," ")
        
        full_move = board.fullmove_number
        
        # 修复后的逻辑：
        if board.turn == chess.WHITE:
            # 白方回合：拼接回合数和点，例如 " 2."
            prompt = f"Result: 1-0 | {pgn} {full_move}."
        else:
            # 黑方回合：如果你的训练集黑方前面只有空格，直接在白方走法后加空格！
            # 绝不能加孤零零的数字 {full_move}
            prompt = f"Result: 0-1 | {pgn}"
            
        return prompt
    elif model_type == 2:
        if board.turn == chess.WHITE:
            return f"<|WHITE|>{get_history_uci(board)}"
        else:
            return f"<|BLACK|>{get_history_uci(board)}"
    elif model_type == 3:
        fen_list = board.fen().split()
        board_content = " ".join(list(fen_list[0]))
        side = "<|WHITE|>" if fen_list[1] == 'w' else "<|BLACK|>"
        castling = fen_list[2]
        en_passant = fen_list[3]
        return f"{board_content} {side} {castling} {en_passant} 0 1 >"


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())
            if move in board.legal_moves:
                return move
        except Exception:
            pass
    return None

@torch.no_grad()
def get_model_move(board: chess.Board, model_key: str):
    tokenizer, model = load_model(model_key)
    model_type = MODEL_TYPE_MAP[model_key]
    prompt  = board_to_prompt(board, model_type)
    print("Current prompt: "+prompt)
    inputs  = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(
        inputs.input_ids,
        max_new_tokens=12,
        do_sample=False,
        # do_sample=True,
        # temperature=0.3,
        # top_k=40,
        # top_p=0.9,
        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)
    if move:
        return move, True
    print(f"Wrong move: {generated}")
    return random.choice(list(board.legal_moves)), False
# @torch.no_grad()
# def get_model_move(board: chess.Board, model_key: str):
#     # 👈 取出 raw_model 和 outlines_model
#     tokenizer, raw_model, outlines_model = load_model(model_key)
#     model_type = MODEL_TYPE_MAP[model_key]
#     prompt = board_to_prompt(board, model_type)
#     print("Current prompt:", prompt)

#     # ------------------ 旧模型 (Model Type 1) ------------------
#     if model_type == 1:
#         legal_sans = [board.san(m) for m in board.legal_moves]
#         choices = [" " + san for san in legal_sans]
#         MoveType = Literal.__getitem__(tuple(choices))

#         # 使用 outlines_model
#         generator = outlines.Generator(outlines_model, MoveType)
#         generated = generator(
#             prompt, temperature=0.3, top_k=40, top_p=0.9
#         ).strip()

#         move = extract_move(generated, board)
#         if move:
#             return move, True

#     # ------------------ 新模型 (Model Type 2) ------------------
#     elif model_type == 2:
#         legal_moves = list(board.legal_moves)
#         legal_ucis = [m.uci() for m in legal_moves]

#         # 1. 编码 Prompt
#         input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)

#         # 👈 核心修改：使用原生 raw_model 计算 Tensor 输入
#         outputs = raw_model(input_ids)
#         next_token_logits = outputs.logits[0, -1, :].clone()

#         # 2. 🚀 重复走法惩罚 (防死循环摆烂)
#         for move in legal_moves:
#             uci_str = move.uci()
#             if uci_str in tokenizer.vocab:
#                 token_id = tokenizer.vocab[uci_str]
#                 board.push(move)
#     return random.choice(list(board.legal_moves)), False

# ──────────────────────────────────────────────────────────────────────────────
# 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 game_status(board: chess.Board, player_color: str):
    empty_board = chess.Board()
    print(f"Current PGN: \n{empty_board.variation_san(board.move_stack)}")
    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 — AI 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 "AI is thinking…"
    return whose, "playing"

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

def update_model_description(model_key: str):
    desc = MODEL_DESCRIPTIONS.get(model_key, "")
    hf_id = AVAILABLE_MODELS.get(model_key, "")
    return (
        f"<div id='model-desc'>"
        f"<span style='color:#8a7a5a;font-style:italic;'>{desc}</span><br>"
        f"<a href='https://huggingface.co/{hf_id}' target='_blank' "
        f"style='color:#8a6a30;font-size:0.85em;text-decoration:none;'>"
        f"🤗 {hf_id}</a>"
        f"</div>"
    )


def new_game(player_color_choice: str, model_key: str):
    """Reset the board and, if player chose Black, let the model move first."""
    board        = chess.Board()
    player_color = "white" if player_color_choice == "⬜ White (move first)" else "black"
    flipped      = (player_color == "black")
    last_move    = None
    log_lines    = []

    # If player chose Black, model plays White first
    if player_color == "black":
        move, legal = get_model_move(board, model_key)
        san         = board.san(move)
        board.push(move)
        last_move   = move
        log_lines.append(f"{model_key} opens with **{san}**")

    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)
    log_html       = "<br>".join(log_lines) if log_lines else "<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,
        "model_key":     model_key,
    }

    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):
    """Apply the player's chosen move, then let the model respond."""
    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"]
    model_key    = state.get("model_key", "ChessSLM")
    flipped      = (player_color == "black")
    log_lines    = []

    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 **{move_san}**")
    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:
        model_move_obj, legal = get_model_move(board, model_key)
        model_san = board.san(model_move_obj)
        board.push(model_move_obj)
        last_move = model_move_obj
        flag = "" if legal else " *(random fallback)*"
        log_lines.append(f"{model_key} plays **{model_san}**{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,
        "model_key":     model_key,
    }

    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]
    )

    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;
}

/* ── Model selector card ── */
#model-selector-card {
    background: linear-gradient(135deg, #161005 0%, #1e1810 100%);
    border: 1px solid #5a4020;
    border-radius: 8px;
    padding: 1rem 1.2rem;
    margin-bottom: 0.6rem;
    box-shadow: inset 0 1px 0 rgba(212,168,67,0.08);
}

#model-desc {
    margin-top: 0.4rem;
    font-family: 'Crimson Text', serif;
    font-size: 0.92em;
    line-height: 1.5;
    color: #8a7a5a;
}

#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
# ──────────────────────────────────────────────────────────────────────────────

SECTION = "<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>"
LABEL   = lambda t: f"<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>{t}</h3>"

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

    state = gr.State({})

    # ── Header ────────────────────────────────────────────────────────────────
    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;
        ">♛ ChessQween</h1>
        <p style="
            font-family:'Crimson Text',serif; color:#8a7a5a;
            font-size:1.1em; font-style:italic; margin:0;
        ">Play against a GPT-2 model trained on 200,000 chess games</p>
    </div>
    """)

    # ── Main layout ───────────────────────────────────────────────────────────
    with gr.Row():

        # Left: board
        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>"
            )

        # Right: controls
        with gr.Column(scale=2):

            # ── Opponent selector ──────────────────────────────────────────
            gr.HTML(LABEL("CHOOSE OPPONENT"))
            gr.HTML("<div id='model-selector-card'>")

            model_dropdown = gr.Dropdown(
                choices=list(AVAILABLE_MODELS.keys()),
                value="ChessQween",
                label="Opponent model",
                interactive=True,
            )

            model_desc_display = gr.HTML(
                value=update_model_description("ChessQween"),
            )

            gr.HTML("</div>")  # close card

            gr.HTML(SECTION)

            # ── New game ───────────────────────────────────────────────────
            gr.HTML(LABEL("NEW GAME"))

            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(SECTION)

            # ── Your move ─────────────────────────────────────────────────
            gr.HTML(LABEL("YOUR MOVE"))

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

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

            gr.HTML(SECTION)

            # ── Move log ──────────────────────────────────────────────────
            gr.HTML(LABEL("MOVE LOG"))

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

            gr.HTML(SECTION)

            # ── Game history ──────────────────────────────────────────────
            gr.HTML(LABEL("GAME HISTORY"))

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

    # ── Footer ────────────────────────────────────────────────────────────────
    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 by
        <a href="https://huggingface.co/zhoudoe23" target="_blank"
            style="color:#8a6a30;text-decoration:none;">FlameF0X</a>
        &nbsp;·&nbsp; GPT-2 pre-trained on PGN games
        &nbsp;·&nbsp; Move selection uses top-k sampling (temp=0.3)
    </div>
    """)

    # ── Wiring ────────────────────────────────────────────────────────────────

    # Live-update description when model changes
    model_dropdown.change(
        fn=update_model_description,
        inputs=[model_dropdown],
        outputs=[model_desc_display],
    )

    new_game_btn.click(
        fn=new_game,
        inputs=[color_choice, model_dropdown],
        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()
    demo.launch(mcp_server=True)