space_app/main.py

from __future__ import annotations

import os
import random
from pathlib import Path
from typing import Any

from fastapi import Depends, FastAPI, Header, HTTPException, Query
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from .board_utils import normalize_board_for_display, normalize_board_for_submission
from .config import Settings
from .dataset import DatasetStore
from .db import SessionStore
from .llm_results import LLMResultsStore
from .models import engine_for_puzzle, normalize_player_name
from .schemas import (
    CreateSessionRequest,
    LeaderboardEntry,
    LeaderboardResponse,
    LLMResultsResponse,
    PuzzleOptionResponse,
    SessionPayload,
    SessionResponse,
    SubmitResponse,
    SubmitSessionRequest,
)
from .verification import VerificationService


def _session_to_response(
    session: dict[str, Any],
    *,
    problem_ascii: str,
    image_base64: str | None,
) -> SessionResponse:
    current_board_ascii = normalize_board_for_display(
        puzzle_type=str(session["puzzle_type"]),
        problem_ascii=problem_ascii,
        board_ascii=str(session["submitted_artifact"] or problem_ascii),
    )
    return SessionResponse(
        session_id=str(session["id"]),
        engine=str(session["engine"]),
        puzzle_type=str(session["puzzle_type"]),
        difficulty=str(session["difficulty"]),
        puzzle_id=str(session["puzzle_filename"]),
        args=str(session["args"]),
        status=str(session["status"]),
        started_at=session["started_at"],
        payload=SessionPayload(
            problem_ascii=problem_ascii,
            current_board_ascii=current_board_ascii,
            image_base64=image_base64,
        ),
    )


