"""Data Experiments dashboard, served from a Hugging Face Space.

The page is a static shell (static/index.html + static/app.js) that holds no data. Everything it shows comes from
/api/*, and every /api route needs a signed-in, allow-listed viewer:

  sign-in   "Sign in with Hugging Face" (the Space's OAuth app: hf_oauth in README.md). Server-side authorization-code
            flow with PKCE and a signed, 10-minute state cookie; the callback trades the code for a token and reads
            /oauth/userinfo server to server. The viewer then gets a signed HttpOnly session cookie (24 h).
  allow     experiments/_access.json in R2 (edited with `python dataops/experiments_db.py access ...`), re-read at most
            every ACCESS_TTL seconds, checked on EVERY request, so a removal takes effect within a minute:
            an allowed HF username, membership of an allowed HF org, an allowed email, or an email in an allowed
            domain. Emails only count when Hugging Face reports them verified.
  data      JSON records under experiments/ (projects/, docs/, sft/, _access.json), the same ones sessions write with
            dataops/experiments_db.py. SFT marks are written here, recording who marked them. Two interchangeable
            stores, picked by DATA_BACKEND:
              r2  (default) bucket ops-dash; R2_ACCOUNT_ID / R2_ACCESS_KEY_ID / R2_SECRET_ACCESS_KEY
              hf  a private Hugging Face dataset repo, same paths; DATA_REPO (default Extrian/data-experiments-db)
                  and HF_DATA_TOKEN (read + write on that one repo)

Settings come from the environment: OAUTH_CLIENT_ID / OAUTH_CLIENT_SECRET / OAUTH_SCOPES / SPACE_HOST are set by the
Space; the store's credentials are Space secrets. SESSION_SECRET (signs the cookies) is optional: without it a random one
is made at startup, so everyone is signed out whenever the Space restarts. The R2 key only ever touches
ops-dash/experiments/ (every key is checked), so a key limited to bucket ops-dash is enough.
"""
import base64, hashlib, html, json, os, re, secrets, threading, time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlencode

import boto3, botocore, httpx
from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, RedirectResponse
from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer

HF = "https://huggingface.co"
BUCKET, PREFIX = "ops-dash", "experiments/"
ACCESS_KEY = PREFIX + "_access.json"
ID_RE = re.compile(r"^[a-z0-9][a-z0-9._-]{0,99}$")
SESSION_COOKIE, STATE_COOKIE = "dx_session", "dx_oauth"
SESSION_TTL, STATE_TTL, ACCESS_TTL, ACCESS_STALE_MAX = 24 * 3600, 600, 60, 600
STATIC = Path(__file__).resolve().parent / "static"

CSP = ("default-src 'self'; script-src 'self' https://cdnjs.cloudflare.com; "
       "style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; font-src https://fonts.gstatic.com; "
       "img-src 'self' data:; connect-src 'self'; frame-ancestors https://huggingface.co; base-uri 'none'; "
       "form-action 'self' https://huggingface.co")


def env(name, default=None):
    v = os.environ.get(name, default)
    if not v:
        raise RuntimeError(f"{name} is not set")
    return v


def public_origin():
    return "https://" + env("SPACE_HOST").strip().rstrip("/")


_SESSION_SECRET = os.environ.get("SESSION_SECRET") or ""
if not _SESSION_SECRET:
    _SESSION_SECRET = secrets.token_urlsafe(48)
    print("SESSION_SECRET is not set: using a random one for this run, so sign-ins reset whenever the Space restarts",
          flush=True)


def signer(purpose):
    return URLSafeTimedSerializer(_SESSION_SECRET, salt="dx-" + purpose)


def inside(key):
    """Every store access stays under experiments/ (a narrow R2 key may not reach anything else)."""
    if not key.startswith(PREFIX) or ".." in key:
        raise ValueError(f"refusing to touch {key!r}: only {PREFIX} is used")
    return key


def now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


# ------------------------------------------------------------------------------------------------ stores
class R2Store:
    def __init__(self):
        self.c = boto3.client("s3", endpoint_url=f"https://{env('R2_ACCOUNT_ID')}.r2.cloudflarestorage.com",
                              aws_access_key_id=env("R2_ACCESS_KEY_ID"), aws_secret_access_key=env("R2_SECRET_ACCESS_KEY"),
                              region_name="auto", config=botocore.config.Config(retries={"max_attempts": 3},
                                                                                 connect_timeout=5, read_timeout=15))

    def get(self, key):
        try:
            return json.loads(self.c.get_object(Bucket=BUCKET, Key=inside(key))["Body"].read())
        except botocore.exceptions.ClientError as e:
            if e.response["Error"]["Code"] in ("NoSuchKey", "404"):
                return None
            raise

    def keys(self, prefix):
        out, kw = [], {"Bucket": BUCKET, "Prefix": inside(prefix)}
        while True:
            r = self.c.list_objects_v2(**kw)
            out += [o["Key"] for o in r.get("Contents", [])]
            if not r.get("IsTruncated"):
                return out
            kw["ContinuationToken"] = r["NextContinuationToken"]

    def put(self, key, obj, note):
        self.c.put_object(Bucket=BUCKET, Key=inside(key), Body=json.dumps(obj).encode(), ContentType="application/json; charset=utf-8")


