#!/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 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)

    fields = oak_runner.runner_fields()

    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

    from oakbench import scale_calibration

    calibration = None

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

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

        problems = scale_calibration.validate(calibration, fields.get("runner_id"))

        if problems:

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

    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")

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

                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)

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

        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 that produces a quiet-host calibration record.

    Runs the chosen tier with git, git_aa, oak_installed, oak_installed_aa and

    records host load, per-cell A/A floors and A/A false claims. The sealed

    record is written even when it does not validate (reasons are printed);

    only a valid one makes later `run --calibration-record` rows claimable.

    """

    import os as _os

    import statistics as _st

    from oakbench import cellstats, scale_calibration

    from oakbench import runner as oak_runner

    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 not produced:

        print("[scale] calibration produced no rows", file=sys.stderr)

        return 2

    rows = [json.loads(line) for path in produced for line in path.read_text().splitlines() if line.strip()]

    measured = [r for r in rows if r.get("returncode") == 0 and not r.get("diagnostic")]

    value = lambda r: r.get("elapsed_ms")  # noqa: E731

    floors, false_claims = [], 0

    for arm, control in (("git", "git_aa"), ("oak_installed", "oak_installed_aa")):

        cell_floors = cellstats.cell_floors(measured, arm, control, value)

        floors.extend(cell_floors.values())

        verdicts = cellstats.classify_cells(measured, control, arm, cell_floors, value_of=value)

        false_claims += sum(1 for item in verdicts.values() if item.get("verdict") in ("regression", "improvement"))

    cpus = _os.cpu_count() or 1

    loads = [r["environment_sample"]["loadavg"][0] for r in rows

             if isinstance((r.get("environment_sample") or {}).get("loadavg"), list)]

    record = scale_calibration.seal({

        "runner_id": oak_runner.runner_fields().get("runner_id"),

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

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

        "loadavg_per_core_max": round(max(loads) / cpus, 4) if loads else None,

        "aa_cells": len(floors), "aa_floor_median_pct": round(_st.median(floors), 3) if floors else None,

        "aa_false_claims": false_claims, "source_rows": [str(p) for p in produced],

    })

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

    problems = scale_calibration.validate(record, record["runner_id"])

    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())