import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
import os

TOKEN = os.environ.get("HF_TOKEN")

model_repo = "typical-cyber/breakthrough-model"
model_path = hf_hub_download(
    model_repo,
    "breakthrough_mcvs.py",
    repo_type="model",
    token=TOKEN
)

ns = {}
with open(model_path, "r", encoding="utf-8-sig") as f:
    exec(f.read(), ns)

db_repo = "typical-cyber/breakthrough-data"
db_path = hf_hub_download(
    db_repo,
    "breakthrough_zone_db.npz",
    repo_type="dataset",
    token=TOKEN
)
zonedb_data = np.load(db_path, allow_pickle=True)

HilbertOrderedZoneDatabase = ns["HilbertOrderedZoneDatabase"]
Breakthrough = ns["Breakthrough"]
MCVSSearcher = ns["MCVSSearcher"]

zonedb = HilbertOrderedZoneDatabase("breakthroughzonedb.npz", maxsize=10000)
zonedb.winningmatrices = list(zonedb_data.get("winning", []))
zonedb.losingmatrices = list(zonedb_data.get("losing", []))
zonedb.drawmatrices = list(zonedb_data.get("draw", []))

def parse_board(board_text):
    rows = [line.strip() for line in board_text.strip().splitlines() if line.strip()]
    if len(rows) != 8:
        raise ValueError("Board must have exactly 8 rows.")
    board = np.zeros((8, 8), dtype=np.int32)
    mapping = {".": 0, "1": 1, "2": 2}
    for r, line in enumerate(rows):
        parts = line.split()
        if len(parts) != 8:
            raise ValueError("Each row must have 8 space-separated cells.")
        for c, cell in enumerate(parts):
            if cell not in mapping:
                raise ValueError("Use '.', '1', or '2' only.")
            board[r, c] = mapping[cell]
    return board

def get_move(board_text, player):
    game = Breakthrough()
    game.board = parse_board(board_text)
    game.movecount = 0
    game.cachedmatrix = None

    searcher = MCVSSearcher(None, None, zonedb, lambda_zone=1.0, k_zone=5)
    visits, _ = searcher.search_with_time_budget(game, 1.0)
    if not visits:
        return "No legal move"
    best_move = max(visits, key=visits.get)
    return str(best_move)

demo = gr.Interface(
    fn=get_move,
    inputs=[
        gr.Textbox(
            label="Board",
            lines=8,
            value=(
                "1 1 1 1 1 1 1 1\n"
                "1 1 1 1 1 1 1 1\n"
                ". . . . . . . .\n"
                ". . . . . . . .\n"
                ". . . . . . . .\n"
                ". . . . . . . .\n"
                "2 2 2 2 2 2 2 2\n"
                "2 2 2 2 2 2 2 2"
            )
        ),
        gr.Dropdown(choices=["1", "2"], value="1", label="Player")
    ],
    outputs=gr.Textbox(label="Best Move"),
    title="Breakthrough MCVS Secure"
)

demo.launch()