# -*- coding: utf-8 -*-
"""Open Materials Challenge — 접수와 순위표.

**이 서비스가 하지 않는 일**: 채점. 스페이스에는 Materials Project 자료도 GPU 도 없다.
제출을 받아 원장(비공개 데이터셋)에 적고, 워커가 채점해 넣은 결과를 보여줄 뿐이다.
여기서 물성을 판정하는 척하면 참가자는 근거 없는 반려를 받게 된다.

순위표는 TTL 당 한 번만 만들어 미리 압축해 두고, 캐시는 단일 갱신자로 돌린다.
첫 페이지는 async 로 서빙하고, 표는 500 행씩 나눠 그린다.
"""
import base64
import gzip
import hashlib
import hmac
import json
import os
import re
import secrets
import time
import urllib.parse
import urllib.request
import uuid

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import (FileResponse, JSONResponse, RedirectResponse,
                               Response)
from pydantic import BaseModel

import gates
import seasons
import store

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "data")
# 페이지 폴링 주기(20초)보다 크게 잡는다.
CACHE = store.Cached(ttl=int(os.environ.get("OMC_CACHE_TTL", "60")))
SALT = os.environ.get("OMC_SALT", "omc-season1")

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 = "omc_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("OMC_SESSION_KEY") or secrets.token_hex(16)

DAILY_CAP = int(os.environ.get("OMC_DAILY_CAP", "30"))
CAP_TZ_OFFSET = 9 * 3600          # 참가자가 가정할 하루 경계 (KST)

app = FastAPI(title="Open Materials Challenge")
app.add_middleware(GZipMiddleware, minimum_size=1024)


class Submission(BaseModel):
    formula: str                  # 조성. 예: Li3YCl6, LiZr2(PO4)3
    display_name: str             # 순위표에 표시될 이름
    model_name: str = ""          # 어떤 모델이 제안했는지. 자유 기재
    rationale: str = ""           # 설계 근거 (선택)
    structure: str = ""           # CIF (선택) - 내면 채점이 빠르고 정확해진다
    season: int = 0
    # 조성을 공개할지는 제출자가 정한다. 되돌릴 수 없는 쪽이 기본값이어야 한다 -
    # 한 번 공개된 조성은 되감을 수 없고, 공개는 특허성을 깎는다.
    visibility: str = "private"


# ---------------------------------------------------------------- 원장 접근
def _submissions(s):
    return CACHE.get("ids:%d" % s["number"],
                     lambda: store.listdir(seasons.path(s, "submissions")))


def _board(s):
    return CACHE.get("board:%d" % s["number"],
                     lambda: store.read(seasons.path(s, "leaderboard.json"),
                                        default={}) or {})


def public_id(key):
    return "OMC-" + hashlib.sha256((SALT + key).encode()).hexdigest()[:10].upper()


def mask(rec):
    """비공개 제출에서 조성을 가린다. 원소 종류까지는 보여준다 - 그것만으로는
    조성을 되돌릴 수 없고, 어떤 화학계가 상위인지는 대회의 공개 지식이어야 한다."""
    out = dict(rec)
    if rec.get("visibility") != "public":
        comp = rec.get("composition") or {}
        out["formula"] = "·".join(sorted(comp)) if comp else "비공개"
        out["masked"] = True
        out.pop("structure", None)
        # 🔴 formula 만 가리면 안 된다. composition 은 기계가 읽는 같은 정보이고,
        #    이것을 남겨두면 가림이 화면에서만 되고 API 로는 전부 새어나간다.
        #    (2026-09-07 참가자 제보로 발견 - 비공개 521건의 조성이 무인증으로 노출)
        out.pop("composition", None)
    return out


# ---------------------------------------------------------------- 세션
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: Request):
    return _unsign(request.cookies.get(COOKIE))


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


# ---------------------------------------------------------------- 화면
@app.get("/")
async def index():
    # 정적 페이지는 요청 스레드풀을 타지 않게 async 로 둔다.
    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 가 설정되지 않았습니다.")
    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, q))
    r.set_cookie("omc_state", state, max_age=600, httponly=True, **COOKIE_KW)
    return r


