"""Single source of truth for the Wearable-AI leaderboard.

Shared by the Gradio Space (submit-time validation + board rendering) and the
offline eval worker (re-scoring) so the two can never drift. If you change a
schema here, both sides pick it up.
"""

from __future__ import annotations

import os

# --- HF repos -------------------------------------------------------------
ORG = os.environ.get("ORG", "facebook")
REQUESTS_REPO = f"{ORG}/wearable-ai-leaderboard-requests"
RESULTS_REPO = f"{ORG}/wearable-ai-leaderboard-results"

# Token used by the Space (submit -> requests) and worker (results write).
# Set as a Space secret; falls back to the standard HF env vars locally.
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")

# --- Phase / tracks / divisions ------------------------------------------
PHASE = "val"  # validation phase; test phase reuses these schemas later.

TRACKS = ["convqa", "longqa", "proactive"]

# Why a track is not offered in the TEST phase. Delete an entry here and the
# track becomes available; TEST_TRACKS is derived, and nothing else changes.
#
# This is the ONE line to change. An earlier version of this file hardcoded
# TEST_TRACKS and claimed the same thing, which was false: a test pinned
# proactive absent, so widening the list alone failed the suite. Deriving the
# list means the test asserts the DERIVATION and follows automatically.
# Tracks not offered in the TEST phase. Add an entry to block one; TEST_TRACKS
# is derived, so that is the only edit needed.
TEST_TRACKS_BLOCKED: dict[str, str] = {}

# Tracks the TEST phase can actually score. Offering one we cannot score would
# let a team register a submission that burns one of their three per-track
# slots and returns nothing, so the cost of the gap falls on the participant.
TEST_TRACKS = [t for t in TRACKS if t not in TEST_TRACKS_BLOCKED]

# Pretty names for the UI.
TRACK_LABELS = {
    "convqa": "EgoConv",
    "longqa": "EgoLongQA",
    "proactive": "EgoProactive",
}

# Map a track to the dataset config folder used by the gold jsonls / starter_kit.
TRACK_CONFIG = {
    "convqa": "egoconv",
    "longqa": "egolongqa",
    "proactive": "egoproactive",
}

DIVISIONS = ["small", "large"]

# Divisions are capped on TOTAL params (not active). small: 0 < x <= 2B, large: x > 2B (no upper cap).
DIVISION_CAPS = {
    "small": (0, 2_000_000_000),
    "large": (2_000_000_000, float("inf")),
}

# --- Test phase -----------------------------------------------------------
# The test phase does not take predictions. Shortlisted teams push a container
# image to their own ECR repository and REGISTER its digest here; the organizers'
# worker pulls that digest, runs it against the held-out split on the cluster and
# publishes the score. Registration rather than the registry is the throttle: a
# team can push all it likes, and nothing is evaluated until it is registered.
# Whether the Submit tab still accepts validation predictions. Defaults CLOSED,
# so the freeze takes effect the moment this ships and does not depend on anyone
# remembering to set a secret. Set VAL_PHASE_OPEN=1 to reopen it.
#
# The button being greyed out is not the control: submit.validate_and_submit()
# checks this too, because a disabled button is a suggestion and the endpoint is
# still reachable.
VAL_PHASE_OPEN = os.environ.get("VAL_PHASE_OPEN", "").lower() in ("1", "true", "yes")

# Rendered INSIDE the Validation sub-tab on both the Leaderboard and the Submit
# tab, never above the sub-tabs: the Test phase is not closed, and a banner
# outside them would say it was. It therefore says nothing about "below" or
# "above" either, since it has to read correctly in both places.
VAL_CLOSED_NOTICE = (
    # Dated so a participant arriving later can see WHEN it closed rather than
    # wondering whether the notice is stale.
    "**[08/08/2026] The Validation Phase is closed.** The validation leaderboard is final "
    "and no further validation predictions are accepted. Shortlisted teams "
    "continue in the Test Phase, under **Test** on the Submit tab."
)

TEST_PHASE = "test"

# Whether the Submit tab accepts test-phase registrations. Off until the test
# window opens, so the tab can ship ahead of it; flip with a Space secret rather
# than a code change.
TEST_PHASE_OPEN = os.environ.get("TEST_PHASE_OPEN", "").lower() in ("1", "true", "yes")

