"""Read-only Nodus Goodput v0.2 evidence validation and per-run scoring.

This checks a supplied evidence bundle, not whether a GPU actually produced it.
Hardware qualification and independent review remain separate release gates.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
from pathlib import Path
import re
import statistics

SCHEMA = "nodus.goodput.run/0.2"
SUITES = {"nodus/serve", "nodus/train", "nodus/kernels"}


class InvalidRun(ValueError):
    """The bundle cannot support a measured score."""


def require(condition, message):
    if not condition:
        raise InvalidRun(message)


def number(value, name, *, positive=False):
    require(type(value) in (int, float) and math.isfinite(value), f"{name}: finite number required")
    require(value > 0 if positive else value >= 0, f"{name}: outside range")
    return value


def integer(value, name, *, positive=False):
    require(type(value) is int, f"{name}: integer required")
    return number(value, name, positive=positive)


def text(value, name):
    require(isinstance(value, str) and bool(value.strip()), f"{name}: nonempty text required")
    return value


def digest(value, name, length=64):
    require(isinstance(value, str) and re.fullmatch(f"[0-9a-f]{{{length}}}", value), f"{name}: digest required")


def pairs(items):
    result = {}
    for key, value in items:
        require(key not in result, f"duplicate JSON key: {key}")
        result[key] = value
    return result


def parse(raw):
    return json.loads(raw, object_pairs_hook=pairs,
                      parse_constant=lambda value: (_ for _ in ()).throw(InvalidRun(f"nonfinite JSON: {value}")))


def artifact(root, ref):
    require(isinstance(ref, dict), "artifact reference required")
    rel = Path(text(ref.get("path"), "artifact path"))
    require(not rel.is_absolute() and ".." not in rel.parts, "artifact must stay inside bundle")
    path = (root / rel).resolve()
    require(path.is_relative_to(root.resolve()), "artifact escapes bundle")
    digest(ref.get("sha256"), "artifact sha256")
    raw = path.read_bytes()
    require(hashlib.sha256(raw).hexdigest() == ref["sha256"], "artifact digest mismatch")
    return raw


def load(root):
    """Load a complete, hash-bound measured bundle; fail closed on missing fields."""
    root = Path(root)
    try:
        meta = parse((root / "run.json").read_bytes())
        require(meta["schema"] == SCHEMA, "unsupported schema")
        require(meta["suite"] in SUITES, "unsupported suite")
        require(meta["record"] in ("measured", "imported"), "MODEL/contextual records cannot receive measured scores")
        text(meta["run_id"], "run_id")
        text(meta["arm"], "arm")
        require(meta["arm_class"] in ("nodus", "baseline", "reference"), "unknown arm class")
        integer(meta["seed"], "seed")
        require(meta["cohort"] in ("development", "holdout"), "unknown cohort")
        require(meta["validity"]["passed"] is True, "run validity failed")
        require(meta["correctness"]["passed"] is True, "correctness failed or unmeasured")
        text(meta["correctness"]["contract"], "correctness contract")
        artifact(root, meta["correctness"]["evidence"])
        provenance = meta["provenance"]
        text(provenance["repository"], "source repository")
        digest(provenance["commit"], "source commit", 40)
        require(isinstance(provenance["sources"], list) and provenance["sources"], "raw sources required")
        for ref in provenance["sources"]:
            artifact(root, ref)
        # The frozen configuration contains model/data revisions, software pins,
        # precision, warmup, tuning budget, arm order and the numerical contract.
        config = parse(artifact(root, meta["configuration"]))
        for key in ("models", "data", "software", "precision", "engine", "warmup", "search_budget", "arm_order", "parallelism"):
            require(key in config and config[key] is not None, f"configuration missing {key}")
        for key in ("models", "data"):
            require(isinstance(config[key], list) and config[key], f"{key}: pinned inputs required")
            for pin in config[key]:
                text(pin["id"], f"{key} id")
                text(pin["revision"], f"{key} revision")
                digest(pin["sha256"], f"{key} content hash")
        require(isinstance(config["software"], dict) and config["software"], "software versions required")
        for name, version in config["software"].items():
            text(name, "software name")
            text(version, "software version")
        text(config["precision"], "precision")
        require(isinstance(config["parallelism"], dict), "parallelism record required")
        for axis in ("tensor", "pipeline", "data"):
            integer(config["parallelism"][axis], axis, positive=True)
        env = meta["environment"]
        for key in ("gpu_model", "driver", "cuda", "interconnect", "image_digest"):
            text(env[key], key)
        require(re.fullmatch(r"sha256:[0-9a-f]{64}", env["image_digest"]), "image must use digest")
        gpu_count = integer(env["gpu_count"], "gpu_count", positive=True)
        window = meta["window"]
        start = number(window["started_s"], "window start")
        stop = number(window["ended_s"], "window end")
        allocations = meta["allocations"]
        require(isinstance(allocations, list) and allocations, "allocation records required")
        gpu_seconds = count = 0
        allocation_ids = set()
        for alloc in allocations:
            allocation_id = text(alloc["id"], "allocation id")
            require(allocation_id not in allocation_ids, "duplicate allocation id")
            allocation_ids.add(allocation_id)
            begin = number(alloc["started_s"], "allocation start")
            end = number(alloc["released_s"], "allocation release")
            require(alloc["release_confirmed"] is True and end > begin, "allocation release unconfirmed or nonpositive")
            require(begin <= start < stop <= end, "window outside allocation")
            n = integer(alloc["gpu_count"], "allocation GPU count", positive=True)
            gpu_seconds += n * (end - begin)
            count += n
        require(count == gpu_count, "allocation GPU count differs from cohort")
        rows = [parse(line) for line in artifact(root, meta["events"]).splitlines() if line.strip()]
        require(len(rows) == integer(meta["expected_events"], "expected_events", positive=True), "missing or extra events")
        ids = [text(row["id"], "event id") for row in rows]
        require(len(set(ids)) == len(ids), "duplicate event ids")
        return meta, rows, gpu_seconds / 3600
    except (KeyError, TypeError, OSError, json.JSONDecodeError) as error:
        raise InvalidRun(f"incomplete or malformed bundle: {type(error).__name__}") from error


def quantile(values, q):
    return sorted(values)[max(0, math.ceil(q * len(values)) - 1)] if values else None


def score(root):
    """Compute suite-specific metrics; never combine unlike suites into one score."""
    meta, rows, gpu_hours = load(root)
    start, stop = meta["window"]["started_s"], meta["window"]["ended_s"]
    result = {"schema": "nodus.goodput.metrics/0.2", "suite": meta["suite"], "run_id": meta["run_id"],
              "record": meta["record"], "cohort": meta["cohort"], "seed": meta["seed"],
              "gpu_hours_allocated": gpu_hours, "events": len(rows), "claim_scope": "research measurement",
              "manifest_sha256": hashlib.sha256((Path(root) / "run.json").read_bytes()).hexdigest()}
    try:
        if meta["suite"] == "nodus/serve":
            ttft_limit = number(meta["slo"]["ttft_s"], "TTFT SLO", positive=True)
            tpot_limit = number(meta["slo"]["tpot_s"], "TPOT SLO", positive=True)
            good_tokens = good_requests = completed = 0
            latencies = []
            for row in rows:
                require(row["kind"] == "request" and type(row["ok"]) is bool, "request status required")
                arrival = number(row["arrival_s"], "arrival")
                done = number(row["done_s"], "done")
                require(start <= arrival <= done <= stop, "request outside measurement window")
                tokens = row["token_times_s"]
                require(isinstance(tokens, list), "token timestamps required")
                for token in tokens:
                    number(token, "token time")
                require(tokens == sorted(tokens) and all(arrival <= t <= done for t in tokens), "invalid token chronology")
                require(len(tokens) == integer(row["output_tokens"], "output tokens"), "token count mismatch")
                if not row["ok"]:
                    continue  # Failure remains in offered-request denominator.
                require(bool(tokens), "successful request has no tokens")
                completed += 1
                latencies.append(done - arrival)
                ttft = tokens[0] - arrival
                # One token has no decode interval and satisfies the TPOT gate vacuously.
                tpot = (tokens[-1] - tokens[0]) / (len(tokens) - 1) if len(tokens) > 1 else 0
                if ttft <= ttft_limit and tpot <= tpot_limit:
                    good_tokens += len(tokens)
                    good_requests += 1
            result.update(completed_requests=completed, slo_requests=good_requests,
                          attainment=good_requests / len(rows), completion_rate=completed / len(rows),
                          slo_output_tokens=good_tokens, slo_output_tokens_per_allocated_gpu_hour=good_tokens / gpu_hours,
                          completed_e2e_p95_s=quantile(latencies, .95))
        elif meta["suite"] == "nodus/train":
            total_tokens = 0
            previous_end = start
            for row in rows:
                require(row["kind"] == "step" and row["validated"] is True, "unvalidated training step")
                begin = number(row["started_s"], "step start")
                end = number(row["ended_s"], "step end")
                require(previous_end <= begin < end <= stop, "overlapping or unordered global training steps")
                previous_end = end
                total_tokens += integer(row["tokens"], "global tokens", positive=True)
                number(row["loss"], "loss")
            result.update(validated_tokens=total_tokens, validated_tokens_per_allocated_gpu_hour=total_tokens / gpu_hours,
                          steady_tokens_per_second_per_gpu=total_tokens / (stop - start) / meta["environment"]["gpu_count"],
                          convergence_claim=False)
        else:
            require(text(meta["baseline"], "production baseline"), "baseline required")
            candidate, baseline = [], []
            shape = text(rows[0]["shape_id"], "shape id")
            for row in rows:
                require(row["kind"] == "kernel" and row["correct"] is True, "incorrect kernel sample")
                require(row["shape_id"] == shape, "different shapes require separate runs")
                candidate.append(number(row["candidate_ms"], "candidate_ms", positive=True))
                baseline.append(number(row["baseline_ms"], "baseline_ms", positive=True))
            result.update(shape_id=shape, baseline=meta["baseline"], candidate_median_ms=statistics.median(candidate),
                          candidate_p95_ms=quantile(candidate, .95), baseline_median_ms=statistics.median(baseline),
                          median_speedup=statistics.median(baseline) / statistics.median(candidate),
                          application_speedup_claim=False)
    except (KeyError, TypeError) as error:
        raise InvalidRun(f"malformed suite events: {type(error).__name__}") from error
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("bundle", type=Path)
    args = parser.parse_args()
    try:
        print(json.dumps(score(args.bundle), sort_keys=True, indent=2, allow_nan=False))
    except (InvalidRun, ValueError) as error:
        parser.exit(2, f"invalid benchmark bundle: {error}\n")


if __name__ == "__main__":
    main()