#!/usr/bin/env python3
"""G2 — census sanity gate (design doc §3, gate table row "G2").

Runs after the `census` stage. Reads `<work>/census.json` (written by
`garment_pipeline.py` stage_census) and checks it against the resolved config's
`expect.islands` block.

Catches the silent killer the design calls out: an MD re-export reorders or
renumbers islands, the config's `parts[*].islands` indices now point at the
wrong geometry, and every downstream stage happily builds the wrong garment.

    python clothing/gates/g2_census.py --config work/<name>/resolved.json --work work/<name>

Exit codes: 0 pass · 2 fail (thresholds violated) · 3 error (could not evaluate).
Writes `<work>/qc/g2.json` in the contract shape.

census.json shape (verified against work/piupiu_sb, work/piupiu, work/pari):

    [ {"island": 0, "verts": 3620, "tris": 7104,
       "z": [5.184, 11.126], "x": [-2.228, 2.517],
       "centroid": [-0.034, 0.315, 8.416], "color": [0.9, 0.1, 0.1]} ]

A list, one record per island, `island` is the index the config's
`parts[*].islands` refers to. UNITS: raw source-mesh units, NOT metres — the
piupiu census reads z 5.18..11.13 because `align.scale_z` (0.1) is applied later
in `prepare`. Author `expect.islands` z bounds in census units (divide the metre
figure by the align scale, or just read them off census.json).

`expect.islands` schema understood by this gate (superset of PIPELINE-CONTRACT.md):

    "islands": {
      "count": 1,                     // exact island count in census.json
      "require_all_mapped": false,    // unmapped non-trivial islands => fail (default warn)
      "mapped": {
        "0": {                        // key: island index, or a part name from parts{}
          "verts":  [3500, 3700],     // inclusive range on vert count
          "tris":   [7000, 7200],     // optional
          "z":      [5.0, 11.3],      // island z extent must be CONTAINED in this range
          "z_min":  [5.1, 5.3],       // optional: range on the low z endpoint
          "z_max":  [11.0, 11.2],     // optional: range on the high z endpoint
          "z_span": [5.8, 6.1],       // optional: range on (zmax - zmin)
          "x":      [-2.4, 2.6]       // optional: x extent containment
        }
      }
    }

Every range accepts `[lo, hi]` (either endpoint may be null for one-sided), a
bare number (exact match), or `{"min": lo, "max": hi}`.

Independent of `expect`, the gate ALWAYS cross-checks that every island index in
`parts[*].islands` exists in census.json — that check needs no thresholds, so a
config with no `expect.islands` still gets it (and passes with a note rather
than erroring; the orchestrator only schedules this gate when the key exists).

Stdlib only.
"""

from __future__ import annotations

import argparse
import json
import os
import sys

GATE_ID = "g2"
STAGE = "census"

EXIT_PASS, EXIT_FAIL, EXIT_ERROR = 0, 2, 3


# ── report plumbing ──────────────────────────────────────────────────────────

class Report:
    def __init__(self):
        self.metrics = {}
        self.failures = []
        self.warnings = []
        self.notes = []
        self.artifacts = []

    def fail(self, part, detail, **extra):
        entry = {"part": part, "detail": detail}
        entry.update(extra)
        self.failures.append(entry)

    def warn(self, part, detail, **extra):
        entry = {"part": part, "detail": detail}
        entry.update(extra)
        self.warnings.append(entry)

    def to_dict(self):
        return {
            "gate": GATE_ID,
            "pass": not self.failures,
            "checked_at_stage": STAGE,
            "metrics": self.metrics,
            "failures": self.failures,
            "warnings": self.warnings,
            "notes": self.notes,
            "artifacts": self.artifacts,
        }


def write_report(work, report, out=None):
    """Write <work>/qc/g2.json (or --out). Never raises; returns the path or None."""
    path = out
    if not path:
        if not work:
            return None
        path = os.path.join(work, "qc", "%s.json" % GATE_ID)
    try:
        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(report, fh, indent=2)
            fh.write("\n")
        return path
    except OSError as exc:  # pragma: no cover - disk/permission edge
        sys.stderr.write("[g2] WARNING could not write report %s: %s\n" % (path, exc))
        return None


