import importlib.util
import gradio as gr
import chess
import chess.svg
import torch
import spaces
from huggingface_hub import hf_hub_download
from transformers import AutoModel, AutoConfig

MODEL_ID = "Datdanboi25/Charles-the-Chess-Bot"
REVISION = "63e2574705794b40899f642aaa619d1668e1e731"

config = AutoConfig.from_pretrained(MODEL_ID, revision=REVISION, trust_remote_code=True)
model = AutoModel.from_pretrained(
    MODEL_ID, revision=REVISION, config=config, trust_remote_code=True
).eval()

_TOK_PATH = hf_hub_download(MODEL_ID, "tokenization_chess_policy.py", revision=REVISION)
_SPEC = importlib.util.spec_from_file_location("charles_tokenization", _TOK_PATH)
charles_tokenization = importlib.util.module_from_spec(_SPEC)
_SPEC.loader.exec_module(charles_tokenization)

board = chess.Board()


@spaces.GPU
def get_bot_move_uci(fen: str) -> str | None:
    """สำคัญ: รับ FEN เป็น argument แทนการอ่าน global `board`
    เพราะฟังก์ชันนี้รันใน subprocess แยกของ ZeroGPU
    ต้อง reconstruct board ในนั้นเอง ไม่งั้นจะได้สถานะเก่าค้าง"""
    b = chess.Board(fen)
    legal_moves = list(b.legal_moves)
    if not legal_moves:
        return None

    inputs = charles_tokenization.build_model_inputs(
        [b.copy(stack=False)], include_candidate_uci=True
    )
    ucis = inputs.pop("candidate_uci")[0]

    with torch.inference_mode():
        logits = model(**inputs).logits[0].float()

    moves = [chess.Move.from_uci(u) for u in ucis]
    best_idx = int(torch.argmax(logits))
    return moves[best_idx].uci()


def make_move(user_move):
    global board
    try:
        if board.is_game_over():
            return "เกมจบแล้ว! กรุณากดเริ่มเกมใหม่", chess.svg.board(board)

        move = chess.Move.from_uci(user_move.strip())
        if move not in board.legal_moves:
            check_hint = ""
            if board.is_check():
                check_hint = " ⚠️ ตอนนี้คุณโดนรุกอยู่! ต้องเดินแก้รุกก่อน (ย้ายราชาหนี / เอาตัวบัง / กินตัวที่รุก)"
            return f"ตาเดินไม่ถูกต้องตามกฎ!{check_hint} (ตัวอย่าง: e2e4, g1f3)", chess.svg.board(board)
        board.push(move)

    except ValueError:
        return "รูปแบบตาเดินไม่ถูกต้อง (ต้องเป็นแบบ UCI เช่น e2e4)", chess.svg.board(board)

    if board.is_game_over():
        if board.is_checkmate():
            winner = "คุณ" if board.turn != chess.WHITE else "บอท"
            # board.turn คือฝ่ายที่ต้องเดินแต่เดินไม่ได้ (โดนรุกฆาต) = ฝ่ายแพ้
            loser = "บอท" if board.turn != chess.WHITE else "คุณ"
            return f"♛ รุกฆาต! {loser}แพ้", chess.svg.board(board)
        return "เกมจบแล้ว! เสมอ (Stalemate/Draw)", chess.svg.board(board)

    bot_move_uci = get_bot_move_uci(board.fen())

    if bot_move_uci:
        bot_move = chess.Move.from_uci(bot_move_uci)
        if bot_move in board.legal_moves:
            board.push(bot_move)
            check_note = " (รุก!)" if board.is_check() else ""
            if board.is_checkmate():
                msg = f"บอทเดิน: {bot_move.uci()} — ♛ รุกฆาต! คุณแพ้"
            else:
                msg = f"บอทเดิน: {bot_move.uci()}{check_note}"
        else:
            msg = f"บอทพยายามเดิน {bot_move_uci} แต่ผิดกฎ - ข้ามตานี้"
    else:
        msg = "บอทไม่สามารถเดินหมากได้ - ตาของคุณเดินต่อ"

    return msg, chess.svg.board(board)


def reset_game():
    global board
    board = chess.Board()
    return "เริ่มเกมใหม่เรียบร้อย", chess.svg.board(board)


with gr.Blocks() as demo:
    gr.Markdown("# ♟️ ซ้อมหมากรุกกับ Charles the Chess Bot")
    board_html = gr.HTML(chess.svg.board(board))
    with gr.Row():
        move_input = gr.Textbox(label="ใส่ตาเดินแบบ UCI (เช่น e2e4, g1f3):", placeholder="e2e4")
        submit_btn = gr.Button("เดินหมาก", variant="primary")
        reset_btn = gr.Button("เริ่มเกมใหม่")
    status_text = gr.Textbox(label="สถานะ / ตาเดินของบอท")
    submit_btn.click(make_move, inputs=[move_input], outputs=[status_text, board_html])
    reset_btn.click(reset_game, outputs=[status_text, board_html])

demo.launch()