def create_app(
    *,
    settings: Settings | None = None,
    dataset_store: DatasetStore | None = None,
    session_store: SessionStore | None = None,
    verifier: VerificationService | None = None,
    llm_results_store: LLMResultsStore | None = None,
) -> FastAPI:
    resolved_settings = settings or Settings.from_env()
    resolved_settings.prepare()
    app = FastAPI(title="TopoBench Space")
    app.state.settings = resolved_settings
    app.state.dataset_store = dataset_store or DatasetStore.load_huggingface()
    app.state.session_store = session_store or SessionStore(resolved_settings.database_path)
    app.state.verifier = verifier or VerificationService()
    app.state.llm_results_store = llm_results_store or LLMResultsStore.from_file(
        Path(__file__).parent / "llm_results.json"
    )

    def get_settings() -> Settings:
        return app.state.settings

    def get_dataset_store() -> DatasetStore:
        return app.state.dataset_store

    def get_session_store() -> SessionStore:
        return app.state.session_store

    def get_verifier() -> VerificationService:
        return app.state.verifier

    def get_llm_results_store() -> LLMResultsStore:
        return app.state.llm_results_store

    def require_admin(
        authorization: str | None = Header(default=None),
        current_settings: Settings = Depends(get_settings),
    ) -> None:
        token = f"Bearer {current_settings.admin_api_token}"
        if authorization != token:
            raise HTTPException(status_code=401, detail="Unauthorized")

    @app.get("/api/health")
    def health() -> dict[str, str]:
        return {"status": "ok"}

    @app.get("/api/puzzles", response_model=list[PuzzleOptionResponse])
    def list_puzzles(
        puzzle_type: str = Query(...),
        difficulty: str = Query(...),
        current_dataset_store: DatasetStore = Depends(get_dataset_store),
    ) -> list[PuzzleOptionResponse]:
        try:
            rows = current_dataset_store.list_rows(
                puzzle_type=puzzle_type,
                difficulty=difficulty,
                limit=50,
            )
        except ValueError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        return [
            PuzzleOptionResponse(
                puzzle_id=row.filename,
                title=row.filename,
                sequence=index,
                puzzle_type=row.puzzlename,
                difficulty=row.difficulty,
                args=row.args,
            )
            for index, row in enumerate(rows, start=1)
        ]

    @app.post("/api/sessions", response_model=SessionResponse)
    def create_session(
        request: CreateSessionRequest,
        current_dataset_store: DatasetStore = Depends(get_dataset_store),
        current_session_store: SessionStore = Depends(get_session_store),
    ) -> SessionResponse:
        player_name_raw = request.player_name.strip()
        player_name_norm = normalize_player_name(player_name_raw)
        if not player_name_norm:
            raise HTTPException(status_code=400, detail="Player name is required.")
        if request.puzzle_id is None:
            try:
                pool = current_dataset_store.list_rows(
                    puzzle_type=request.puzzle_type,
                    difficulty=request.difficulty,
                    limit=50,
                )
            except ValueError as exc:
                raise HTTPException(status_code=404, detail=str(exc)) from exc
            already_solved = current_session_store.get_solved_filenames(
                player_name_norm=player_name_norm,
                puzzle_type=request.puzzle_type,
                difficulty=request.difficulty,
            )
            unseen = [item for item in pool if item.filename not in already_solved]
            row = random.choice(unseen) if unseen else random.choice(pool)
        else:
            try:
                row = current_dataset_store.get_row(request.puzzle_id)
            except KeyError:
                raise HTTPException(status_code=404, detail="Puzzle not found.") from None
            if row.puzzlename != request.puzzle_type or row.difficulty != request.difficulty:
                raise HTTPException(
                    status_code=400,
                    detail="Selected puzzle does not match the chosen family and difficulty.",
                )
        session = current_session_store.create_session(
            player_name_raw=player_name_raw,
            player_name_norm=player_name_norm,
            puzzle_type=request.puzzle_type,
            difficulty=request.difficulty,
            dataset_variant="plain",
            puzzle_filename=row.filename,
            args=row.args,
            engine=engine_for_puzzle(request.puzzle_type),
        )
        return _session_to_response(
            session,
            problem_ascii=row.problem,
            image_base64=row.image_base64,
        )

    @app.get("/api/sessions/{session_id}", response_model=SessionResponse)
    def get_session(
        session_id: str,
        current_dataset_store: DatasetStore = Depends(get_dataset_store),
        current_session_store: SessionStore = Depends(get_session_store),
    ) -> SessionResponse:
        session = current_session_store.get_session(session_id)
        if session is None:
            raise HTTPException(status_code=404, detail="Session not found.")
        row = current_dataset_store.get_row(str(session["puzzle_filename"]))
        return _session_to_response(
            session,
            problem_ascii=row.problem,
            image_base64=row.image_base64,
        )

    @app.post("/api/sessions/{session_id}/ready")
    def ready_session(
        session_id: str,
        current_session_store: SessionStore = Depends(get_session_store),
    ) -> dict[str, str | None]:
        try:
            session = current_session_store.mark_ready(session_id)
        except KeyError:
            raise HTTPException(status_code=404, detail="Session not found.") from None
        return {"status": str(session["status"]), "started_at": session["started_at"]}

    @app.post("/api/sessions/{session_id}/submit", response_model=SubmitResponse)
    def submit_session(
        session_id: str,
        request: SubmitSessionRequest,
        current_dataset_store: DatasetStore = Depends(get_dataset_store),
        current_session_store: SessionStore = Depends(get_session_store),
        current_verifier: VerificationService = Depends(get_verifier),
    ) -> SubmitResponse:
        session = current_session_store.get_session(session_id)
        if session is None:
            raise HTTPException(status_code=404, detail="Session not found.")
        if not session["started_at"]:
            raise HTTPException(status_code=409, detail="Session is not ready yet.")
        row = current_dataset_store.get_row(str(session["puzzle_filename"]))
        # Normalization and verification must never crash the request: a failure
        # here (e.g. a native verifier segfault/hang, or a malformed board) is
        # recorded as a failed attempt rather than surfacing as an opaque 500.
        try:
            normalized_board_ascii = normalize_board_for_submission(
                puzzle_type=str(session["puzzle_type"]),
                problem_ascii=row.problem,
                board_ascii=request.board_ascii,
            )
        except Exception:
            normalized_board_ascii = request.board_ascii
        try:
            verification = current_verifier.verify(
                puzzle_type=str(session["puzzle_type"]),
                problem_ascii=row.problem,
                board_ascii=normalized_board_ascii,
                args=str(session["args"]),
            )
        except Exception as exc:  # pragma: no cover - defensive guard
            verification = {
                "board_exists": bool(normalized_board_ascii),
                "board_valid": False,
                "board_modified": False,
                "correct": False,
                "error": f"Verification failed: {exc}"[:500],
            }
        updated = current_session_store.record_submission(
            session_id=session_id,
            solved=bool(verification.get("correct")),
            submitted_artifact=normalized_board_ascii,
            verification_payload=verification,
        )
        return SubmitResponse(
            solved=bool(updated["solved"]),
            elapsed_ms=updated["elapsed_ms"],
            status=str(updated["status"]),
            verification=verification,
        )

    @app.get("/api/llm-results/{puzzle_id}", response_model=LLMResultsResponse)
    def llm_results(
        puzzle_id: str,
        current_llm_results_store: LLMResultsStore = Depends(get_llm_results_store),
    ) -> LLMResultsResponse:
        verdicts = current_llm_results_store.get_for_puzzle(puzzle_id)
        if verdicts is None:
            raise HTTPException(
                status_code=404,
                detail="No LLM results recorded for this puzzle.",
            )
        return LLMResultsResponse(
            puzzle_id=puzzle_id,
            models_solved=verdicts["models_solved"],
            models_failed=verdicts["models_failed"],
        )

    @app.get("/api/admin/check", dependencies=[Depends(require_admin)])
    def admin_check() -> dict[str, bool]:
        return {"ok": True}

    @app.get("/api/admin/leaderboard/{puzzle_id}", response_model=LeaderboardResponse, dependencies=[Depends(require_admin)])
    def admin_leaderboard(
        puzzle_id: str,
        include_test: bool = Query(default=False),
        current_session_store: SessionStore = Depends(get_session_store),
    ) -> LeaderboardResponse:
        rows = current_session_store.list_leaderboard(
            puzzle_id=puzzle_id,
            include_test=include_test,
        )
        entries = [
            LeaderboardEntry(
                player_name=str(row["player_name_raw"]),
                elapsed_ms=row["elapsed_ms"],
                submission_count=int(row["submission_count"]),
                solved=bool(row["solved"]),
                started_at=row["started_at"],
                submitted_at=row["submitted_at"],
            )
            for row in rows
        ]
        return LeaderboardResponse(puzzle_id=puzzle_id, entries=entries)

    @app.get("/api/admin/solves", dependencies=[Depends(require_admin)])
    def admin_solves(
        player_name_norm: str | None = Query(default=None),
        puzzle_type: str | None = Query(default=None),
        difficulty: str | None = Query(default=None),
        status: str | None = Query(default=None),
        date_from: str | None = Query(default=None),
        date_to: str | None = Query(default=None),
        current_session_store: SessionStore = Depends(get_session_store),
    ) -> list[dict[str, Any]]:
        return current_session_store.list_solves(
            {
                "player_name_norm": player_name_norm,
                "puzzle_type": puzzle_type,
                "difficulty": difficulty,
                "status": status,
                "date_from": date_from,
                "date_to": date_to,
            }
        )

    @app.get("/api/admin/aggregates", dependencies=[Depends(require_admin)])
    def admin_aggregates(
        player_name_norm: str | None = Query(default=None),
        puzzle_type: str | None = Query(default=None),
        difficulty: str | None = Query(default=None),
        date_from: str | None = Query(default=None),
        date_to: str | None = Query(default=None),
        current_session_store: SessionStore = Depends(get_session_store),
    ) -> list[dict[str, Any]]:
        aggregates = current_session_store.aggregate_solves(
            {
                "player_name_norm": player_name_norm,
                "puzzle_type": puzzle_type,
                "difficulty": difficulty,
                "date_from": date_from,
                "date_to": date_to,
            }
        )
        return [item.__dict__ for item in aggregates]

    frontend_dist_dir = resolved_settings.frontend_dist_dir
    assets_dir = frontend_dist_dir / "assets"
    if assets_dir.exists():
        app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")

    @app.get("/{path:path}")
    def spa(path: str) -> FileResponse:
        if path.startswith("api/"):
            raise HTTPException(status_code=404, detail="Not found.")
        index_path = frontend_dist_dir / "index.html"
        if index_path.exists():
            return FileResponse(index_path)
        raise HTTPException(
            status_code=503,
            detail="Frontend has not been built yet. Run `npm run build` in frontend/.",
        )

    return app


app = create_app() if os.getenv("TOPOBENCH_SKIP_AUTOAPP") != "1" else FastAPI()