class HFStore:
    """A private Hugging Face dataset repo holding the same paths. Reads go to resolve/main with the token (fresh, no
    local cache); every write is one commit."""

    def __init__(self):
        from huggingface_hub import HfApi
        self.repo, self.token = env("DATA_REPO", "Extrian/data-experiments-db"), env("HF_DATA_TOKEN")
        self.api = HfApi(token=self.token)
        self.http = httpx.Client(timeout=15, headers={"Authorization": "Bearer " + self.token}, follow_redirects=True)

    def get(self, key):
        r = self.http.get(f"{HF}/datasets/{self.repo}/resolve/main/{inside(key)}")
        if r.status_code == 404:
            return None
        r.raise_for_status()
        return r.json()

    def keys(self, prefix):
        from huggingface_hub.errors import EntryNotFoundError
        try:
            items = self.api.list_repo_tree(self.repo, path_in_repo=inside(prefix).rstrip("/"), repo_type="dataset")
            return [i.path for i in items if getattr(i, "size", None) is not None]
        except EntryNotFoundError:
            return []

    def put(self, key, obj, note):
        self.api.upload_file(path_or_fileobj=json.dumps(obj).encode(), path_in_repo=inside(key), repo_id=self.repo,
                             repo_type="dataset", commit_message=note)


_store = None


def store():
    global _store
    if _store is None:
        kind = os.environ.get("DATA_BACKEND", "r2").strip().lower()
        if kind not in ("r2", "hf"):
            raise RuntimeError(f"DATA_BACKEND must be r2 or hf, got {kind!r}")
        _store = R2Store() if kind == "r2" else HFStore()
    return _store


def r2_get(key):  # kept under this name: the tests patch it
    return store().get(key)


def r2_read_all(coll):
    prefix = f"{PREFIX}{coll}/"
    keys = [k for k in store().keys(prefix) if k.endswith(".json") and "/" not in k[len(prefix):]]
    with ThreadPoolExecutor(16) as ex:
        vals = list(ex.map(r2_get, keys))
    return {k[len(prefix):-5]: v for k, v in zip(keys, vals) if v is not None}


# ------------------------------------------------------------------------------------------------ access list
_acl = {"data": None, "at": 0.0, "error": None}
_acl_lock = threading.Lock()
R2_ENV = ("R2_ACCOUNT_ID", "R2_ACCESS_KEY_ID", "R2_SECRET_ACCESS_KEY")


def _error_code(e):
    """A short, value-free reason for a failed store read: the S3 error code or the exception class."""
    if isinstance(e, botocore.exceptions.ClientError):
        return e.response.get("Error", {}).get("Code") or "ClientError"
    if isinstance(e, RuntimeError) and str(e).endswith("is not set"):
        return str(e)  # names the missing variable, never a value
    return type(e).__name__


def access_list():
    """The allow list, re-read at most every ACCESS_TTL s. If R2 is unreachable the last good copy is used for up to
    ACCESS_STALE_MAX s; past that, or with no copy at all, nobody gets in (fail closed)."""
    with _acl_lock:
        if _acl["data"] is not None and time.time() - _acl["at"] < ACCESS_TTL:
            return _acl["data"]
        try:
            _acl["data"], _acl["at"], _acl["error"] = r2_get(ACCESS_KEY) or {}, time.time(), None
        except Exception as e:  # noqa: BLE001
            _acl["error"] = _error_code(e)
            print(f"access list read failed: {_acl['error']}", flush=True)
            if _acl["data"] is None or time.time() - _acl["at"] > ACCESS_STALE_MAX:
                raise
        return _acl["data"]


def allowed(ident, acl):
    """Why this viewer may see the page ("hf_user", "hf_org", "email", "domain"), or None."""
    low = lambda xs: {str(x).strip().lower().lstrip("@") for x in xs or []}
    user = str(ident.get("u") or "").lower()
    if user and user in low(acl.get("hf_users")):
        return "hf_user"
    if low(ident.get("orgs")) & low(acl.get("hf_orgs")):
        return "hf_org"
    email = str(ident.get("e") or "").strip().lower()
    if email and ident.get("ev") is True and email.count("@") == 1:
        if email in low(acl.get("emails")):
            return "email"
        if email.split("@")[1] in low(acl.get("domains")):
            return "domain"
    return None


