#!/usr/bin/env python3
"""User-study server: static pages + a small JSON API. Python standard library only
(huggingface_hub is optional, for backing up data/ to a private HF dataset).

    python3 server.py [--port 8000] [--host 0.0.0.0]

Env:
  PART1_N (12), PART2_N (10)  questions per participant
  ADMIN_KEY                   enables GET /api/admin/summary?key=... and /api/admin/export?key=...
  HF_TOKEN + BACKUP_REPO      e.g. BACKUP_REPO=user/us_responses (private dataset): restore data/ on start,
                              upload it every BACKUP_EVERY seconds (120) when it changed

Data (data/):
  participants.json   pid -> {created, questions: [...with the hidden method order...], done: bool}
  responses.jsonl     one line per answer (raw labels + resolved methods)
"""
import argparse
import json
import os
import random
import re
import threading
import time
from http import HTTPStatus
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlparse

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.environ.get("US_DATA", os.path.join(HERE, "data"))
STATIC = os.path.join(HERE, "static")
PART1_N = int(os.environ.get("PART1_N", 12))
PART2_N = int(os.environ.get("PART2_N", 10))
ADMIN_KEY = os.environ.get("ADMIN_KEY", "")
PID_RE = re.compile(r"^[A-Za-z0-9]{1,5}$")
LABELS = "ABCDE"
# a finished ID that has not been opened again for this long is taken: entering it again is refused
TAKEN_AFTER = int(os.environ.get("TAKEN_AFTER", 30 * 60))

def load_manifest():
    """private/manifest.json if present (local runs); otherwise manifest.json from the private BACKUP_REPO dataset,
    so the method <-> video mapping is never part of a public Space repo."""
    p = os.path.join(HERE, "private", "manifest.json")
    if os.path.exists(p):
        return json.load(open(p))
    from huggingface_hub import hf_hub_download
    return json.load(open(hf_hub_download(os.environ["BACKUP_REPO"], "manifest.json", repo_type="dataset",
                                          token=os.environ["HF_TOKEN"])))


MAN = load_manifest()
LOCK = threading.Lock()
os.makedirs(DATA, exist_ok=True)
P_FILE = os.path.join(DATA, "participants.json")
R_FILE = os.path.join(DATA, "responses.jsonl")
DIRTY = threading.Event()


# ---------------------------------------------------------------- optional HF backup
def hf_restore():
    tok, repo = os.environ.get("HF_TOKEN"), os.environ.get("BACKUP_REPO")
    if not (tok and repo):
        return
    try:
        from huggingface_hub import hf_hub_download
        for name in ("participants.json", "responses.jsonl"):
            if not os.path.exists(os.path.join(DATA, name)):
                try:
                    p = hf_hub_download(repo, name, repo_type="dataset", token=tok, local_dir=DATA)
                    print("[backup] restored", p, flush=True)
                except Exception as e:  # first run: nothing there yet
                    print("[backup] no %s in %s (%s)" % (name, repo, type(e).__name__), flush=True)
    except ImportError:
        print("[backup] huggingface_hub not installed; backup disabled", flush=True)


def hf_backup_loop():
    tok, repo = os.environ.get("HF_TOKEN"), os.environ.get("BACKUP_REPO")
    if not (tok and repo):
        return
    from huggingface_hub import HfApi
    api, every = HfApi(token=tok), int(os.environ.get("BACKUP_EVERY", 120))
    while True:
        time.sleep(every)
        if not DIRTY.is_set():
            continue
        DIRTY.clear()
        try:
            with LOCK:
                snap = {n: open(os.path.join(DATA, n), "rb").read() for n in ("participants.json", "responses.jsonl")
                        if os.path.exists(os.path.join(DATA, n))}
            from huggingface_hub import CommitOperationAdd
            api.create_commit(repo, [CommitOperationAdd(n, b) for n, b in snap.items()], repo_type="dataset",
                              commit_message="backup %s" % time.strftime("%Y-%m-%d %H:%M:%S"))
            print("[backup] uploaded", list(snap), flush=True)
        except Exception as e:
            DIRTY.set()
            print("[backup] failed:", e, flush=True)


# ---------------------------------------------------------------- state
def load_participants():
    if os.path.exists(P_FILE):
        return json.load(open(P_FILE))
    return {}


def save_participants(ps):
    tmp = P_FILE + ".tmp"
    with open(tmp, "w") as f:
        json.dump(ps, f)
    os.replace(tmp, P_FILE)
    DIRTY.set()


def exposure(ps, part, assignable=False):
    excl = set(MAN.get("exclude", {}).get(part, [])) if assignable else set()
    c = {s: 0 for s in MAN[part] if s not in excl}
    for p in ps.values():
        for q in p["questions"]:
            if q["part"] == part and q["scene"] in c:
                c[q["scene"]] += 1
    return c


def pick(counts, n, exclude, rng):
    pool = [s for s in counts if s not in exclude]
    rng.shuffle(pool)                               # random tie-break
    pool.sort(key=lambda s: counts[s])              # least-shown scenes first (stable sort keeps the shuffle)
    return pool[:n]


