#!/usr/bin/env python3
""" 
SFT Audit: Witness Stability Certificate (WSC)

This audit enforces the artifact-level interface required for the declared-regime
soundness/completeness contract.

Checked here (v0.2):

- required fields exist (including target_type, structured baseline, predicate ledger)
- basic margin inequalities (LCM < pi, HMR < P/2)
- stability_modulus internal consistency (when present)
- convergence certificate shape (when present)
- presence + well-formedness of the validator predicate ledger
- minimal applicability coverage of predicates by mechanism/target type
- target/baseline correctness (CTP/CHT) with structured, non-heuristic fields
- evidence references point to existing files (path-level check; anchors are not parsed)

This script performs optional machine recomputation checks when evidence payloads are present.
"""

import json
import os
import math
import re
import sys
from pathlib import Path
from typing import Any, Iterable

PRED_PAT = re.compile(r"^V[0-9]{2}$")


def fail(msg: str) -> None:
    print("FAIL:", msg)
    sys.exit(1)


def load_canonical_wob_codes() -> set[str]:
    here = Path(__file__).resolve()
    codebook = (here.parent / "../contracts/wob_codebook.json").resolve()
    if not codebook.exists():
        fail(f"Missing codebook: {codebook}")
    data = json.loads(codebook.read_text(encoding="utf-8"))
    codes: set[str] = set()
    for item in data.get("codes", []):
        c = str(item.get("code", "")).strip()
        if c:
            codes.add(c)
    if not codes:
        fail("Codebook has no codes")
    return codes


def load_predicate_map() -> dict[str, str]:
    """predicate -> canonical WOB code"""
    here = Path(__file__).resolve()
    spec = (here.parent / "../contracts/validator_predicates.json").resolve()
    if not spec.exists():
        fail(f"Missing predicate spec: {spec}")
    data = json.loads(spec.read_text(encoding="utf-8"))
    m: dict[str, str] = {}
    for p in data.get("predicates", []):
        pid = str(p.get("predicate", "")).strip()
        code = str(p.get("wob_code", "")).strip()
        if pid and code:
            m[pid] = code
    if not m:
        fail("Predicate spec has no predicates")
    return m


CANONICAL_WOB = load_canonical_wob_codes()
PRED_TO_WOB = load_predicate_map()


def project_root() -> Path:
    # engine/audit/audit_wsc.py -> engine/audit -> engine -> PROJECT_ROOT
    return Path(__file__).resolve().parents[2]


def _resolve_repo_path(root: Path, p: str) -> Path:
    """Resolve a reference path in both in-repo and extracted-bundle layouts.

    Artifacts may refer to sources under "finish_ready/..." when audited from the
    repository root. When the same artifacts are audited from the extracted
    release bundle (whose root is the former finish_ready/), those references
    should resolve after stripping the "finish_ready/" prefix.
    """
    cand = (root / p).resolve()
    if cand.exists():
        return cand
    if p.startswith("finish_ready/"):
        cand2 = (root / p[len("finish_ready/"):]).resolve()
        if cand2.exists():
            return cand2
    return cand


def _ref_to_path(ref: str) -> str:
    """Return a candidate filesystem path for a reference string.

    Convention: references may include anchors like 'path/to/file:section' or
    'path/to/file#anchor'. This audit only checks existence of the left path.
    """
    r = (ref or "").strip()
    if not r:
        return ""
    if r.startswith("http://") or r.startswith("https://") or r.startswith("doi:"):
        return ""
    # split on first ':' (used as 'file:anchor')
    if ":" in r:
        r = r.split(":", 1)[0].strip()
    # split on '#'
    if "#" in r:
        r = r.split("#", 1)[0].strip()
    return r


def _anchor_of(ref: str) -> str:
    r = (ref or "").strip()
    if "#" in r:
        return r.split("#", 1)[1].strip()
    if ":" in r:
        return r.split(":", 1)[1].strip()
    return ""


def check_anchor_exists(ref: str) -> None:
    """Optional strict check: ensure the anchor exists in the referenced file."""
    p = _ref_to_path(ref)
    a = _anchor_of(ref)
    if not p or not a:
        fail(f"strict anchors: theorem_ref must include file plus anchor: {ref}")
    root = project_root()
    fp = _resolve_repo_path(root, p)
    try:
        txt = fp.read_text(encoding="utf-8", errors="ignore")
    except Exception:
        fail(f"strict anchors: cannot read referenced file: {p}")
    if p.endswith(".tex") and (a.startswith("thm:") or a.startswith("lem:") or a.startswith("cor:") or a.startswith("def:") or a.startswith("eq:")):
        if f"\\\\label{{{a}}}" not in txt:
            fail(f"strict anchors: missing \\\\label{{{a}}} in {p}")
    else:
        if a not in txt:
            fail(f"strict anchors: anchor '{a}' not found in {p}")