def identity_from_userinfo(info):
    orgs = [o.get("preferred_username") or o.get("name") for o in info.get("orgs") or [] if isinstance(o, dict)]
    return {"u": info.get("preferred_username") or "", "sub": str(info.get("sub") or ""),
            "e": (info.get("email") or "").strip().lower(), "ev": info.get("email_verified") is True,
            "orgs": [o for o in orgs if o], "iat": int(time.time())}


def viewer(request):
    """(identity, None) for an allowed viewer, else (None, response)."""
    raw = request.cookies.get(SESSION_COOKIE)
    if not raw:
        return None, JSONResponse({"error": "sign in", "login": "/login"}, 401)
    try:
        ident = signer("session").loads(raw, max_age=SESSION_TTL)
    except SignatureExpired:
        return None, JSONResponse({"error": "your sign-in expired", "login": "/login"}, 401)
    except BadSignature:
        return None, JSONResponse({"error": "sign in", "login": "/login"}, 401)
    try:
        why = allowed(ident, access_list())
    except Exception:  # noqa: BLE001
        return None, JSONResponse({"error": "could not check the access list; try again shortly"}, 503)
    if not why:
        return None, JSONResponse({"error": "not on the access list", "who": label(ident)}, 403)
    return ident, None


def label(ident):
    return ident.get("e") if ident.get("ev") and ident.get("e") else "@" + (ident.get("u") or "?")


# ------------------------------------------------------------------------------------------------ app
app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)


@app.middleware("http")
async def headers(request, call_next):
    resp = await call_next(request)
    resp.headers.setdefault("Cache-Control", "no-store")
    resp.headers["X-Content-Type-Options"] = "nosniff"
    resp.headers["Referrer-Policy"] = "no-referrer"
    resp.headers["Content-Security-Policy"] = CSP
    return resp


def cookie(resp, name, value, max_age, path="/"):
    resp.set_cookie(name, value, max_age=max_age, path=path, secure=True, httponly=True, samesite="lax")


@app.get("/")
def index():
    return FileResponse(STATIC / "index.html", media_type="text/html; charset=utf-8")


@app.get("/status")
def status():
    """Unauthenticated health check: which store settings are present (true/false, never values) and whether the
    access list can be read (an error code at most). Shows no data."""
    try:
        access_list()
        state = "ok"
    except Exception:  # noqa: BLE001
        state = _acl["error"] or "error"
    return JSONResponse({"store": STORE if "STORE" in globals() else "r2",
                         "r2_settings_present": {n: bool(os.environ.get(n)) for n in R2_ENV},
                         "access_list": state}, headers={"Cache-Control": "no-store"})


@app.get("/app.js")
def app_js():
    return FileResponse(STATIC / "app.js", media_type="text/javascript; charset=utf-8")


@app.get("/login")
def login():
    state, verifier = secrets.token_urlsafe(24), secrets.token_urlsafe(48)
    challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode()
    scopes = " ".join(dict.fromkeys(["openid", "profile", "email"] + os.environ.get("OAUTH_SCOPES", "").split()))
    q = {"client_id": env("OAUTH_CLIENT_ID"), "redirect_uri": public_origin() + "/login/callback",
         "response_type": "code", "scope": scopes, "state": state,
         "code_challenge": challenge, "code_challenge_method": "S256"}
    resp = RedirectResponse(f"{HF}/oauth/authorize?{urlencode(q)}", 302)
    cookie(resp, STATE_COOKIE, signer("state").dumps({"s": state, "v": verifier}), STATE_TTL, path="/login")
    return resp


def page(title, body, status=200):
    return HTMLResponse(f"""<!doctype html><html><head><meta charset="utf-8"><meta name="viewport"
content="width=device-width,initial-scale=1"><title>{html.escape(title)}</title><style>body{{font:15px/1.55 system-ui,
sans-serif;margin:0;padding:48px 16px;background:#f3f5f8;color:#16202c}}main{{max-width:560px;margin:0 auto;
background:#fff;border:1px solid #dde2ea;border-radius:12px;padding:24px}}a{{color:#2a5bd7}}code{{background:#eef1f5;
padding:1px 5px;border-radius:4px}}@media (prefers-color-scheme:dark){{body{{background:#0f141b;color:#e4e9f1}}
main{{background:#161d27;border-color:#283241}}a{{color:#7ea6ff}}code{{background:#1d2633}}}}</style></head>
<body><main>{body}</main></body></html>""", status)