# Scored submissions a team may register per TRACK (not per subtrack, and not per
# day as the val phase is): a team's model has one size, so track and subtrack are
# the same in practice here. This is the board-side half of the cap; the worker
# enforces the same number independently, because the board cannot be the only
# gate on work that costs GPU time.
MAX_TEST_SUBMISSIONS_PER_TRACK = 3

# Which request states hold one of those slots. A submission that was refused
# before anything ran (REJECTED) does not; one that ran and failed to score
# (FAILED) does, because it consumed the compute. Keep this identical to
# count_prior_submissions() in the organizers' test_phase_worker.py.
TEST_SLOT_STATUSES = ("PENDING", "FINISHED", "FAILED")

# The image size ceiling, one number for every division, stated as the compressed
# size in the registry (the number `docker push` reports), not the uncompressed
# size on disk. Over it, a team asks for an exception rather than being blocked
# by the form. Nothing in this Space enforces it; the organizer worker rejects
# from the ECR manifest before any pull.
#
# It used to be two numbers. A per-division dict said 10 GB for small while this
# published 100 GB to everyone, and the dict was never actually read here, so a
# small-division team could follow the published figure and still be rejected at
# scoring time. One number now, with no per-division split to drift back into.
TEST_IMAGE_ADVERTISED_GB = 200


# A registry reference we are willing to pull from. Anchored on purpose: the
# reference reaches a subprocess on the organizers' side, so anything that is not
# plainly an ECR repository path is refused here rather than escaped later. These
# two patterns are the same as _ECR_REF / _DIGEST in test_phase_worker.py; a
# reference this accepts and the worker rejects would strand the team with a
# submission that can never be scored.
ECR_REF_RE = (
    r"^(?P<registry>\d{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com)"
    r"/(?P<repository>[a-z0-9][a-z0-9._/-]*[a-z0-9])$"
)
DIGEST_RE = r"^sha256:[0-9a-f]{64}$"

# An image TAG. Docker's own rule: up to 128 of word character, dot or dash,
# not starting with a dot or dash. Anchored like the two above, because this
# value reaches a subprocess on the organizers' side.
#
# A tag is accepted as an alternative to a digest because participant
# repositories are created with --image-tag-mutability IMMUTABLE and the
# per-team policy grants no ecr:Delete*, so a tag cannot be overwritten and
# cannot be freed and re-pushed. On this setup `repo:v1` names one image
# permanently. The worker still resolves it to a digest at intake and records
# that, so provenance stays digest-based; this only changes what a participant
# has to type.
TAG_RE = r"^[A-Za-z0-9_][A-Za-z0-9._-]{0,127}$"

# Whether the Submit form accepts a TAG instead of a digest.
#
# OFF until the worker can resolve one. The worker validates image_digest
# against its own ^sha256:...$ and would reject a tag at intake, so accepting
# one here first would let a participant spend a submission on a reference that
# can never be scored. Flip this in the same change that teaches the worker to
# resolve tag -> digest via `aws ecr describe-images --image-ids imageTag=`,
# not before. The Space side is already written and tested behind it.
TEST_ACCEPT_IMAGE_TAG = False

# --- Gold / submission constraints ---------------------------------------
N_GOLD = 700
MAX_SUBMISSIONS_PER_DAY = 5  # per (hf_user, track, division) = "subtrack", val phase.

# Required keys each prediction row must contain, per track.
PREDICTION_KEYS = {
    "longqa": {"video_path", "mcq_answer"},
    "convqa": {"video_path", "answers"},
    "proactive": {"video_path", "answers"},
}

# --- Score schema (enforced on submit AND emitted by worker) -------------
# set(scores.keys()) must equal REQUIRED[track]; OPTIONAL keys are allowed
# extras. Every value must be a float in [0, 1].
SCORE_SCHEMA = {
    "longqa": {"required": {"accuracy"}, "optional": set()},
    "convqa": {"required": {"llm_judge", "bleu"}, "optional": set()},
    "proactive": {"required": {"macro_f1"}, "optional": {"gmean_f1"}},
}

