from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import chess
import chess.pgn
import time
from chess_search.sparse.bitmap import BitmapIndex
from chess_search.sparse.bitmap.parser import parse
from chess_search.sparse.bitmap.expressions import execute

idx = BitmapIndex.load_from_hub("christopher/chess-index")

games = {}
pgn = open("games.pgn")
i = 0
while (game := chess.pgn.read_game(pgn)):
    gid = f"{game.headers.get('White', '?')} vs {game.headers.get('Black', '?')} {game.headers.get('Date', '')}"
    board = game.board()
    boards = []
    for move in game.mainline_moves():
        board.push(move)
        boards.append(board.copy())
    games[gid] = {"headers": dict(game.headers), "boards": boards}
    i += 1
    
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

from pyroaring import BitMap
from chess_search.sparse.bitmap.query import (
    on_file, on_rank, in_center, on_kingside, on_queenside,
    empty_square, open_file, anywhere,
)

@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 = [f"{p.symbol()}_{chess.square_name(sq)}" for sq, p in board.piece_map().items()]
    matches = idx.query(tokens)
    resolved = idx.resolve(matches)

    seen_games = {}
    for game_id, move_idx in resolved:
        if game_id not in seen_games:
            seen_games[game_id] = move_idx

    results = []
    for game_id, move_idx in seen_games.items():
        g = games.get(game_id)
        if not g or move_idx >= len(g["boards"]):
            continue
        results.append({
            "FEN": g["boards"][move_idx].fen(),
            "White": g["headers"].get("White", ""),
            "Black": g["headers"].get("Black", ""),
            "Event": g["headers"].get("Event", ""),
            "Date": g["headers"].get("Date", ""),
            "Result": g["headers"].get("Result", ""),
            "MatchedMove": move_idx,
        })

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



@app.post("/text-search")
async def text_search(data: dict):
    start = time.time()
    q = parse(data["query"])
    matches = execute(q, idx)
    resolved = idx.resolve(matches)

    seen_games = {}
    for game_id, move_idx in resolved:
        if game_id not in seen_games:
            seen_games[game_id] = move_idx

    results = []
    for game_id, move_idx in seen_games.items():
        g = games.get(game_id)
        if not g or move_idx >= len(g["boards"]):
            continue
        results.append({
            "FEN": g["boards"][move_idx].fen(),
            "White": g["headers"].get("White", ""),
            "Black": g["headers"].get("Black", ""),
            "Event": g["headers"].get("Event", ""),
            "Date": g["headers"].get("Date", ""),
            "Result": g["headers"].get("Result", ""),
            "MatchedMove": move_idx,
        })

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