def die(work, detail, out=None):
    """Exit 3 with a contract-shaped error report on disk."""
    rep = Report()
    rep.fail("config", detail, code="gate_error")
    doc = rep.to_dict()
    doc["error"] = detail
    write_report(work, doc, out)
    sys.stderr.write("[g2] ERROR %s\n" % detail)
    print(json.dumps(doc, indent=2))
    return EXIT_ERROR


# ── range helpers ────────────────────────────────────────────────────────────

def parse_range(spec):
    """Normalize a threshold spec to (lo, hi); either may be None. Raises ValueError."""
    if spec is None:
        return (None, None)
    if isinstance(spec, (int, float)) and not isinstance(spec, bool):
        return (float(spec), float(spec))
    if isinstance(spec, dict):
        lo, hi = spec.get("min"), spec.get("max")
        return (None if lo is None else float(lo), None if hi is None else float(hi))
    if isinstance(spec, (list, tuple)):
        if len(spec) != 2:
            raise ValueError("range must have exactly 2 entries, got %r" % (spec,))
        lo, hi = spec
        return (None if lo is None else float(lo), None if hi is None else float(hi))
    raise ValueError("unsupported range spec %r" % (spec,))


def fmt_range(lo, hi):
    return "[%s, %s]" % ("-inf" if lo is None else round(lo, 4),
                         "inf" if hi is None else round(hi, 4))


def check_scalar(report, part, island, label, value, spec):
    """value must sit inside the range. Returns True on pass."""
    try:
        lo, hi = parse_range(spec)
    except ValueError as exc:
        report.fail(part, "expect.islands[%s].%s: %s" % (island, label, exc),
                    island=island, code="bad_expect")
        return False
    if (lo is not None and value < lo) or (hi is not None and value > hi):
        report.fail(part, "%s = %s, outside expected %s" % (label, round(value, 4), fmt_range(lo, hi)),
                    island=island, code="out_of_range", value=value,
                    expected=[lo, hi], metric=label)
        return False
    return True


def check_extent(report, part, island, label, extent, spec):
    """The [min,max] extent must be CONTAINED in the range. Returns True on pass."""
    try:
        lo, hi = parse_range(spec)
    except ValueError as exc:
        report.fail(part, "expect.islands[%s].%s: %s" % (island, label, exc),
                    island=island, code="bad_expect")
        return False
    emin, emax = float(extent[0]), float(extent[1])
    bad = (lo is not None and emin < lo) or (hi is not None and emax > hi)
    if bad:
        report.fail(part, "%s extent [%s, %s] not contained in expected %s"
                    % (label, round(emin, 4), round(emax, 4), fmt_range(lo, hi)),
                    island=island, code="out_of_range",
                    value=[emin, emax], expected=[lo, hi], metric=label)
        return False
    return True


# ── census access ────────────────────────────────────────────────────────────

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


def census_index(records):
    """{island index: record}. Tolerates a dict-wrapped census ({"islands": [...]})."""
    out = {}
    for i, rec in enumerate(records):
        if not isinstance(rec, dict):
            raise ValueError("census entry %d is %s, expected an object" % (i, type(rec).__name__))
        idx = rec.get("island", rec.get("index", i))
        out[int(idx)] = rec
    return out


def resolve_mapped_key(key, parts):
    """expect.islands.mapped key -> island index. Accepts an index or a part name."""
    try:
        return int(key), None
    except (TypeError, ValueError):
        pass
    part = parts.get(key)
    if part is None:
        return None, "no island index and no part named %r in parts{}" % (key,)
    islands = part.get("islands") or []
    if len(islands) != 1:
        return None, ("part %r maps %d islands (%r) — key expect.islands.mapped by "
                      "island index instead" % (key, len(islands), islands))
    return int(islands[0]), None


def part_for_island(parts, idx):
    for name, cfg in parts.items():
        if idx in [int(i) for i in (cfg.get("islands") or [])]:
            return name
    return None