def new_assignment(ps, rng):
    """Part 2 first (smaller pool, 27 scenes), then part 1 from the scenes not used in part 2."""
    s2 = pick(exposure(ps, "part2", assignable=True), PART2_N, set(), rng)
    s1 = pick(exposure(ps, "part1", assignable=True), PART1_N, set(s2), rng)
    rng.shuffle(s1)
    rng.shuffle(s2)
    qs = []
    for part, scenes in (("part1", s1), ("part2", s2)):
        for s in scenes:
            methods = list(MAN[part][s])
            rng.shuffle(methods)                    # display order, per question
            qs.append({"part": part, "scene": s, "order": methods})
    return qs


def last_seen(p):
    return p.get("last_seen") or max([q.get("t", 0) for q in p["questions"]] + [p["created"]])


def is_taken(p):
    return bool(p.get("done")) and time.time() - last_seen(p) > TAKEN_AFTER


def public_questions(p):
    out = []
    for i, q in enumerate(p["questions"]):
        vids = [{"label": LABELS[k], "url": "/static/v/%s.mp4" % MAN[q["part"]][q["scene"]][m]}
                for k, m in enumerate(q["order"])]
        out.append({"qi": i, "part": q["part"], "prompt": MAN["prompts"][q["scene"]], "videos": vids,
                    "answered": "answer" in q, "answer": q.get("answer")})
    return out


# ---------------------------------------------------------------- summary
def summary():
    ps = load_participants()
    m1 = {}
    m2 = {}
    n_done = sum(1 for p in ps.values() if p.get("done"))
    for p in ps.values():
        for q in p["questions"]:
            a = q.get("answer")
            if a is None:
                continue
            if q["part"] == "part1":
                for lab, r in a.items():
                    meth = q["order"][LABELS.index(lab)]
                    d = m1.setdefault(meth, {"n": 0, "rank_sum": 0, "top1": 0})
                    d["n"] += 1
                    d["rank_sum"] += r
                    d["top1"] += int(r == 1)
            else:
                d = m2.setdefault(" vs ".join(sorted(q["order"])), {})
                key = "tie" if a == "tie" else q["order"][LABELS.index(a)]
                d[key] = d.get(key, 0) + 1
    part1 = {m: {"mean_rank": round(d["rank_sum"] / d["n"], 3), "top1_rate": round(d["top1"] / d["n"], 3),
                 "n": d["n"]} for m, d in sorted(m1.items())}
    per_scene1 = {s: c for s, c in exposure(ps, "part1").items()}
    per_scene2 = {s: c for s, c in exposure(ps, "part2").items()}
    return {"participants": len(ps), "completed": n_done, "part1": part1, "part2": m2,
            "assigned_per_scene": {"part1": per_scene1, "part2": per_scene2}}