def check_ref_paths(refs: Iterable[str], label: str) -> None:
    root = project_root()
    for ref in refs:
        if not isinstance(ref, str):
            fail(f"{label}: reference is not a string: {ref!r}")
        p = _ref_to_path(ref)
        if not p:
            continue
        # only treat as file path if it looks like one
        if "/" not in p and "." not in p:
            continue
        cand = _resolve_repo_path(root, p)
        if not cand.exists():
            fail(f"{label}: referenced file not found: {p}")


def _load_phase_samples(p: Path) -> list[float]:
    """
    Load phase samples from:
    - a .txt file containing one float per line, or
    - a .json file containing {"phases":[...]}.
    """
    if p.suffix.lower() == ".json":
        obj = json.loads(p.read_text(encoding="utf-8"))
        phases = obj.get("phases", [])
        if not isinstance(phases, list) or not all(isinstance(x, (int, float)) for x in phases):
            fail(f"evidence file {p}: expected phases list")
        return [float(x) for x in phases]
    lines = [ln.strip() for ln in p.read_text(encoding="utf-8").splitlines() if ln.strip()]
    try:
        return [float(x) for x in lines]
    except Exception:
        fail(f"evidence file {p}: could not parse float samples")


def _recompute_degree(samples: list[float], P: float) -> tuple[int, float, float]:
    """
    Return (degree, total_unwrapped, max_step_abs) for a loop sample sequence in [0,P).
    The unwrapping uses nearest-branch increments, producing increments in [-P/2,P/2].
    """
    if len(samples) < 2:
        return 0, 0.0, 0.0
    prev = float(samples[0])
    total = 0.0
    max_step = 0.0
    for x in samples[1:]:
        x = float(x)
        delta = x - (prev % P)
        delta -= round(delta / P) * P
        if abs(delta) > max_step:
            max_step = abs(delta)
        total += delta
        prev = prev + delta
    deg = int(round(total / P)) if P > 0 else 0
    return deg, total, max_step



def check_degradation_block(d: Any) -> None:
    """Verify declared degradation inequalities if present.

    Degradation is optional. When present, it must be explicit and numerically valid.
    """
    blk = d.get("degradation", None)
    if blk is None:
        return
    if not isinstance(blk, dict):
        fail("degradation must be an object when present")
    checks = blk.get("checks", None)
    if not isinstance(checks, list) or len(checks) == 0:
        fail("degradation.checks must be a non-empty list")
    for i, c in enumerate(checks):
        if not isinstance(c, dict):
            fail(f"degradation.checks[{i}] must be an object")
        name = str(c.get("name", "")).strip()
        rel = str(c.get("relation", "")).strip()
        if not name:
            fail(f"degradation.checks[{i}].name is required")
        if rel not in ("<=", "<"):
            fail(f"degradation.checks[{i}].relation must be '<=' or '<'")
        try:
            lhs = float(c.get("lhs"))
            rhs = float(c.get("rhs"))
        except Exception as ex:
            fail(f"degradation.checks[{i}] must include numeric lhs/rhs")
        ok = (lhs <= rhs) if rel == "<=" else (lhs < rhs)
        if not ok:
            fail(f"degradation check failed ({name}): lhs={lhs} {rel} rhs={rhs} is false")


def required_predicates_for_wsc(mechanism: str, target_type: str, witness_kind: str) -> set[str]:
    req = {"V01"}

    if mechanism in ("LDW", "HW"):
        req |= {"V02", "V04"}
        if witness_kind == "holonomy_multiple":
            req.add("V05")

    if mechanism == "AWT":
        req |= {"V08", "V09"}

    if mechanism == "CMT" or target_type == "cst":
        req.add("V06")

    if mechanism == "MXC":
        req.add("V07")
        if target_type == "cst":
            req.add("V06")

    return req


def parse_baseline(b: Any) -> dict[str, Any]:
    if not isinstance(b, dict):
        fail("baseline must be an object with fields {kind, target_type, description?}")
    kind = str(b.get("kind", "")).strip()
    ttype = str(b.get("target_type", "")).strip().lower()
    if kind not in ("haar", "coupled_haar"):
        fail(f"baseline.kind must be haar or coupled_haar (got {kind!r})")
    if ttype not in ("circle", "torus", "cst", "embedded"):
        fail(f"baseline.target_type must be circle/torus/cst/embedded (got {ttype!r})")
    desc = b.get("description", None)
    if desc is not None and not isinstance(desc, str):
        fail("baseline.description must be a string when present")
    return {"kind": kind, "target_type": ttype, "description": desc}


