#!/usr/bin/env python3

"""Scale & coverage lane: every user-facing op, git vs oak, across repo tiers.

Subcommands:

  list                         tiers, core.yaml shapes and the op catalog

  generate --tier T [...]      build (or reuse) fixtures + git-sizer stats only

  run --tier T [...]           prepare workspaces, run the interleaved A/A

                               schedule, write rows + summary

Rows are DEVELOPMENT evidence (evidence_class development_only_loaded_host,

claim_eligible false) until the calibrated campaign runs on a quiet host.

Raw rows go to --results (default results/scale, gitignored; ADR-0003).

Exit: 0 when every executed op passed its oracle, 1 when any row failed

(subject failure, oracle failure or harness error), 2 on usage errors.

"""

from __future__ import annotations

import argparse

import json

import os

import sys

import time

from pathlib import Path

HERE = Path(__file__).resolve().parent

ROOT = HERE.parents[1]

sys.path.insert(0, str(HERE))

sys.path.insert(0, str(ROOT / "scripts"))

from scalesuite import generator as gen  # noqa: E402

from scalesuite import ops as ops_mod  # noqa: E402

from scalesuite import sizer as sz  # noqa: E402

from scalesuite import tiers as tiers_mod  # noqa: E402

DEFAULT_SUBJECTS = "git,git_aa,oak_installed,oak_installed_aa"

def _git_bin(explicit: str | None) -> str:

    if explicit:

        return explicit

    from oakbench.subjects import which_measured_git

    found = which_measured_git()

    if not found:

        raise SystemExit("git not found")

    return found

def _tiers(args: argparse.Namespace) -> list:

    names = []

    for item in args.tier:

        names.extend(part.strip() for part in item.split(",") if part.strip())

    if "shapes" in names:

        names.remove("shapes")

        names.extend(f"shape:{shape['id']}" for shape in tiers_mod.load_core_shapes())

    return [tiers_mod.resolve(name, shape_scale=args.shape_scale, allow_xl=args.allow_xl,

                              allow_xl_combined=args.allow_xl_combined) for name in names]

def cmd_list(_: argparse.Namespace) -> int:

    print("tiers:")

    for name, tier in tiers_mod.TIERS.items():

        spec = tier.spec

        print(f"  {name:12s} files={spec.files} depth={spec.depth} commits={spec.commits} branches={spec.branches} "

              f"tags={spec.tags} binaries={spec.binaries}x{spec.binary_bytes} xl={tier.xl} id={tier.fixture_id}")

    print("core.yaml shapes (shape:<id>, --shape-scale):")

    for shape in tiers_mod.load_core_shapes():

        print(f"  shape:{shape['id']}")

    print("ops:")

    for unit in ops_mod.UNIT_ORDER:

        for op in unit:

            info = ops_mod.OP_INFO[op]

            print(f"  {op:24s} t/perf: {', '.join(info.tperf) or '-'}")

    return 0

def cmd_generate(args: argparse.Namespace) -> int:

    git = _git_bin(args.git_bin)

    work = Path(args.work).resolve()

    out = []

    for tier in _tiers(args):

        started = time.perf_counter()

        manifest = gen.generate(tier.spec, work / tier.fixture_id / "fixture", git)

        pin = gen.check_pin(manifest)  # raises FixturePinMismatch

        stats = sz.measure(work / tier.fixture_id / "fixture" / "src.git", git)

        item = {"tier": tier.tier_id, "fixture_id": tier.fixture_id, "reused": manifest.get("reused"),

                "generation_elapsed_s": manifest["generation_elapsed_s"], "sizer_elapsed_s": stats["elapsed_s"],

                "store_bytes": manifest["store_bytes"], "refs_sha256": manifest["refs_sha256"],

                "head_commit": manifest["head_commit"], "fixture_pin_status": pin, "wall_s": round(time.perf_counter() - started, 3),

                "sizer": stats["values"], "checkout_scope": stats["checkout_scope"],

                "max_level_of_concern": stats["max_level_of_concern"]}

        (work / tier.fixture_id / "fixture" / "SIZER.json").write_text(

            json.dumps({"refs_sha256": manifest["refs_sha256"], "stats": stats}, indent=2) + "\n")

        print(json.dumps(item, sort_keys=True), flush=True)

        out.append(item)

        if args.delete_after:

            import shutil

            shutil.rmtree(work / tier.fixture_id, ignore_errors=True)

    return 0

def _registry() -> dict:

    return json.loads((ROOT / "config" / "fixtures.json").read_text()).get("fixtures", {})