@app.get("/login/callback")
def callback(request: Request, code: str = "", state: str = "", error: str = ""):
    raw = request.cookies.get(STATE_COOKIE)
    try:
        st = signer("state").loads(raw or "", max_age=STATE_TTL)
    except (BadSignature, SignatureExpired):
        st = None
    if error or not st or not code or not state or not secrets.compare_digest(str(st.get("s")), state):
        resp = page("Sign-in failed", "<h1>Sign-in didn't complete</h1><p>The sign-in link expired or came from "
                    "somewhere else. <a href='/login'>Try again</a>.</p>", 400)
        resp.delete_cookie(STATE_COOKIE, path="/login")
        return resp
    try:
        with httpx.Client(timeout=15) as h:
            tok = h.post(f"{HF}/oauth/token", auth=(env("OAUTH_CLIENT_ID"), env("OAUTH_CLIENT_SECRET")),
                         data={"grant_type": "authorization_code", "code": code, "code_verifier": st["v"],
                               "redirect_uri": public_origin() + "/login/callback"})
            tok.raise_for_status()
            info = h.get(f"{HF}/oauth/userinfo", headers={"Authorization": "Bearer " + tok.json()["access_token"]})
            info.raise_for_status()
            info = info.json()
    except Exception:  # noqa: BLE001
        return page("Sign-in failed", "<h1>Couldn't reach Hugging Face</h1><p><a href='/login'>Try again</a>.</p>", 502)
    ident = identity_from_userinfo(info)
    try:
        why = allowed(ident, access_list())
    except Exception:  # noqa: BLE001
        return page("Try again", "<h1>Couldn't check the access list</h1><p><a href='/login'>Try again</a> in a minute.</p>", 503)
    if not why:
        who = html.escape(label(ident))
        mail = "" if ident.get("ev") or not ident.get("e") else (
            "<p>Your Hugging Face email isn't verified, so it can't be used to let you in. Verify it in your Hugging Face "
            "settings, or ask for your username to be added.</p>")
        resp = page("No access yet", f"<h1>You're signed in, but not on the access list</h1><p>Signed in as "
                    f"<code>{who}</code> (Hugging Face user <code>{html.escape(ident.get('u') or '?')}</code>).</p>{mail}"
                    "<p>Ask the dashboard owner to add your email or Hugging Face username, then "
                    "<a href='/login'>sign in again</a>.</p>", 403)
        resp.delete_cookie(STATE_COOKIE, path="/login")
        return resp
    resp = RedirectResponse("/", 302)
    resp.delete_cookie(STATE_COOKIE, path="/login")
    cookie(resp, SESSION_COOKIE, signer("session").dumps(ident), SESSION_TTL)
    return resp


@app.get("/logout")
def logout():
    resp = RedirectResponse("/", 302)
    resp.delete_cookie(SESSION_COOKIE, path="/")
    return resp


@app.get("/api/me")
def me(request: Request):
    ident, err = viewer(request)
    return err or {"me": label(ident), "user": ident.get("u")}


@app.get("/api/data")
def data(request: Request):
    ident, err = viewer(request)
    if err:
        return err
    with ThreadPoolExecutor(3) as ex:
        projects, docs, sft = ex.map(r2_read_all, ("projects", "docs", "sft"))
    index = {i: {"title": d.get("title"), "path": d.get("path"), "has_markdown": bool(d.get("markdown")),
                 "updated_at": d.get("updated_at")} for i, d in docs.items()}
    return {"projects": projects, "docs": index, "sft": sft, "me": label(ident)}


@app.get("/api/doc/{doc_id}")
def doc(doc_id: str, request: Request):
    ident, err = viewer(request)
    if err:
        return err
    if not ID_RE.match(doc_id):
        return JSONResponse({"error": "bad doc id"}, 400)
    d = r2_get(f"{PREFIX}docs/{doc_id}.json")
    return d if d is not None else JSONResponse({"error": "no such guide"}, 404)


async def sft_write(slug, request, done):
    ident, err = viewer(request)
    if err:
        return err
    # Writes must come from this page: same Origin and a custom header a cross-site form can't send.
    if request.headers.get("origin") != public_origin() or request.headers.get("x-dashboard") != "1":
        return JSONResponse({"error": "writes must come from the dashboard page"}, 403)
    if not ID_RE.match(slug):
        return JSONResponse({"error": "bad project id"}, 400)
    if r2_get(f"{PREFIX}projects/{slug}.json") is None:
        return JSONResponse({"error": "no such project"}, 404)
    note = ""
    if done:
        try:
            body = await request.json()
        except Exception:  # noqa: BLE001
            return JSONResponse({"error": "send a JSON body"}, 400)
        note = str((body or {}).get("note") or "")[:300]
    mark = {"done": done, "by": label(ident), "by_hf": ident.get("u"), "at": now(), "note": note}
    store().put(f"{PREFIX}sft/{slug}.json", mark, f"SFT {'mark' if done else 'undo'} {slug} by {mark['by']}")
    return mark


@app.post("/api/sft/{slug}")
async def sft_mark(slug: str, request: Request):
    return await sft_write(slug, request, True)


@app.delete("/api/sft/{slug}")
async def sft_undo(slug: str, request: Request):
    return await sft_write(slug, request, False)