"""Shared study protocol: booklet assignment, trial order and response validation.

Used unchanged by the Hugging Face backend (hf_backend_app.py), the local server
(serve.py), the response simulator and the analysis, so every path enforces the same
rules. Depends on the standard library only.
"""

import copy
import hashlib
import json
import os
import random
import re
import threading
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path

SCHEMA_VERSION = 1
MAX_PAYLOAD_BYTES = 1024 * 1024
MAX_FEEDBACK_CHARS = 2000
# Crockford base32 without I, L, O, U, grouped as XXXX-XXXX (40 random bits).
PARTICIPANT_PATTERN = re.compile(r"^[0-9A-HJKMNP-TV-Z]{4}-[0-9A-HJKMNP-TV-Z]{4}$")
CHOICES = {"a", "b", "tie"}
SKIP_REASONS = {"video_error"}
# A session that saved within this window still holds its booklet.
ACTIVE_WINDOW = timedelta(hours=2)
# Codes like TEST-0001 are for piloting the live deployment: they never count toward
# booklet balancing and analyze.py ignores them.
TEST_PREFIX = "TEST-"


# ---------------------------------------------------------------------------
# Generic helpers


def read_json(path):
    with Path(path).open("r", encoding="utf-8") as handle:
        return json.load(handle)


def write_json_atomic(value, path):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
    temporary.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    temporary.replace(path)


def stable_seed(*parts):
    payload = ":".join(str(part) for part in parts)
    return int.from_bytes(hashlib.sha256(payload.encode("utf-8")).digest()[:8], "big")


def utc_now():
    return datetime.now(timezone.utc)