def cmd_run(args: argparse.Namespace) -> int:

    from oakbench import runner as oak_runner

    from oakbench.environment import base_env

    from oakbench.envwatch import sample_environment

    from oakbench.results import ResultsStore

    from oakbench.runlock import measurement_lock

    from oakbench.subjects import load_subjects, subject_details

    from scalesuite import runner, workspace

    if args.runs < 5:

        print("[scale] warning: runs < 5 leaves every cell UNMEASURED (no A/A floor)", file=sys.stderr)

    git = _git_bin(args.git_bin)

    ns = argparse.Namespace(config=Path(args.config), subjects=args.subjects, git_bin=Path(git),

                            oak_installed_bin=Path(args.oak_bin) if args.oak_bin else None, oak_local_bin=None,

                            oak_aa_bin=Path(args.oak_bin) if args.oak_bin else None, allow_git_shim=False)

    subjects = load_subjects(ns)

    names = [s.name for s in subjects]

    kinds = tuple(sorted({s.kind for s in subjects}))

    oak_bin = next((str(s.bin) for s in subjects if s.kind == "oak"), args.oak_bin or "oak")

    details = subject_details(subjects)

    from oakbench import scale_calibration

    fields = {**oak_runner.runner_fields(), **scale_calibration.host_facts()}

    work = Path(args.work).resolve()

    home = work / "home"

    env = base_env(home=home)

    env["PATH"] = str(Path(git).parent) + os.pathsep + env.get("PATH", "")

    ops = ops_mod.ALL_OPS if args.ops == "all" else tuple(op.strip() for op in args.ops.split(",") if op.strip())

    unknown = sorted(set(ops) - set(ops_mod.ALL_OPS))

    if unknown:

        raise SystemExit(f"unknown ops: {', '.join(unknown)}")

    results = Path(args.results).resolve()

    status = 0

    calibration = None

    if getattr(args, "calibration_record", None):

        # Same verification publish_gate performs: recompute the campaign

        # from its own JSONL (matched by content sha256); never trust the record.

        record = json.loads(Path(args.calibration_record).read_text())

        cal_path = Path(str(record.get("calibration_rows_path") or ""))

        cal_rows = scale_calibration.load_rows(cal_path) if cal_path.is_file() else None

        cal_sha = scale_calibration.file_sha256(cal_path) if cal_path.is_file() else None

        facts, problems = scale_calibration.verify_campaign(record, cal_rows, cal_sha, _registry())

        if problems:

            print(f"[scale] calibration rejected (rows stay development-only): {problems}", file=sys.stderr)

        calibration = (record, facts, problems)

    for tier in _tiers(args):

        bench_id = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())

        try:

            prepared = workspace.prepare(tier, work, git_bin=git, oak_bin=oak_bin, env=env, kinds=kinds,

                                         oak_branch_cap=args.oak_branch_cap)

        except Exception as exc:  # noqa: BLE001 - one tier's prep failure must not hide the others

            import traceback

            traceback.print_exc()

            print(f"[scale] {tier.tier_id}: preparation failed: {exc}", file=sys.stderr, flush=True)

            status = 1

            continue

        try:

            load_start = sample_environment().get("loadavg")

            probe_base = {"bench_id": bench_id, "profile": "scale", "scenario": tier.tier_id, **fields,

                          **prepared.fixture_fields()}

            with measurement_lock(f"scale:{tier.tier_id}") as lock:

                probes = scale_calibration.probe_rows(scale_calibration.disk_probe(prepared.root / "probe"),

                                                      probe_base, "start")

                rows = runner.run_schedule(prepared, subjects, runs=args.runs, ops=ops, bench_id=bench_id,

                                           seed=args.seed, scratch=prepared.root / "scratch", details=details,

                                           runner_fields=fields, timeout_s=args.timeout_s,

                                           lock_wait_ms=lock.wait_ms, calibration=calibration)

                probes += scale_calibration.probe_rows(scale_calibration.disk_probe(prepared.root / "probe"),

                                                       probe_base, "end")

            load_end = sample_environment().get("loadavg")

            disk_reasons, _ = scale_calibration.probe_verdict(probes)

            if disk_reasons and calibration is not None:

                # The published run's own disk probes must be quiet too.

                for row in rows:

                    if row.get("claim_eligible") is True:

                        scale_calibration.stamp(row, calibration[0], calibration[1], disk_reasons)

            rows = probes + rows

        finally:

            prepared.close()

        oak_subject = next((s for s in subjects if s.name == "oak_installed"), None)

        if oak_subject is not None:

            prepared.oak_identity = runner._subject_meta(oak_subject, details)

        rows = runner.ingest_rows(prepared, bench_id, fields) + rows

        report = runner.cell_report(rows, names)

        meta = {"bench_id": bench_id, "runs": args.runs, "subjects": names, "loadavg_start": load_start,

                "loadavg_end": load_end, "calibrated": any(r.get("claim_eligible") is True for r in rows)}

        summary = runner.summary_markdown(prepared, report, meta)

        store = ResultsStore(results, "scale", filename_suffix=tier.tier_id)

        raw_path, summary_path = store.write(bench_id, rows, summary)

        (results / f"{bench_id}.{tier.tier_id}.cells.json").write_text(

            json.dumps({"meta": meta, "cells": report, "prep_phases": prepared.phases,

                        "sizer": prepared.sizer, "manifest": prepared.manifest}, indent=2, default=str) + "\n")

        print(summary)

        print(f"[scale] rows: {raw_path}")

        if any(r.get("returncode") == 1 for r in rows):

            status = 1

        if args.delete_workspaces:

            import shutil

            shutil.rmtree(prepared.root / "git", ignore_errors=True)

            shutil.rmtree(prepared.root / "oak", ignore_errors=True)

            shutil.rmtree(prepared.root / "scratch", ignore_errors=True)

    return status