@app.get("/auth/callback")
def callback(request: Request, code: str = "", state: str = ""):
    if not code or state != request.cookies.get("omc_state"):
        return RedirectResponse("/?auth=failed")
    body = urllib.parse.urlencode({
        "client_id": OAUTH_ID, "client_secret": OAUTH_SECRET,
        "grant_type": "authorization_code", "code": code,
        "redirect_uri": _redirect_uri(request)}).encode()
    try:
        req = urllib.request.Request("%s/oauth/token" % OAUTH_ISS, data=body)
        with urllib.request.urlopen(req, timeout=30) as x:
            tok = json.loads(x.read())
        req = urllib.request.Request(
            "%s/oauth/userinfo" % OAUTH_ISS,
            headers={"Authorization": "Bearer " + tok["access_token"]})
        with urllib.request.urlopen(req, timeout=30) as x:
            info = json.loads(x.read())
    except Exception:
        return RedirectResponse("/?auth=failed")
    u = {"sub": info.get("sub"), "name": info.get("preferred_username") or info.get("name"),
         "picture": info.get("picture")}
    r = RedirectResponse("/?auth=ok")
    r.set_cookie(COOKIE, _sign(u), max_age=30 * 86400, httponly=True, **COOKIE_KW)
    r.delete_cookie("omc_state")
    return r


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


@app.get("/api/seasons")
def season_list():
    return {"seasons": seasons.all_public()}


@app.get("/api/season")
def season(request: Request):
    return seasons.public(seasons.get(request.query_params.get("season")))


