# -*- coding: utf-8 -*-
"""Open Quantum Challenge (OQC) — submission intake + live leaderboard.

Thin by design, exactly like the Open Discovery Challenge it mirrors. Nothing here scores.
Entries land in a private dataset ledger (FINAL-Bench/oqc-ledger); a worker on the GPU box
polls that ledger, scores each entry against a frozen private answer key, and writes the
result and the rolled-up leaderboard back. One entrant can never block the service, and the
yardstick stays private while the harness that produces entries is fully public.

Layout in the ledger (shared with the worker):
    submissions/<board>/<id>.json   written here when an entry is accepted
    results/<board>/<id>.json       written by the worker
    leaderboard/<board>.json        rolled up by the worker; the page reads one file

Boards: qsim (Season 1), qec (Season 2).
"""
import base64
import hashlib
import hmac
import json
import os
import secrets
import time
import urllib.parse
import urllib.request

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
from pydantic import BaseModel

import gates
import seasons
import store

HERE = os.path.dirname(os.path.abspath(__file__))
CACHE = store.Cached(ttl=int(os.environ.get("OQC_CACHE_TTL", "30")))

OAUTH_ID = os.environ.get("OAUTH_CLIENT_ID", "")
OAUTH_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "")
OAUTH_ISS = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co")
SPACE_HOST = os.environ.get("SPACE_HOST", "")
COOKIE = "oqc_session"
IN_FRAME = bool(SPACE_HOST)
COOKIE_KW = ({"samesite": "none", "secure": True} if IN_FRAME
             else {"samesite": "lax", "secure": False})
SESSION_KEY = os.environ.get("OQC_SESSION_KEY") or secrets.token_hex(16)
DAILY_CAP = int(os.environ.get("OQC_DAILY_CAP", "50"))
CAP_TZ_OFFSET = 9 * 3600

app = FastAPI(title="Open Quantum Challenge")


class Submission(BaseModel):
    board: str                       # "qsim" or "qec"
    display_name: str
    method: str = ""
    predictions: dict = {}


# ----------------------------------------------------------------- session (HF OAuth, optional)
def _sign(payload):
    raw = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
    sig = hmac.new(SESSION_KEY.encode(), raw.encode(), hashlib.sha256).hexdigest()[:32]
    return raw + "." + sig


def _unsign(token):
    try:
        raw, sig = (token or "").rsplit(".", 1)
    except ValueError:
        return None
    good = hmac.new(SESSION_KEY.encode(), raw.encode(), hashlib.sha256).hexdigest()[:32]
    if not hmac.compare_digest(sig, good):
        return None
    pad = "=" * (-len(raw) % 4)
    try:
        return json.loads(base64.urlsafe_b64decode(raw + pad))
    except Exception:
        return None


def current_user(request):
    return _unsign(request.cookies.get(COOKIE))


def _redirect_uri(request):
    if SPACE_HOST:
        return "https://%s/auth/callback" % SPACE_HOST
    return str(request.base_url).rstrip("/") + "/auth/callback"


@app.get("/")
async def index():
    return FileResponse(os.path.join(HERE, "index.html"))


@app.get("/api/me")
def me(request: Request):
    u = current_user(request)
    return {"signed_in": bool(u), "user": u,
            "oauth_configured": bool(OAUTH_ID and OAUTH_SECRET)}


@app.get("/login")
def login(request: Request):
    if not (OAUTH_ID and OAUTH_SECRET):
        raise HTTPException(503, "OAuth is not configured (add hf_oauth in README)")
    state = secrets.token_urlsafe(16)
    q = urllib.parse.urlencode({
        "client_id": OAUTH_ID, "redirect_uri": _redirect_uri(request),
        "response_type": "code", "scope": "openid profile", "state": state})
    r = RedirectResponse("%s/oauth/authorize?%s" % (OAUTH_ISS.rstrip("/"), q))
    r.set_cookie("oqc_state", state, max_age=600, httponly=True, **COOKIE_KW)
    return r