# Which key the board ranks on, and the proxy column (None if no proxy).
PRIMARY_METRIC = {"longqa": "accuracy", "convqa": "llm_judge", "proactive": "macro_f1"}
PROXY_METRIC = {"longqa": None, "convqa": "bleu", "proactive": None}

# --- Self-reported scores (val phase) ------------------------------------
# ConvQA's main metric (LLM-Judge) is too costly to run on every submission on
# the HF CPU worker (Llama API ~10 RPM), so in the val phase participants run the
# judge themselves and report the score, which the board ranks on (badged
# "self-reported"); the worker still computes verified BLEU as a cross-check, and
# a future internal vLLM judge run can override the self-report. Other tracks are
# fully organizer-scored and need no self-report.
#
# CONVQA_JUDGE_MODEL is the ONE judge every team must use so self-reports are
# comparable: the official Llama-4-Maverick FP8 model, run via the starter_kit
# (starter_kit/run_evaluation.py --task convqa, _build_judge_prompt /
# _parse_judge_score, 0 / 0.5 / 1.0 rubric, temperature 0).
CONVQA_JUDGE_MODEL = "Llama-4-Maverick-17B-128E-Instruct-FP8"

# Required self-reported score keys per track (empty = none collected).
SELF_REPORT_KEYS = {"convqa": {"llm_judge"}, "longqa": set(), "proactive": set()}


_GOLD_IDS_CACHE: dict[str, set[str]] = {}


def load_gold_ids(track: str) -> set[str] | None:
    """Set of expected gold video_path ids for a track, or None if not bundled.

    Shipped under gold_ids/<track>.txt (one id per line). Used to validate that a
    submission covers exactly the gold set (catches wrong/duplicate/missing ids).
    """
    if track in _GOLD_IDS_CACHE:
        return _GOLD_IDS_CACHE[track]
    path = os.path.join(os.path.dirname(__file__), "gold_ids", f"{track}.txt")
    if not os.path.exists(path):
        return None
    with open(path, encoding="utf-8") as fh:
        ids = {line.strip() for line in fh if line.strip()}
    _GOLD_IDS_CACHE[track] = ids
    return ids


def validate_scores(track: str, scores: dict) -> tuple[bool, str]:
    """Return (ok, reason). Shared by Space and worker so they can't drift."""
    if track not in SCORE_SCHEMA:
        return False, f"unknown track {track!r}"
    required = SCORE_SCHEMA[track]["required"]
    optional = SCORE_SCHEMA[track]["optional"]
    keys = set(scores.keys())
    missing = required - keys
    if missing:
        return False, f"missing score keys: {sorted(missing)}"
    unknown = keys - required - optional
    if unknown:
        return False, f"unknown score keys: {sorted(unknown)}"
    for k, v in scores.items():
        if not isinstance(v, (int, float)) or isinstance(v, bool):
            return False, f"score {k!r} is not a number"
        if not (0.0 <= float(v) <= 1.0):
            return False, f"score {k!r}={v} out of range [0,1]"
    return True, ""


def validate_self_report(track: str, scores: dict) -> tuple[bool, str]:
    """Validate participant self-reported scores. Shared by Space and worker.

    For ConvQA the LLM-Judge score is required; other tracks accept none. Every
    value must be a float in [0, 1]. Returns (ok, reason).
    """
    required = SELF_REPORT_KEYS.get(track, set())
    scores = scores or {}
    keys = set(scores.keys())
    missing = required - keys
    if missing:
        return False, f"missing required self-reported score(s): {sorted(missing)}"
    unknown = keys - required
    if unknown:
        return False, f"unexpected self-reported score key(s): {sorted(unknown)}"
    for k, v in scores.items():
        if isinstance(v, bool) or not isinstance(v, (int, float)):
            return False, f"self-reported {k!r} is not a number"
        if not (0.0 <= float(v) <= 1.0):
            return False, f"self-reported {k!r}={v} out of range [0,1]"
    return True, ""


def division_for_params(total_params: int) -> str | None:
    """Map a total param count to its division, or None if over the cap."""
    for div, (lo, hi) in DIVISION_CAPS.items():
        if lo < total_params <= hi:
            return div
    return None