# ---------------------------------------------------------------- 제출
@app.post("/api/submit")
def submit(s: Submission, request: Request):
    season = seasons.get(s.season or None)
    if not season.get("open"):
        raise HTTPException(400, "이 시즌은 제출을 받지 않습니다.")
    user = current_user(request)
    if (OAUTH_ID and OAUTH_SECRET) and not user:
        raise HTTPException(401, "제출하려면 Hugging Face 로그인이 필요합니다.")

    g = seasons.gate(season)
    ok, why, info = gates.check(s.formula, cif=s.structure or None,
                                require_elements=g.get("require_elements", ("Li",)),
                                max_atoms=g.get("max_atoms", 60))
    if not ok:
        raise HTTPException(400, why)

    comp = info["composition"]
    # 열쇠는 **기약 조성**으로 만든다. Li3YCl6 와 Li9Y3Cl18 은 같은 물질이다.
    # 원본 조성으로 만들면 배수만 바꿔 같은 물질을 몇 번이고 올릴 수 있다.
    red = info.get("reduced") or comp
    key = "-".join("%s%g" % (k, red[k]) for k in sorted(red))
    cid = public_id(key)

    ids = _submissions(season)
    if cid in ids:
        raise HTTPException(409, "이미 제출된 조성입니다 (%s)." % cid)

    if user:
        today = int((time.time() + CAP_TZ_OFFSET) // 86400)
        mine = CACHE.get("cap:%s:%d:%d" % (user.get("sub"), season["number"], today),
                         lambda: 0)
        if mine >= DAILY_CAP:
            raise HTTPException(429, "하루 제출 상한(%d건)에 도달했습니다." % DAILY_CAP)

    rec = {"candidate_id": cid, "formula": s.formula.strip(),
           "composition": comp, "n_atoms": info["n_atoms"],
           "display_name": (s.display_name or "").strip()[:60],
           "model_name": (s.model_name or "").strip()[:80],
           "rationale": (s.rationale or "").strip()[:1000],
           "visibility": "public" if s.visibility == "public" else "private",
           "has_structure": bool(s.structure),
           "hf_user": (user or {}).get("name", ""),
           "hf_sub": (user or {}).get("sub", ""),
           "season": season["number"],
           "submitted_at": int(time.time()),
           "uuid": uuid.uuid4().hex}
    store.write(seasons.path(season, "submissions/%s.json" % cid), rec,
                summary="submit %s" % cid)
    if s.structure:
        # 구조는 따로 둔다. 목록을 훑을 때 같이 끌려오면 순위표가 무거워진다.
        store.write(seasons.path(season, "structures/%s.json" % cid),
                    {"candidate_id": cid, "cif": s.structure[:400_000]},
                    summary="structure %s" % cid)
    CACHE.drop("ids:%d" % season["number"])

    scored = set((_board(season).get("entries") or {}).keys())
    ahead = len([i for i in _submissions(season) if i not in scored])
    return {"accepted": True, "candidate_id": cid, "queue_position": ahead,
            "note": "채점은 별도 계산 작업으로 처리되며 완료까지 시간이 걸립니다."}


# ---------------------------------------------------------------- 순위표
def _anchors(season):
    p = os.path.join(DATA, season.get("anchors") or "")
    if not os.path.exists(p):
        return []
    out = []
    for a in json.load(open(p, encoding="utf-8")):
        if not a.get("admitted"):
            continue
        out.append({"candidate_id": a["label"], "is_anchor": True,
                    "display_name": "기준물질", "model_name": "",
                    "formula": a.get("formula", ""),
                    "total": a.get("total"), "total_max": a.get("total_max"),
                    "tier": a.get("tier", 1),
                    "axes": {k: v["points"] for k, v in (a.get("axes") or {}).items()},
                    "detail": {k: v["detail"] for k, v in (a.get("axes") or {}).items()},
                    "note": a.get("note", "")})
    return out


# 음이온 골격이 물질의 성질을 크게 가른다. 같은 계열끼리 견주어야 후보의 위치가 읽힌다.
# 부문을 나누되 **점수는 그대로 둔다** - 어떤 제출도 무효가 되지 않는다.
_ANION_ORDER = ["F", "Cl", "Br", "I", "S", "Se", "O", "N", "P", "H"]
DIVISIONS = {"F": "fluoride", "Cl": "halide", "Br": "halide", "I": "halide",
             "S": "sulfide", "Se": "sulfide", "O": "oxide", "N": "nitride",
             "P": "oxide", "H": "hydride"}


def division_of(comp):
    """골격 음이온으로 부문을 정한다.

    섞여 있으면 **가장 많은 것**을 고른다. Li6PS5Cl 은 S 가 5, Cl 이 1 이므로 황화물이고,
    실제로 아지로다이트는 황화물로 분류된다. 개수가 같으면 위 순서에서 앞선 것을 쓴다.
    """
    if not comp:
        return "other"
    best, bestn = None, -1
    for el in _ANION_ORDER:
        n = comp.get(el, 0)
        if n > bestn:
            best, bestn = el, n
    if bestn <= 0:
        return "other"
    return DIVISIONS.get(best, "other")


def _leaderboard_body(season):
    rows = [mask(r) for r in (_board(season).get("entries") or {}).values()]
    anchors = _anchors(season)
    merged = rows + anchors
    # 점수가 없는 것은 맨 뒤로. 그 안에서는 낸 순서를 지킨다.
    # 전 축이 채워진 것부터, 그 안에서 총점 순. 축이 빈 것은 뒤로 - 총점이 있어도
    # 만점이 달라 같은 줄에 세울 수 없다.
    merged.sort(key=lambda e: (e.get("tier", 1),
                               1 if (e.get("total") is None or e.get("held_axes")) else 0,
                               -(e.get("total") if e.get("total") is not None else 0),
                               e.get("submitted_at") or 0))
    n = 0
    div_n = {}
    for e in merged:
        comp = e.get("composition") or {}
        e["division"] = division_of(comp)
        # 등수는 **전 축이 채워진 것**에만 준다. 축이 하나라도 비면 만점이 달라 총점을
        # 나란히 세울 수 없다 - 점수는 남기되 순위에는 넣지 않는다.
        unscored = e.get("total") is None or bool(e.get("held_axes"))
        if e.get("is_anchor") or e.get("tier", 1) != 1 or unscored:
            e["rank"] = None          # 기준물질은 척도를 표시할 뿐 등수를 갖지 않는다
            e["div_rank"] = None
        else:
            n += 1
            e["rank"] = n
            # 부문 순위는 전체 순위와 따로 매긴다. 점수는 손대지 않는다.
            d = e["division"]
            div_n[d] = div_n.get(d, 0) + 1
            e["div_rank"] = div_n[d]
    tally = {}
    for e in merged:
        if e.get("is_anchor"):
            continue
        tally[e["division"]] = tally.get(e["division"], 0) + 1
    # 🔴 거부와 보류를 한 칸에 합치지 않는다.
    # 합쳐 놓으면 관문에서 **탈락한 것**이 "심사 대기 중"으로 읽힌다.
    # 실측(2026-08-30): 거부 45 · 보류 2 인데 요약은 held 47 로 나갔다.
    # 개별 항목에는 status 가 제대로 실려 있었으므로, 틀린 것은 요약뿐이다.
    real = [e for e in merged if not e.get("is_anchor")]
    n_rejected = sum(1 for e in real if e.get("status") == "rejected")
    n_held = sum(1 for e in real
                 if e.get("status") != "rejected"
                 and (e.get("total") is None or e.get("held_axes")))
    return {"season": seasons.public(season), "entries": merged,
            "divisions": tally,
            "counts": {"scored": len(real) - n_held - n_rejected,
                       "rejected": n_rejected, "held": n_held,
                       "anchors": len(anchors)}}


def _leaderboard_bytes(season):
    raw = json.dumps(_leaderboard_body(season), ensure_ascii=False).encode("utf-8")
    return raw, gzip.compress(raw, 6)


@app.get("/api/leaderboard")
def leaderboard(request: Request):
    """응답은 시즌에만 의존하므로 TTL 당 한 번만 만들고 한 번만 압축한다."""
    season = seasons.get(request.query_params.get("season"))
    raw, gz = CACHE.get("lbresp:%d" % season["number"],
                        lambda: _leaderboard_bytes(season))
    if "gzip" in (request.headers.get("accept-encoding") or ""):
        return Response(content=gz, media_type="application/json",
                        headers={"Content-Encoding": "gzip", "Vary": "Accept-Encoding"})
    return Response(content=raw, media_type="application/json",
                    headers={"Vary": "Accept-Encoding"})


_FILLER = {"ai", "model", "llm", "the", "v", "ver", "version", "latest", "preview",
           "chat", "instruct", "it", "api", "official", "new"}


def canon_model(name):
    """표기가 갈린 같은 모델을 하나로 묶는다.

    'Anthropic Claude Opus 5' 와 'claude-opus-5' 는 같은 것이고, 따로 세면 둘 다
    실제보다 적게 쓰인 것처럼 보인다. 다만 변종을 구분하는 낱말(Pro/Flash 등)은
    지우지 않는다 - 그것까지 뭉치면 다른 모델이 한 칸에 들어간다.
    """
    s = re.sub(r"[^a-z0-9]+", " ", (name or "").lower()).strip()
    toks = [t for t in s.split() if t and t not in _FILLER]
    return " ".join(toks) or (name or "").strip().lower()


@app.get("/api/models")
def models(request: Request):
    """모델별 성적: 어느 모델이 가장 좋은 후보를 냈는가, 그리고 얼마나 쓰이는가.

    두 수를 **따로** 낸다. 많이 쓰인다고 잘하는 것이 아니다 - 다들 집는 모델은 점수와
    무관하게 제출 수가 쌓이고, 세 번 쓰인 모델이 1위 후보를 갖고 있을 수 있다.
    평균만 내면 둘 다 가려진다. 그래서 최고·평균·건수를 나란히 둔다.
    기준물질은 제외한다 - 참가자의 모델이 만든 것이 아니다.
    """
    season = seasons.get(request.query_params.get("season"))
    agg = {}
    for r in (_board(season).get("entries") or {}).values():
        if r.get("total") is None:
            continue
        name = (r.get("model_name") or "").strip() or "미기재 / unspecified"
        key = canon_model(name)
        a = agg.setdefault(key, {"n": 0, "best": None, "sum": 0.0,
                                 "top_candidate": None, "spellings": {}})
        a["spellings"][name] = a["spellings"].get(name, 0) + 1
        a["n"] += 1
        a["sum"] += r["total"]
        if a["best"] is None or r["total"] > a["best"]:
            a["best"] = r["total"]
            a["top_candidate"] = r.get("candidate_id")
    out = []
    for k, a in agg.items():
        # 참가자들이 실제로 가장 많이 쓴 표기를 이름으로 삼는다
        label = max(a["spellings"].items(), key=lambda kv: kv[1])[0]
        out.append({"model": label, "n": a["n"], "best": round(a["best"], 2),
                    "mean": round(a["sum"] / a["n"], 2),
                    "top_candidate": a["top_candidate"]})
    out.sort(key=lambda x: (-x["best"], -x["n"]))
    return {"models": out, "counts": {"models": len(out),
                                      "entries": sum(m["n"] for m in out)}}


@app.get("/api/queue")
def queue(request: Request):
    season = seasons.get(request.query_params.get("season"))
    ids = _submissions(season)
    scored = set((_board(season).get("entries") or {}).keys())
    return {"queued": len([i for i in ids if i not in scored]),
            "scored": len(scored), "total": len(ids)}


@app.get("/api/health")
def health():
    return {"ok": True, "seasons": [s["number"] for s in seasons.SEASONS]}