def check_target_baseline_compatibility(target_type: str, target: str, baseline: dict[str, Any]) -> None:
    """Artifact-level meaning rule: baseline must match the declared target.

    This is the operational form of the CTP/CHT target-correctness axiom.
    """
    bt = baseline["target_type"]
    bk = baseline["kind"]

    if bt != target_type:
        fail(f"baseline.target_type must match target_type (got {bt!r}, expected {target_type!r})")

    lt = (target or "").lower()

    # If target string declares CST, the structured target_type must agree.
    if "cst" in lt and target_type != "cst":
        fail("target string indicates CST but target_type is not cst (V06)")

    # CST regime is strict.
    if target_type == "cst":
        if bk != "coupled_haar":
            fail("CST target declared but baseline.kind is not coupled_haar (V06)")
        return

    # Non-CST regimes must not advertise coupled baselines.
    if bk == "coupled_haar":
        fail("Non-CST target declared but baseline.kind is coupled_haar (V06)")


def main() -> None:
    if len(sys.argv) != 2:
        fail("Usage: audit_wsc.py path/to/WSC.json")

    p = Path(sys.argv[1])
    if not p.exists():
        fail(f"File not found: {p}")

    data = json.loads(p.read_text(encoding="utf-8"))

    required = [
        "sft_version",
        "certificate_id",
        "witness_tier",
        "mechanism",
        "target",
        "period_P",
        "witness_value",
        "perturbation_model",
        "margins",
        "assumptions",
        "artifacts",
        "target_type",
        "baseline",
        "validator_predicates",
    ]
    missing = [k for k in required if k not in data]
    if missing:
        fail("Missing required fields: " + ", ".join(missing))

    P = float(data["period_P"])
    if not (P > 0):
        fail("period_P must be > 0")

    mechanism = str(data.get("mechanism", "")).strip()
    target = str(data.get("target", "")).strip()

    target_type = str(data.get("target_type", "")).strip().lower()
    if target_type not in ("circle", "torus", "cst", "embedded"):
        fail(f"target_type must be one of circle/torus/cst/embedded. Got: {target_type}")

    baseline = parse_baseline(data.get("baseline"))

    # --- basic margin checks ---
    margins = data.get("margins", {}) or {}

    if "LCM_pi" in margins:
        lcm = float(margins["LCM_pi"])
        if not (0 < lcm < math.pi):
            fail(f"LCM_pi must satisfy 0 < LCM_pi < pi. Got {lcm}")

    if "HMR_P" in margins:
        hmr = float(margins["HMR_P"])
        if not (0 < hmr < P / 2):
            fail(f"HMR_P must satisfy 0 < HMR_P < P/2. Got {hmr}, P={P}")

    # If holonomy is invoked, HMR_P must be present.
    wv = data.get("witness_value", {}) or {}
    wv_kind = str(wv.get("kind", "")).strip()
    if wv_kind == "holonomy_multiple" and "HMR_P" not in margins:
        fail("holonomy_multiple witness requires margins.HMR_P (V05)")

    # --- validator predicate ledger checks ---
    led = data.get("validator_predicates")
    if not isinstance(led, list) or len(led) == 0:
        fail("validator_predicates must be a non-empty list")

    seen: set[str] = set()
    for i, it in enumerate(led):
        if not isinstance(it, dict):
            fail(f"validator_predicates[{i}] must be an object")
        pid = str(it.get("predicate", "")).strip()
        if not PRED_PAT.match(pid):
            fail(f"validator_predicates[{i}].predicate must match V##. Got: {pid}")
        if pid not in PRED_TO_WOB:
            fail(f"Unknown predicate id in ledger: {pid}")
        if pid in seen:
            fail(f"Duplicate predicate in ledger: {pid}")
        seen.add(pid)

        holds = it.get("holds", None)
        if holds is not True:
            fail(f"validator_predicates[{i}].holds must be true for WSC (got {holds})")

        wob_code = str(it.get("wob_code", "")).strip()
        if wob_code != PRED_TO_WOB[pid]:
            fail(
                f"validator_predicates[{i}].wob_code mismatch for {pid}: "
                f"got {wob_code}, expected {PRED_TO_WOB[pid]}"
            )
        if wob_code not in CANONICAL_WOB:
            fail(f"validator_predicates[{i}].wob_code not in canonical catalog: {wob_code}")

        ev = it.get("evidence_refs", None)
        if not isinstance(ev, list) or len(ev) == 0 or not all(isinstance(x, str) and x.strip() for x in ev):
            fail(f"validator_predicates[{i}].evidence_refs must be a non-empty list of strings")
        check_ref_paths(ev, f"validator_predicates[{i}].evidence_refs")

    req = required_predicates_for_wsc(mechanism, target_type, wv_kind)
    missing_req = sorted(list(req - seen))
    if missing_req:
        fail(
            "validator_predicates missing required entries "
            f"for mechanism={mechanism}, target_type={target_type}: {missing_req}"
        )

    # --- target/baseline correctness (artifact-level meaning) ---
    check_target_baseline_compatibility(target_type, target, baseline)

    # --- stability modulus consistency (optional) ---
    sm = data.get("stability_modulus", None)
    if sm is not None:
        if not isinstance(sm, dict):
            fail("stability_modulus must be an object when present")

        # check k against witness_value.integer if both exist
        if "k" in sm:
            k = int(sm["k"])
            if "integer" in wv:
                w_int = int(wv["integer"])
                if k != w_int:
                    fail(f"stability_modulus.k must equal witness_value.integer (got {k}, expected {w_int})")

        if "LCM_pi" in sm:
            lcm2 = float(sm["LCM_pi"])
            if not (0 < lcm2 < math.pi):
                fail(f"stability_modulus.LCM_pi must satisfy 0 < LCM_pi < pi. Got {lcm2}")

        if "HMR_P" in sm:
            hmr2 = float(sm["HMR_P"])
            if not (0 < hmr2 < P / 2):
                fail(f"stability_modulus.HMR_P must satisfy 0 < HMR_P < P/2. Got {hmr2}, P={P}")

    # --- optional degradation inequalities ---
    check_degradation_block(data)

    # --- artifacts references: path-level existence check ---
    art = data.get("artifacts", {}) or {}
    theorem_refs = art.get("theorem_refs", []) or []
    files = art.get("files", []) or []
    if not isinstance(theorem_refs, list) or not all(isinstance(x, str) for x in theorem_refs):
        fail("artifacts.theorem_refs must be a list of strings")
    if len([x for x in theorem_refs if str(x).strip()]) == 0:
        fail("artifacts.theorem_refs must be a non-empty list")
    if not isinstance(files, list) or not all(isinstance(x, str) for x in files):
        fail("artifacts.files must be a list of strings")
    if len([x for x in files if str(x).strip()]) == 0:
        fail("artifacts.files must be a non-empty list")

    check_ref_paths(theorem_refs, "artifacts.theorem_refs")
    check_ref_paths(files, "artifacts.files")
    evidence_files = art.get("evidence_files", []) or []
    if evidence_files:
        if not isinstance(evidence_files, list) or not all(isinstance(x, str) for x in evidence_files):
            fail("artifacts.evidence_files must be a list of strings when present")
        check_ref_paths(evidence_files, "artifacts.evidence_files")

        sm = data.get("stability_modulus", {}) or {}
        witness_kind = str(sm.get("witness_kind", "")).strip().lower()

        if str(data.get("mechanism", "")).strip().upper() == "LDW" and witness_kind == "degree":
            P = float(data.get("period_P"))
            # use the first evidence payload as the canonical sample source
            ev_path = _ref_to_path(evidence_files[0])
            if not ev_path:
                fail("artifacts.evidence_files contains an unparseable path reference")
            root = project_root()
            p = (root / ev_path).resolve()
            samples = _load_phase_samples(p)
            deg, total, max_step = _recompute_degree(samples, P)

            wv = data.get("witness_value", {}) or {}
            want = int(wv.get("integer"))
            if deg != want:
                fail(f"degree recomputation mismatch: computed {deg}, certificate witness_value {want}")

            k_sm = sm.get("k", None)
            if k_sm is not None and int(k_sm) != want:
                fail(f"degree recomputation mismatch: stability_modulus.k {int(k_sm)} vs witness_value {want}")

            # local step bound check against LCM_pi when present
            margins = data.get("margins", {}) or {}
            lcm = margins.get("LCM_pi", None)
            if lcm is not None:
                lcm = float(lcm)
                if max_step > lcm + 1e-12:
                    fail(f"degree evidence violates LCM_pi: max_step={max_step} > LCM_pi={lcm}")

    strict = (os.environ.get("SFT_AUDIT_STRICT_ANCHORS", "").strip() == "1")
    if strict:
        for ref in theorem_refs:
            check_anchor_exists(ref)

    print("OK:", p.name)


if __name__ == "__main__":
    main()