# ── main ─────────────────────────────────────────────────────────────────────

def main(argv=None):
    ap = argparse.ArgumentParser(description="G2 census sanity gate")
    ap.add_argument("--config", required=True, help="resolved config JSON (work/<name>/resolved.json)")
    ap.add_argument("--work", help="work dir (default: the config's directory)")
    ap.add_argument("--census", help="census.json override (default: <work>/census.json)")
    ap.add_argument("--out", help="report path override (default: <work>/qc/g2.json)")
    args = ap.parse_args(argv)

    work = args.work or os.path.dirname(os.path.abspath(args.config))

    try:
        cfg = load_json(args.config)
    except (OSError, ValueError) as exc:
        return die(work, "cannot read config %s: %s" % (args.config, exc), args.out)
    if not isinstance(cfg, dict):
        return die(work, "config %s is not a JSON object" % args.config, args.out)

    census_path = args.census or os.path.join(work, "census.json")
    try:
        raw = load_json(census_path)
    except OSError as exc:
        return die(work, "census.json not found — run the census stage first (%s)" % exc, args.out)
    except ValueError as exc:
        return die(work, "census.json is not valid JSON: %s" % exc, args.out)

    if isinstance(raw, dict):  # tolerate a future {"islands": [...]} wrapper
        raw = raw.get("islands", raw.get("records"))
    if not isinstance(raw, list):
        return die(work, "census.json must be a list of island records (got %s)"
                   % type(raw).__name__, args.out)

    try:
        islands = census_index(raw)
    except ValueError as exc:
        return die(work, str(exc), args.out)

    parts = cfg.get("parts") or {}
    if not isinstance(parts, dict):
        return die(work, "config parts must be an object", args.out)

    expect = ((cfg.get("expect") or {}).get("islands")) or {}
    if not isinstance(expect, dict):
        return die(work, "expect.islands must be an object", args.out)

    rep = Report()
    try:
        rel_census = os.path.relpath(census_path, work).replace("\\", "/")
    except ValueError:  # different drive
        rel_census = census_path.replace("\\", "/")
    rep.artifacts.append(rel_census)
    for name in ("qa_00_census_front.png", "qa_00_census_quarter.png"):
        if os.path.exists(os.path.join(work, name)):
            rep.artifacts.append(name)

    rep.metrics["census_islands"] = len(islands)
    rep.metrics["config_parts"] = len(parts)

    # ── 1. config parts[*].islands must exist in census.json (always) ────────
    referenced = []
    for part_name, part_cfg in sorted(parts.items()):
        if not isinstance(part_cfg, dict):
            rep.fail(part_name, "parts[%r] is not an object" % part_name, code="bad_config")
            continue
        idxs = part_cfg.get("islands")
        if idxs is None:
            rep.warn(part_name, "part declares no islands[] — nothing to cross-check",
                     code="no_islands")
            continue
        if not isinstance(idxs, list):
            rep.fail(part_name, "parts[%r].islands must be a list" % part_name, code="bad_config")
            continue
        for raw_idx in idxs:
            try:
                idx = int(raw_idx)
            except (TypeError, ValueError):
                rep.fail(part_name, "island index %r is not an integer" % (raw_idx,),
                         code="bad_config")
                continue
            referenced.append(idx)
            if idx not in islands:
                rep.fail(part_name,
                         "config maps island %d but census.json has no such island "
                         "(present: %s) — the source mesh was probably re-exported and "
                         "islands renumbered" % (idx, sorted(islands)),
                         island=idx, code="missing_island")
    rep.metrics["config_islands_referenced"] = len(referenced)

    # ── 2. island count ─────────────────────────────────────────────────────
    if "count" in expect:
        try:
            want = int(expect["count"])
        except (TypeError, ValueError):
            rep.fail("config", "expect.islands.count must be an integer, got %r"
                     % (expect["count"],), code="bad_expect")
        else:
            rep.metrics["expected_islands"] = want
            if want != len(islands):
                rep.fail("config",
                         "island count = %d, expected %d — source mesh changed "
                         "(island->part mapping is no longer trustworthy)"
                         % (len(islands), want),
                         code="island_count", value=len(islands), expected=want,
                         metric="count")

    # ── 3. per-island thresholds ────────────────────────────────────────────
    mapped = expect.get("mapped") or {}
    if not isinstance(mapped, dict):
        return die(work, "expect.islands.mapped must be an object", args.out)

    checked = 0
    for key in sorted(mapped, key=lambda k: str(k)):
        spec = mapped[key]
        if not isinstance(spec, dict):
            rep.fail("config", "expect.islands.mapped[%r] must be an object" % (key,),
                     code="bad_expect")
            continue
        idx, err = resolve_mapped_key(key, parts)
        if err:
            rep.fail("config", "expect.islands.mapped[%r]: %s" % (key, err), code="bad_expect")
            continue
        rec = islands.get(idx)
        label = part_for_island(parts, idx) or ("island_%d" % idx)
        if rec is None:
            rep.fail(label, "expect.islands.mapped[%r] targets island %d, absent from "
                            "census.json (present: %s)" % (key, idx, sorted(islands)),
                     island=idx, code="missing_island")
            continue
        checked += 1

        if "verts" in spec:
            check_scalar(rep, label, idx, "verts", float(rec.get("verts", -1)), spec["verts"])
        if "tris" in spec:
            check_scalar(rep, label, idx, "tris", float(rec.get("tris", -1)), spec["tris"])

        z = rec.get("z")
        if z and len(z) == 2:
            if "z" in spec:
                check_extent(rep, label, idx, "z", z, spec["z"])
            if "z_min" in spec:
                check_scalar(rep, label, idx, "z_min", float(z[0]), spec["z_min"])
            if "z_max" in spec:
                check_scalar(rep, label, idx, "z_max", float(z[1]), spec["z_max"])
            if "z_span" in spec:
                check_scalar(rep, label, idx, "z_span", float(z[1]) - float(z[0]), spec["z_span"])
        elif any(k in spec for k in ("z", "z_min", "z_max", "z_span")):
            rep.fail(label, "census island %d has no usable z extent (%r)" % (idx, z),
                     island=idx, code="bad_census")

        x = rec.get("x")
        if "x" in spec:
            if x and len(x) == 2:
                check_extent(rep, label, idx, "x", x, spec["x"])
            else:
                rep.fail(label, "census island %d has no usable x extent (%r)" % (idx, x),
                         island=idx, code="bad_census")

    rep.metrics["mapped_checked"] = checked

    # ── 4. unmapped islands ─────────────────────────────────────────────────
    min_verts = cfg.get("min_island_verts", 0) or 0
    unmapped = [i for i in sorted(islands)
                if i not in referenced and int(islands[i].get("verts", 0)) >= min_verts]
    rep.metrics["unmapped_islands"] = len(unmapped)
    if unmapped:
        detail = ("islands %s (>= min_island_verts %s) are in census.json but no part "
                  "claims them" % (unmapped, min_verts))
        if expect.get("require_all_mapped"):
            rep.fail("config", detail, code="unmapped_island")
        else:
            rep.warn("config", detail + " — set expect.islands.require_all_mapped to fail on this",
                     code="unmapped_island")

    if not expect:
        rep.notes.append("no expect.islands block in the config — ran the config/census "
                         "island cross-check only; add expect.islands for vert/z thresholds")

    doc = rep.to_dict()
    path = write_report(work, doc, args.out)
    if path:
        doc["report_path"] = path.replace("\\", "/")
    print(json.dumps(doc, indent=2))

    if rep.failures:
        sys.stderr.write("[g2] FAIL %d check(s)\n" % len(rep.failures))
        return EXIT_FAIL
    sys.stderr.write("[g2] PASS (%d island(s), %d mapped checked, %d warning(s))\n"
                     % (len(islands), checked, len(rep.warnings)))
    return EXIT_PASS


if __name__ == "__main__":
    sys.exit(main())