def iso(moment):
    return moment.astimezone(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")


def parse_iso(value):
    return datetime.fromisoformat(str(value).replace("Z", "+00:00"))


def deep_merge(base, override):
    merged = copy.deepcopy(base)
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = deep_merge(merged[key], value)
        else:
            merged[key] = copy.deepcopy(value)
    return merged


def load_config(path, override=None):
    config = read_json(path)
    if override:
        config = deep_merge(config, read_json(override))
    return config


def normalize_participant_id(value):
    if not isinstance(value, str):
        raise ValueError("Participant code must be text")
    code = value.strip().upper().replace(" ", "")
    if len(code) == 8 and "-" not in code:
        code = f"{code[:4]}-{code[4:]}"
    if not PARTICIPANT_PATTERN.fullmatch(code):
        raise ValueError("Participant codes look like 7K2M-QX4P")
    return code


# ---------------------------------------------------------------------------
# Trial order


def session_order(manifest, booklet_id, participant_id):
    """Deterministic per-participant order of one booklet's trials.

    Main trials are shuffled with a participant-specific seed; consecutive trials
    never share a prompt when that can be avoided. The first attention check is the
    third pair, so a session abandoned after two pairs can still be screened; the
    other checks are spread evenly over the rest.
    """
    booklet = booklet_by_id(manifest, booklet_id)
    trials = manifest["trials"]
    rng = random.Random(stable_seed(manifest["study_id"], "order", participant_id))

    main = list(booklet["trials"])
    rng.shuffle(main)
    for index in range(1, len(main)):
        if trials[main[index]]["prompt"] != trials[main[index - 1]]["prompt"]:
            continue
        for swap in range(index + 1, len(main)):
            if trials[main[swap]]["prompt"] != trials[main[index - 1]]["prompt"]:
                main[index], main[swap] = main[swap], main[index]
                break

    checks = list(booklet["checks"])
    rng.shuffle(checks)
    total = len(main) + len(checks)
    first = min(2, total - 1)
    positions = [first] if checks else []
    for index in range(1, len(checks)):
        target = first + round(index * (total - first) / len(checks)) + rng.randint(-1, 1)
        positions.append(max(positions[-1] + 2, min(target, total - 1)))

    order = []
    main_iter = iter(main)
    check_iter = iter(checks)
    position_set = set(positions)
    for slot in range(total):
        order.append(next(check_iter) if slot in position_set else next(main_iter, None))
    order = [trial for trial in order if trial is not None]
    order.extend(check_iter)
    return order


def booklet_by_id(manifest, booklet_id):
    for booklet in manifest["booklets"]:
        if booklet["id"] == booklet_id:
            return booklet
    raise ValueError(f"Unknown booklet {booklet_id!r}")


# ---------------------------------------------------------------------------
# Booklet assignment


class Assigner:
    """Hands each new participant the least-filled booklet, in booklet order on ties.

    The load of a booklet counts completed sessions, sessions saved within the
    active window, and the answered fraction of abandoned sessions, so coverage of
    every (prompt, comparison) cell stays balanced even when volunteers drop out.
    Retired booklets (manifest "retired_booklets") go to nobody new. Someone rating
    another set gets the next booklet they have not done, by the same rule, so the
    booklets are used strictly in order and every pair is shown once before any pair
    again. State lives in one JSON file next to the responses.
    """

    def __init__(self, manifest, state_path):
        self.manifest = manifest
        self.state_path = Path(state_path)
        self.lock = threading.Lock()
        self.state = {"participants": {}}
        if self.state_path.exists():
            self.state = read_json(self.state_path)

    def _load(self, now):
        load = {booklet["id"]: 0.0 for booklet in self.manifest["booklets"]}
        for participant, record in self.state["participants"].items():
            if record["booklet"] not in load or participant.startswith(TEST_PREFIX):
                continue
            if record.get("complete"):
                weight = 1.0
            elif now - parse_iso(record["updated_at"]) < ACTIVE_WINDOW:
                weight = 1.0
            else:
                weight = record.get("answered", 0) / max(1, record.get("total", 1))
            load[record["booklet"]] += weight
        return load

    def _person_booklets(self, root):
        return {
            record["booklet"]
            for record in self.state["participants"].values()
            if record.get("root") == root
        }

    def assign(self, participant_id, linked_to=None, now=None):
        now = now or utc_now()
        with self.lock:
            record = self.state["participants"].get(participant_id)
            if record is not None:
                return record["booklet"]
            root = participant_id
            if linked_to and linked_to in self.state["participants"]:
                root = self.state["participants"][linked_to].get("root", linked_to)
            seen = self._person_booklets(root)
            load = self._load(now)
            retired = set(self.manifest.get("retired_booklets", []))
            order = [booklet["id"] for booklet in self.manifest["booklets"] if booklet["id"] not in retired]
            candidates = [booklet for booklet in order if booklet not in seen] or order
            booklet = min(candidates, key=lambda key: (load[key], order.index(key)))
            self.state["participants"][participant_id] = {
                "booklet": booklet,
                "root": root,
                "linked_to": linked_to,
                "assigned_at": iso(now),
                "updated_at": iso(now),
                "answered": 0,
                "total": len(booklet_by_id(self.manifest, booklet)["trials"])
                + len(booklet_by_id(self.manifest, booklet)["checks"]),
                "complete": False,
            }
            write_json_atomic(self.state, self.state_path)
            return booklet

    def booklet_of(self, participant_id):
        record = self.state["participants"].get(participant_id)
        return None if record is None else record["booklet"]

    def record_progress(self, participant_id, answered, complete, now=None):
        now = now or utc_now()
        with self.lock:
            record = self.state["participants"][participant_id]
            record.update(answered=answered, complete=bool(complete), updated_at=iso(now))
            write_json_atomic(self.state, self.state_path)


# ---------------------------------------------------------------------------
# Response validation


def _is_count(value):
    return isinstance(value, int) and not isinstance(value, bool) and value >= 0


def validate_payload(payload, manifest, booklet_id):
    """Raise ValueError unless payload is a well-formed session for booklet_id."""
    if not isinstance(payload, dict):
        raise ValueError("Payload must be an object")
    expected = {
        "schema_version": SCHEMA_VERSION,
        "study_id": manifest["study_id"],
        "booklet_id": booklet_id,
    }
    for key, value in expected.items():
        if payload.get(key) != value:
            raise ValueError(f"{key} does not match this study build")
    # Builds made with build_study.py --extend keep every earlier trial, so pages
    # loaded before an extension can keep saving.
    if payload.get("build_id") not in [manifest["build_id"], *manifest.get("compatible_builds", [])]:
        raise ValueError("build_id does not match this study build")
    normalize_participant_id(payload.get("participant_id"))
    linked_to = payload.get("linked_to")
    if linked_to is not None:
        normalize_participant_id(linked_to)
    for key in ("started_at", "updated_at"):
        parse_iso(payload.get(key))
    if not isinstance(payload.get("complete"), bool):
        raise ValueError("complete must be true or false")

    profile = payload.get("profile", {})
    allowed_profile = {option["id"] for option in manifest["profile"]["options"]} | {None}
    if not isinstance(profile, dict) or profile.get("experience") not in allowed_profile:
        raise ValueError("Invalid profile answer")
    client = payload.get("client", {})
    if not isinstance(client, dict) or len(json.dumps(client)) > 1000:
        raise ValueError("Invalid client information")
    feedback = payload.get("feedback", "")
    if not isinstance(feedback, str) or len(feedback) > MAX_FEEDBACK_CHARS:
        raise ValueError("Feedback is too long")

    booklet = booklet_by_id(manifest, booklet_id)
    allowed = set(booklet["trials"]) | set(booklet["checks"])
    criteria = {criterion["id"] for criterion in manifest["criteria"]}
    choices_allowed = CHOICES if manifest["allow_tie"] else CHOICES - {"tie"}
    trials = payload.get("trials")
    if not isinstance(trials, list):
        raise ValueError("trials must be a list")
    seen = set()
    for entry in trials:
        trial_id = entry.get("trial_id") if isinstance(entry, dict) else None
        if trial_id not in allowed or trial_id in seen:
            raise ValueError(f"Invalid or repeated trial {trial_id!r}")
        seen.add(trial_id)
        trial = manifest["trials"][trial_id]
        if entry.get("a") != trial["a"] or entry.get("b") != trial["b"]:
            raise ValueError(f"Trial {trial_id} videos do not match the study")
        skipped = entry.get("skipped")
        choices = entry.get("choices")
        if skipped is not None:
            if skipped not in SKIP_REASONS or choices != {}:
                raise ValueError(f"Trial {trial_id} has an invalid skip")
        elif not isinstance(choices, dict) or set(choices) != criteria or not set(
            choices.values()
        ) <= choices_allowed:
            raise ValueError(f"Trial {trial_id} must answer every question")
        for key in ("position", "plays", "watch_ms", "elapsed_ms"):
            if not _is_count(entry.get(key)):
                raise ValueError(f"Trial {trial_id} has an invalid {key}")
        parse_iso(entry.get("answered_at"))
    if payload["complete"] and seen != allowed:
        raise ValueError("A complete session must contain every trial")


# ---------------------------------------------------------------------------
# API operations shared by the hosted backend and the local server


class StudyService:
    def __init__(self, manifest, response_root):
        self.manifest = manifest
        root = Path(response_root) / manifest["study_id"]
        self.sessions_dir = root / "sessions"
        self.assigner = Assigner(manifest, root / "assignments.json")

    def start_session(self, participant_id, linked_to=None):
        participant_id = normalize_participant_id(participant_id)
        linked_to = normalize_participant_id(linked_to) if linked_to else None
        booklet_id = self.assigner.assign(participant_id, linked_to)
        saved_path = self.sessions_dir / f"{participant_id}.json"
        saved = read_json(saved_path) if saved_path.exists() else None
        return {
            "participant_id": participant_id,
            "booklet_id": booklet_id,
            "build_id": self.manifest["build_id"],
            "order": session_order(self.manifest, booklet_id, participant_id),
            "saved": saved,
        }

    def save_session(self, payload):
        if isinstance(payload, str):
            payload = json.loads(payload)
        if len(json.dumps(payload).encode("utf-8")) > MAX_PAYLOAD_BYTES:
            raise ValueError("Payload is too large")
        participant_id = normalize_participant_id(payload.get("participant_id"))
        if payload["participant_id"] != participant_id:
            raise ValueError("Participant code is not normalized")
        booklet_id = self.assigner.booklet_of(participant_id)
        if booklet_id is None:
            raise ValueError("Unknown participant; start a session first")
        validate_payload(payload, self.manifest, booklet_id)
        stored = dict(payload, received_at=iso(utc_now()))
        write_json_atomic(stored, self.sessions_dir / f"{participant_id}.json")
        self.assigner.record_progress(participant_id, len(payload["trials"]), payload["complete"])
        return {"saved": True, "answered": len(payload["trials"]), "complete": payload["complete"]}