@app.get("/auth/callback")
def callback(request: Request, code: str = "", state: str = ""):
    saved = request.cookies.get("oqc_state")
    if not code or not saved or state != saved:
        raise HTTPException(400, "sign-in failed — please try again (third-party cookies may be blocked)")
    body = urllib.parse.urlencode({
        "client_id": OAUTH_ID, "client_secret": OAUTH_SECRET,
        "grant_type": "authorization_code", "code": code,
        "redirect_uri": _redirect_uri(request)}).encode()
    with urllib.request.urlopen(urllib.request.Request(
            OAUTH_ISS.rstrip("/") + "/oauth/token", data=body,
            headers={"Content-Type": "application/x-www-form-urlencoded"}), timeout=20) as r:
        tok = json.loads(r.read().decode())
    with urllib.request.urlopen(urllib.request.Request(
            OAUTH_ISS.rstrip("/") + "/oauth/userinfo",
            headers={"Authorization": "Bearer " + tok["access_token"]}), timeout=20) as r:
        ui = json.loads(r.read().decode())
    user = {"name": ui.get("preferred_username") or ui.get("name"),
            "picture": ui.get("picture"), "at": int(time.time())}
    resp = RedirectResponse("/")
    resp.set_cookie(COOKIE, _sign(user), max_age=60 * 60 * 24 * 14, httponly=True, **COOKIE_KW)
    resp.delete_cookie("oqc_state")
    return resp


@app.get("/logout")
def logout():
    r = RedirectResponse("/")
    r.delete_cookie(COOKIE)
    return r


# ----------------------------------------------------------------- seasons / leaderboard
@app.get("/api/seasons")
def season_list():
    return {"seasons": seasons.listing(), "default": seasons.DEFAULT}


def _board_json(board):
    return CACHE.get("board:%s" % board,
                     lambda: store.read("leaderboard/%s.json" % board, default={}) or {})


def _sub_ids(board):
    return CACHE.get("ids:%s" % board, lambda: store.listdir("submissions/%s" % board))


@app.get("/api/leaderboard")
def leaderboard(request: Request):
    board = (request.query_params.get("board") or "qsim").lower()
    s = seasons.by_board(board)
    if not s:
        raise HTTPException(404, "unknown board")
    data = _board_json(board)
    entries = data.get("entries") or []
    anchors, ranked = [], []
    for e in entries:
        u = e.get("user", "") or ""
        row = {"participant": u, "method": e.get("method", ""),
               "score": e.get("score"), "verified": bool(e.get("verified"))}
        if u.startswith(("baseline-", "reference-", "anchor-")):
            row["anchor"] = True
            anchors.append(row)
        else:
            ranked.append(row)
    # reference anchors mark the ladder; they do not climb it (ODC-style)
    if board == "qsim":
        anchors.insert(0, {"participant": "exact diagonalization", "anchor": True,
                           "method": "theoretical floor (reference)", "score": 0.0, "verified": True})
    elif board == "qec":
        anchors.append({"participant": "near-optimal MLE bound", "anchor": True,
                        "method": "information-theoretic reference (not a submission)", "score": None, "verified": True})
    ranked.sort(key=lambda r: (r["score"] is None, r["score"] if r["score"] is not None else 1e18))
    for i, r in enumerate(ranked, 1):
        r["rank"] = i
    return {"season": seasons.public(s), "anchors": anchors, "entries": ranked,
            "updated": data.get("updated"), "n": len(ranked)}


@app.get("/api/queue")
def queue(request: Request):
    board = (request.query_params.get("board") or "qsim").lower()
    ids = _sub_ids(board)
    scored = {e.get("user") for e in (_board_json(board).get("entries") or [])}
    return {"submitted": len(ids), "ranked": len(scored)}


