from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import chess
from datasets import load_dataset
import time
from chess_search.sparse.bitmap import BitmapIndex
from chess_search.utils import position_to_tokens, replay_moves

dset = load_dataset("Lichess/chess-puzzles", split="train")
puzzle_lookup = {row["PuzzleId"]: i for i, row in enumerate(dset)}
idx = BitmapIndex.load("index")

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

@app.get("/")
def read_root(request: Request):
    return templates.TemplateResponse(request=request, name="index.html")

@app.post("/search")
async def search(data: dict):
    start = time.time()
    board = chess.Board(data["fen"])
    tokens = position_to_tokens(board)
    matches = idx.query(tokens)
    resolved = idx.resolve(matches)

    seen_puzzles = {}
    for puzzle_id, move_idx in resolved:
        if puzzle_id not in seen_puzzles:
            seen_puzzles[puzzle_id] = move_idx

    results = []
    for puzzle_id, move_idx in seen_puzzles.items():
        row = dset[puzzle_lookup[puzzle_id]]
        boards = replay_moves(row["FEN"], row["Moves"].split())
        matched_board = boards[move_idx]
        results.append({
            "PuzzleId": row["PuzzleId"],
            "FEN": matched_board.fen(),
            "Moves": row["Moves"],
            "Rating": row["Rating"],
            "Popularity": row["Popularity"],
            "Themes": row["Themes"],
            "MatchedMove": move_idx,
        })

    elapsed_ms = (time.time() - start) * 1000
    return {"count": len(results), "results": results, "time_ms": elapsed_ms}