# ---------------------------------------------------------------- HTTP
class Handler(SimpleHTTPRequestHandler):
    def translate_path(self, path):
        p = urlparse(path).path
        p = p[len("/static"):] if p.startswith("/static/") else "/__none__"
        return super().translate_path(p)        # drops '..' segments, stays inside static/

    def __init__(self, *a, **kw):
        super().__init__(*a, directory=STATIC, **kw)   # only static/ is ever served

    def log_message(self, fmt, *args):
        if "/static/v/" not in str(args[0] if args else ""):
            super().log_message(fmt, *args)

    def end_headers(self):
        if self.path.startswith("/static/v/"):
            self.send_header("Cache-Control", "public, max-age=86400")
            if not self.headers.get("Range"):
                self.send_header("Accept-Ranges", "bytes")
        else:
            self.send_header("Cache-Control", "no-store")
        super().end_headers()

    def send_json(self, obj, code=200):
        b = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def do_GET(self):
        u = urlparse(self.path)
        if u.path in ("/", "/index.html"):
            self.path = "/static/index.html"
        elif u.path in ("/survey", "/survey.html"):
            self.path = "/static/survey.html"
        elif u.path.startswith("/api/admin/"):
            key = parse_qs(u.query).get("key", [""])[0]
            if not ADMIN_KEY or key != ADMIN_KEY:
                return self.send_json({"error": "forbidden"}, 403)
            if u.path == "/api/admin/summary":
                return self.send_json(summary())
            if u.path == "/api/admin/export":
                with LOCK:
                    return self.send_json({"participants": load_participants(),
                                           "responses": [json.loads(l) for l in open(R_FILE)] if os.path.exists(R_FILE) else []})
            return self.send_json({"error": "not found"}, 404)
        elif not u.path.startswith("/static/"):         # never serve data/ or private/
            return self.send_json({"error": "not found"}, 404)
        if u.path.startswith("/static/v/") and self.headers.get("Range"):
            return self.send_range(u.path)
        return super().do_GET()

    def send_range(self, path):
        """HTTP Range for the videos (Safari / iOS will not play an mp4 without 206 responses)."""
        fp = self.translate_path(path)
        if not os.path.isfile(fp):
            return self.send_json({"error": "not found"}, 404)
        size = os.path.getsize(fp)
        m = re.match(r"bytes=(\d*)-(\d*)$", self.headers["Range"].strip())
        if not m or (m.group(1) == "" and m.group(2) == ""):
            self.send_response(HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE)
            self.send_header("Content-Range", "bytes */%d" % size)
            self.end_headers()
            return
        if m.group(1) == "":                         # suffix range: last N bytes
            start, end = max(0, size - int(m.group(2))), size - 1
        else:
            start = int(m.group(1))
            end = min(int(m.group(2)), size - 1) if m.group(2) else size - 1
        if start >= size or start > end:
            self.send_response(HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE)
            self.send_header("Content-Range", "bytes */%d" % size)
            self.end_headers()
            return
        self.send_response(HTTPStatus.PARTIAL_CONTENT)
        self.send_header("Content-Type", "video/mp4")
        self.send_header("Accept-Ranges", "bytes")
        self.send_header("Content-Range", "bytes %d-%d/%d" % (start, end, size))
        self.send_header("Content-Length", str(end - start + 1))
        self.end_headers()
        with open(fp, "rb") as f:
            f.seek(start)
            self.wfile.write(f.read(end - start + 1))

    def do_POST(self):
        u = urlparse(self.path)
        try:
            n = int(self.headers.get("Content-Length", 0))
            body = json.loads(self.rfile.read(min(n, 65536)) or b"{}")
        except Exception:
            return self.send_json({"error": "bad request"}, 400)
        pid = str(body.get("pid", "")).strip()
        if not PID_RE.match(pid):
            return self.send_json({"error": "ID must be 1-5 letters or digits"}, 400)
        key = pid.upper()
        if u.path == "/api/check":                # landing page: is this ID free / resumable / taken?
            with LOCK:
                p = load_participants().get(key)
            st = "new" if p is None else "taken" if is_taken(p) else "done" if p.get("done") else "resume"
            return self.send_json({"status": st})
        if u.path == "/api/start":
            with LOCK:
                ps = load_participants()
                if key in ps and is_taken(ps[key]):
                    return self.send_json({"error": "taken"}, 409)
                if key not in ps:
                    ps[key] = {"pid": pid, "created": time.time(), "ua": self.headers.get("User-Agent", "")[:200],
                               "questions": new_assignment(ps, random.Random()), "done": False}
                ps[key]["last_seen"] = time.time()
                save_participants(ps)
                p = ps[key]
            return self.send_json({"pid": p["pid"], "done": p["done"], "questions": public_questions(p)})
        if u.path == "/api/answer":
            try:
                qi = int(body["qi"])
                ans = body["answer"]
                ms = int(body.get("ms", 0))
            except Exception:
                return self.send_json({"error": "bad request"}, 400)
            with LOCK:
                ps = load_participants()
                if key not in ps or not 0 <= qi < len(ps[key]["questions"]):
                    return self.send_json({"error": "unknown participant or question"}, 400)
                q = ps[key]["questions"][qi]
                k = len(q["order"])
                if q["part"] == "part1":
                    ok = isinstance(ans, dict) and sorted(ans) == list(LABELS[:k]) and \
                        sorted(ans.values()) == list(range(1, k + 1))
                else:
                    ok = ans in ("A", "B", "tie")
                if not ok:
                    return self.send_json({"error": "invalid answer"}, 400)
                if is_taken(ps[key]):
                    return self.send_json({"error": "taken"}, 409)
                q["answer"], q["ms"], q["t"] = ans, ms, time.time()
                ps[key]["last_seen"] = q["t"]
                if q["part"] == "part1":
                    resolved = {q["order"][LABELS.index(l)]: r for l, r in ans.items()}
                else:
                    resolved = "tie" if ans == "tie" else q["order"][LABELS.index(ans)]
                ps[key]["done"] = all("answer" in x for x in ps[key]["questions"])
                save_participants(ps)
                with open(R_FILE, "a") as f:
                    f.write(json.dumps({"pid": key, "qi": qi, "part": q["part"], "scene": q["scene"],
                                        "order": q["order"], "answer": ans, "resolved": resolved, "ms": ms,
                                        "t": q["t"]}) + "\n")
                done = ps[key]["done"]
            return self.send_json({"ok": True, "done": done})
        return self.send_json({"error": "not found"}, 404)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=int(os.environ.get("PORT", 8000)))
    a = ap.parse_args()
    hf_restore()
    threading.Thread(target=hf_backup_loop, daemon=True).start()
    print("part1 %d scenes, part2 %d scenes; %d + %d questions per participant; http://%s:%d" % (
        len(MAN["part1"]), len(MAN["part2"]), PART1_N, PART2_N, a.host, a.port), flush=True)
    ThreadingHTTPServer((a.host, a.port), Handler).serve_forever()


if __name__ == "__main__":
    main()