def cmd_calibrate(args: argparse.Namespace) -> int:

    """Identical-binary A/A campaign + host probes -> a sealed calibration record.

    The record names the campaign JSONL by content sha256; its numbers are

    recomputed from that JSONL wherever it is used (run, publish_gate). It is

    written even when the host is not quiet (reasons are printed).

    """

    from oakbench import scale_calibration

    if len(_tiers(args)) != 1:

        raise SystemExit("calibrate takes exactly one tier")

    args.subjects = DEFAULT_SUBJECTS

    args.calibration_record = None

    results = Path(args.results).resolve()

    before = set(results.glob("*.jsonl")) if results.exists() else set()

    rc = cmd_run(args)

    produced = sorted(p for p in set(results.glob("*.jsonl")) - before if not p.name.startswith("latest"))

    if len(produced) != 1:

        print(f"[scale] calibration expected one JSONL, got {produced}", file=sys.stderr)

        return 2

    rows = scale_calibration.load_rows(produced[0])

    facts, problems = scale_calibration.recompute(rows, _registry())

    record = scale_calibration.seal({

        **{key: facts.get(key) for key in ("runner_id", "host_cpu_count", "host_mem_bytes", "oak_binary_sha256",

                                           "git_binary_sha256", "aa_cells", "aa_floor_median_pct",

                                           "aa_false_claims", "loadavg_per_core_max", "fullfsync_median_ms",

                                           "create_files_per_s_min")},

        "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),

        "calibration_bench_id": rows[0].get("bench_id"),

        "calibration_rows_path": str(produced[0]),

        "calibration_rows_sha256": scale_calibration.file_sha256(produced[0]),

    })

    Path(args.out).write_text(json.dumps(record, indent=2, sort_keys=True) + "\n")

    print(json.dumps({"record": args.out, "valid": not problems, "reasons": problems}, indent=2))

    return rc if rc else (0 if not problems else 1)

def parser() -> argparse.ArgumentParser:

    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)

    sub = ap.add_subparsers(dest="command", required=True)

    sub.add_parser("list")

    for name in ("generate", "run", "calibrate"):

        p = sub.add_parser(name)

        p.add_argument("--tier", action="append", required=True,

                       help="tier id (fresh, default, large, xl-*), shape:<core.yaml id>, or 'shapes'")

        p.add_argument("--shape-scale", type=float, default=1.0)

        p.add_argument("--allow-xl", action="store_true")

        p.add_argument("--allow-xl-combined", action="store_true")

        p.add_argument("--work", default=str(ROOT / "results" / "scale-work"),

                       help="fixture + workspace root (large; keep it on the benchmark volume)")

        p.add_argument("--git-bin")

        if name == "generate":

            p.add_argument("--delete-after", action="store_true")

        else:

            if name == "calibrate":

                p.add_argument("--out", required=True, help="where to write the sealed calibration record")

            else:

                p.add_argument("--calibration-record",

                               help="validated quiet-host record (from `calibrate`); without it rows are dev-only")

            p.add_argument("--subjects", default=DEFAULT_SUBJECTS)

            p.add_argument("--config", default=str(ROOT / "config" / "subjects.toml"))

            p.add_argument("--oak-bin")

            p.add_argument("--runs", type=int, default=5)

            p.add_argument("--ops", default="all")

            p.add_argument("--seed", type=int, default=20260928)

            p.add_argument("--timeout-s", type=float, default=1800.0)

            p.add_argument("--oak-branch-cap", type=int)

            p.add_argument("--results", default=str(ROOT / "results" / "scale"))

            p.add_argument("--delete-workspaces", action="store_true")

    return ap

def main(argv: list | None = None) -> int:

    args = parser().parse_args(argv)

    return {"list": cmd_list, "generate": cmd_generate, "run": cmd_run,

            "calibrate": cmd_calibrate}[args.command](args)

if __name__ == "__main__":

    raise SystemExit(main())