@app.get("/api/totals")
def totals():
    subs = 0
    users = set()
    for n in seasons.SEASONS:
        b = seasons.SEASONS[n]["board"]
        for e in (_board_json(b).get("entries") or []):
            subs += 1
            if e.get("user"):
                users.add(e["user"])
    return {"ranked": subs, "participants": len(users)}


# ----------------------------------------------------------------- submit
@app.post("/api/submit")
def submit(sub: Submission, request: Request):
    board = (sub.board or "").lower()
    s = seasons.by_board(board)
    if not s:
        raise HTTPException(400, "unknown board %r" % sub.board)
    if not seasons.accepting(s):
        if s.get("open"):
            raise HTTPException(403, "Season %d is closed (deadline %s)." % (s["number"], s["closes"]))
        raise HTTPException(403, "Season %d is not open yet (opens %s)." % (s["number"], s["opens"]))
    user = current_user(request)
    if (OAUTH_ID and OAUTH_SECRET) and not user:
        raise HTTPException(401, "Hugging Face sign-in required")
    who = (user or {}).get("name") or (sub.display_name or "").strip()[:60]
    if not who:
        raise HTTPException(400, "display name required")

    v = gates.check(board, sub.predictions)
    if not v["admitted"]:
        return JSONResponse({"accepted": False, "reasons": v["reject"]}, status_code=422)

    # wrong-track guard: predictions whose keys are clearly for the other board
    keys = list(sub.predictions.keys())
    if keys:
        qec_like = sum(k.startswith("shot_") for k in keys) > len(keys) // 2
        qsim_like = sum(k.startswith("hub_") for k in keys) > len(keys) // 2
        if board == "qsim" and qec_like:
            return JSONResponse({"accepted": False, "reasons": [
                "This looks like a QEC submission (shot_* keys). Please select Season 2 - QEC Decoder."]}, status_code=422)
        if board == "qec" and qsim_like:
            return JSONResponse({"accepted": False, "reasons": [
                "This looks like a QSIM submission (hub_* keys). Please select Season 1 - Quantum Simulation."]}, status_code=422)

    # daily cap via a small tally file (not a tree walk)
    day = time.strftime("%Y-%m-%d", time.gmtime(time.time() + CAP_TZ_OFFSET))
    quota_path = "quota/%s/%s.json" % (board, day)
    quota = store.read(quota_path, default={}) or {}
    if DAILY_CAP > 0 and int(quota.get(who, 0)) >= DAILY_CAP:
        return JSONResponse({"accepted": False, "daily_cap": DAILY_CAP,
                             "reasons": ["daily submission cap reached (%d). Resets at KST midnight." % DAILY_CAP]},
                            status_code=429)

    safe = "".join(c for c in who if c.isalnum() or c in "-_")[:40] or "anon"
    sid = "%s_%s_%s" % (time.strftime("%Y%m%d_%H%M%S"), safe, secrets.token_hex(2))
    rec = {"hf_user": who, "display_name": (sub.display_name or who)[:60],
           "method": (sub.method or "").strip()[:200], "predictions": sub.predictions,
           "board": board, "season": s["number"], "spec": "oqc-v1",
           "n_items": v["n_items"], "submitted_at": int(time.time())}
    quota[who] = int(quota.get(who, 0)) + 1
    try:
        store.write_many([("submissions/%s/%s.json" % (board, sid), rec),
                          (quota_path, quota)],
                         summary="oqc %s entry %s" % (board, sid))
    except store.Busy:
        return JSONResponse({"accepted": False, "retryable": True,
                             "reasons": ["the ledger was busy — please submit again in a moment"]},
                            status_code=503)
    CACHE.drop()
    ids = _sub_ids(board)
    scored = {e.get("user") for e in (_board_json(board).get("entries") or [])}
    ahead = len([i for i in ids if i not in scored])
    return {"accepted": True, "submission_id": sid, "queue_position": ahead,
            "note": "scored by a GPU worker against the private key; the board updates within a few minutes."}