#!/usr/bin/env python3
"""Audit staged round-one SR artifacts without mutating canonical history."""
from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable, Mapping

import numpy as np

SR_ROOT = Path(__file__).resolve().parents[1]
REPO_ROOT = SR_ROOT.parents[1]
if str(REPO_ROOT / "src") not in sys.path:
    sys.path.insert(0, str(REPO_ROOT / "src"))

from SR_analysis.utils.provenance import atomic_write_json, hash_file, hash_json, sha256_text

AUDIT_SCHEMA_VERSION = "sr-round1-audit-v1"
CANDIDATE_SCHEMA_VERSION = "sr-candidate-registry-v1"
LINK_HASH_KEYS = {
    "npz_hash", "controlled_hash", "target_hash", "telemetry_hash", "validation_hash",
}
# Composite hashes are validated by their producing schema, not as file links.
PROVENANCE_KEYS = ("git_sha", "environment", "command", "created_at")
IGNORED_JSON = {"candidate_registry.json", "audit_report.json"}


@dataclass(frozen=True)
class Artifact:
    kind: str
    path: Path
    relative_path: str
    digest: str
    document: Mapping[str, Any] | None = None


def _issue(code: str, message: str, path: str | None = None) -> dict[str, Any]:
    result: dict[str, Any] = {"code": code, "message": message}
    if path is not None:
        result["path"] = path
    return result


def _read_json(path: Path) -> Mapping[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError("JSON root must be an object")
    return value


def _classify_json(document: Mapping[str, Any], path: Path) -> str:
    if "deployment_expression" in document or "fitted_expression" in document:
        return "formula"
    if "record_hash" in document or ("metric_version" in document and "telemetry" in document):
        return "validation"
    if "telemetry" in path.name.lower() or document.get("artifact_type") == "telemetry":
        return "telemetry"
    return str(document.get("artifact_type", "metadata"))


def _artifact_id(artifact: Artifact) -> str:
    if artifact.document:
        value = artifact.document.get("artifact_id")
        if isinstance(value, str) and value.strip():
            return value
    return f"{artifact.kind}:{artifact.digest[:16]}"


def _path_value(value: Any) -> str | None:
    if isinstance(value, str) and value:
        return value
    if isinstance(value, Mapping):
        for key in ("path", "file", "raw_path", "telemetry_path"):
            candidate = value.get(key)
            if isinstance(candidate, str) and candidate:
                return candidate
    return None


def _resolve_reference(reference: str, owner: Path, roots: Iterable[Path]) -> Path | None:
    candidate = Path(reference).expanduser()
    attempts = [candidate] if candidate.is_absolute() else [owner.parent / candidate]
    if not candidate.is_absolute():
        attempts.extend(root / candidate for root in roots)
    for attempt in attempts:
        resolved = attempt.resolve()
        if resolved.is_file():
            return resolved
    return None


def _walk_hash_links(value: Any, prefix: str = "") -> Iterable[tuple[str, str]]:
    if isinstance(value, Mapping):
        for key, child in value.items():
            name = f"{prefix}.{key}" if prefix else str(key)
            if key in LINK_HASH_KEYS and isinstance(child, str):
                yield name, child
            yield from _walk_hash_links(child, name)
    elif isinstance(value, list):
        for index, child in enumerate(value):
            yield from _walk_hash_links(child, f"{prefix}[{index}]")


def _validate_npz(artifact: Artifact, errors: list[dict[str, Any]]) -> None:
    try:
        with np.load(artifact.path, allow_pickle=False) as data:
            keys = set(data.files)
            if "schema_version" not in keys:
                errors.append(_issue("npz_schema_missing", "NPZ lacks schema_version", artifact.relative_path))
                return
            schema = str(np.asarray(data["schema_version"]).item())
            if not schema.strip():
                errors.append(_issue("npz_schema_invalid", "NPZ schema_version is empty", artifact.relative_path))
            if artifact.path.name == "controlled.npz":
                required = {"sensors_post", "forces_post", "actions_norm", "dt_c"}
                missing = sorted(required - keys)
                if missing:
                    errors.append(_issue("npz_fields_missing", f"controlled NPZ missing {missing}", artifact.relative_path))
                elif (data["sensors_post"].ndim != 2 or data["sensors_post"].shape[1] != 6
                      or data["forces_post"].shape != data["sensors_post"].shape
                      or data["actions_norm"].shape != (data["sensors_post"].shape[0], 3)):
                    errors.append(_issue("npz_shape_invalid", "controlled NPZ array shapes violate the v2 contract", artifact.relative_path))
            if "target" in artifact.path.name and not ({"target_states", "layout"} <= keys):
                errors.append(_issue("npz_fields_missing", "target NPZ requires target_states and layout", artifact.relative_path))
    except Exception as exc:
        errors.append(_issue("npz_unreadable", str(exc), artifact.relative_path))


def _missing_provenance(provenance: Any) -> list[str]:
    if not isinstance(provenance, Mapping):
        return list(PROVENANCE_KEYS)
    return [key for key in PROVENANCE_KEYS if not provenance.get(key)]


def _validate_formula(artifact: Artifact, errors: list[dict[str, Any]]) -> None:
    doc = artifact.document or {}
    required = {
        "schema_version", "artifact_id", "run_id", "feature_names", "fitted_expression",
        "deployment_expression", "fitted_expression_hash", "deployment_expression_hash",
        "alignment_metadata", "dt_c_metadata", "g_metadata", "provenance",
    }
    missing = sorted(key for key in required if key not in doc)
    if missing:
        errors.append(_issue("formula_schema_missing", f"formula missing {missing}", artifact.relative_path))
        return
    for expression_key in ("fitted_expression", "deployment_expression"):
        expression = doc.get(expression_key)
        hash_key = f"{expression_key}_hash"
        if not isinstance(expression, str) or not expression.strip():
            errors.append(_issue("formula_expression_invalid", f"{expression_key} must be non-empty", artifact.relative_path))
        else:
            actual = sha256_text(expression.strip())
            if doc.get(hash_key) != actual:
                errors.append(_issue("expression_hash_mismatch", f"{hash_key} does not hash the actual expression", artifact.relative_path))
    provenance = doc.get("provenance")
    if not isinstance(provenance, Mapping) or not all(isinstance(provenance.get(k), Mapping) for k in ("fit", "pysr", "data")):
        errors.append(_issue("formula_provenance_missing", "formula provenance requires fit, pysr, and data objects", artifact.relative_path))
    else:
        fit_required = ("mode", "seed", "niterations", "scene_sample_counts", "diagnostic_metrics")
        missing_fit = [key for key in fit_required if key not in provenance["fit"]]
        if missing_fit or not doc.get("git_sha") or not isinstance(doc.get("environment"), Mapping):
            errors.append(_issue(
                "required_provenance_missing",
                f"formula provenance lacks top-level git/environment or fit fields {missing_fit}",
                artifact.relative_path,
            ))


def _validate_validation(
    artifact: Artifact, errors: list[dict[str, Any]], roots: Iterable[Path]
) -> None:
    doc = artifact.document or {}
    required = {"schema_version", "formula_hash", "metric_version", "scene", "mode", "n_steps", "metrics", "telemetry", "provenance", "record_hash"}
    missing = sorted(key for key in required if key not in doc)
    if missing:
        errors.append(_issue("validation_schema_missing", f"validation missing {missing}", artifact.relative_path))
        return
    expected_record_hash = hash_json({key: value for key, value in doc.items() if key != "record_hash"})
    if doc.get("record_hash") != expected_record_hash:
        errors.append(_issue("record_hash_mismatch", "validation record_hash is invalid", artifact.relative_path))
    if _missing_provenance(doc.get("provenance")):
        errors.append(_issue("required_provenance_missing", "validation provenance lacks git_sha/environment/command/created_at", artifact.relative_path))
    telemetry_ref = _path_value(doc.get("telemetry"))
    if telemetry_ref is None:
        errors.append(_issue("telemetry_reference_missing", "validation must reference raw telemetry", artifact.relative_path))
    elif _resolve_reference(telemetry_ref, artifact.path, roots) is None:
        errors.append(_issue("telemetry_missing", f"telemetry does not exist: {telemetry_ref}", artifact.relative_path))


def _discover(data_dir: Path, results_dir: Path, errors: list[dict[str, Any]]) -> list[Artifact]:
    artifacts: list[Artifact] = []
    for root, prefix in ((data_dir, "data"), (results_dir, "results")):
        if not root.is_dir():
            errors.append(_issue("run_directory_missing", f"required run directory does not exist: {root}"))
            continue
        for path in sorted(item for item in root.rglob("*") if item.is_file()):
            relative = f"{prefix}/{path.relative_to(root).as_posix()}"
            if path.name in IGNORED_JSON:
                continue
            if path.suffix == ".npz":
                kind = "telemetry" if any("telemetry" in part.lower() for part in path.parts) else "npz"
                artifacts.append(Artifact(kind, path, relative, hash_file(path)))
            elif path.suffix == ".json":
                try:
                    document = _read_json(path)
                    artifacts.append(Artifact(_classify_json(document, path), path, relative, hash_file(path), document))
                except Exception as exc:
                    errors.append(_issue("json_unreadable", str(exc), relative))
            else:
                artifacts.append(Artifact("telemetry" if "telemetry" in path.name.lower() else "artifact", path, relative, hash_file(path)))
    return artifacts


def _historical_snapshot(sr_root: Path) -> dict[str, Any]:
    registry = sr_root / "scene_registry.json"
    validations = sr_root / "results" / "validations"
    return {
        "status": "historical_frozen",
        "canonical_registry": {"path": "scene_registry.json", "sha256": hash_file(registry) if registry.is_file() else None},
        "validation_files": sorted(str(path.relative_to(sr_root)) for path in validations.glob("*.json")) if validations.is_dir() else [],
        "known_conflicts": [
            "Illusion similarity numbers differ between historical README prose/tables and scene_registry.json; round one records both and does not choose a winner.",
            "Karman fitted expressions and deployment expressions differ in historical artifacts/documentation; round one records both and does not choose a winner.",
        ],
        "adjudication": "none",
    }


def audit_run(
    run_id: str,
    data_run_id: str | None = None,
    *,
    formula_run_id: str | None = None,
    sr_root: Path = SR_ROOT,
) -> tuple[dict[str, Any], dict[str, Any]]:
    if not run_id or Path(run_id).name != run_id or run_id in {".", ".."}:
        raise ValueError("run_id must be one non-empty path component")
    data_run_id = data_run_id or run_id
    if not data_run_id or Path(data_run_id).name != data_run_id or data_run_id in {".", ".."}:
        raise ValueError("data_run_id must be one non-empty path component")
    formula_run_id = formula_run_id or run_id
    if not formula_run_id or Path(formula_run_id).name != formula_run_id or formula_run_id in {".", ".."}:
        raise ValueError("formula_run_id must be one non-empty path component")
    data_dir = sr_root / "data" / "runs" / data_run_id
    results_dir = sr_root / "results" / "runs" / run_id
    formula_dir = sr_root / "results" / "runs" / formula_run_id
    results_dir.mkdir(parents=True, exist_ok=True)
    errors: list[dict[str, Any]] = []
    warnings: list[dict[str, Any]] = []
    artifacts = _discover(data_dir, results_dir, errors)
    if formula_dir.resolve() != results_dir.resolve():
        formula_artifacts = _discover(data_dir, formula_dir, errors)
        artifacts.extend(artifact for artifact in formula_artifacts if artifact.kind == "formula")
    roots = (data_dir, results_dir, formula_dir, sr_root)
    digest_index = {artifact.digest: artifact for artifact in artifacts}

    for artifact in artifacts:
        if artifact.kind == "npz" or (artifact.kind == "telemetry" and artifact.path.suffix == ".npz"):
            _validate_npz(artifact, errors)
        elif artifact.kind == "formula":
            _validate_formula(artifact, errors)
        elif artifact.kind == "validation":
            _validate_validation(artifact, errors, roots)
        elif artifact.kind == "telemetry" and artifact.document is not None:
            if not artifact.document.get("schema_version"):
                errors.append(_issue("telemetry_schema_missing", "telemetry JSON lacks schema_version", artifact.relative_path))
        if artifact.document:
            for key, digest in _walk_hash_links(artifact.document):
                if len(digest) == 64 and digest not in digest_index:
                    errors.append(_issue("hash_link_broken", f"{key} references absent artifact hash {digest}", artifact.relative_path))

    kinds = {artifact.kind for artifact in artifacts}
    for kind in ("npz", "formula", "validation", "telemetry"):
        if kind not in kinds:
            errors.append(_issue("artifact_kind_missing", f"run has no {kind} artifact"))

    audit_passed = not errors
    validation_artifacts = [artifact for artifact in artifacts if artifact.kind == "validation"]
    validations_complete = bool(validation_artifacts) and all(
        artifact.document is not None and artifact.document.get("status") not in {"failed", "incomplete"}
        for artifact in validation_artifacts
    )
    chain_scientifically_eligible = audit_passed and validations_complete
    candidates = [
        {
            "artifact_id": _artifact_id(artifact),
            "path": artifact.relative_path,
            "sha256": artifact.digest,
            "run_id": run_id,
            "status": (
                "candidate"
                if chain_scientifically_eligible
                else "rejected"
            ),
        }
        for artifact in artifacts
        if artifact.kind in {"formula", "validation"}
    ]
    candidate_registry = {
        "schema_version": CANDIDATE_SCHEMA_VERSION,
        "run_id": run_id,
        "status": "pass" if audit_passed else "fail",
        "promotion_eligible": False,
        "scientifically_eligible": chain_scientifically_eligible,
        "artifacts": candidates,
    }
    report = {
        "schema_version": AUDIT_SCHEMA_VERSION,
        "run_id": run_id,
        "data_run_id": data_run_id,
        "formula_run_id": formula_run_id,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "status": "pass" if not errors else "fail",
        "strict_failure": bool(errors),
        "summary": {"artifact_count": len(artifacts), "candidate_count": len(candidates), "error_count": len(errors), "warning_count": len(warnings)},
        "errors": errors,
        "warnings": warnings,
        "artifacts": [{"kind": a.kind, "path": a.relative_path, "sha256": a.digest} for a in artifacts],
        "historical": _historical_snapshot(sr_root),
        "promotion": {"automatic": False, "required_action": "manual review before any canonical registry edit"},
        "scope_exclusions": ["Vortex", "V5", "Stage 4 figure generation"],
    }
    atomic_write_json(results_dir / "candidate_registry.json", candidate_registry)
    atomic_write_json(results_dir / "audit_report.json", report)
    return report, candidate_registry


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-id", required=True, help="Stage 2/3 results run identifier")
    parser.add_argument("--data-run-id", help="Stage 1 data run identifier; defaults to --run-id")
    parser.add_argument("--formula-run-id", help="Stage 2 formula run identifier; defaults to --run-id")
    parser.add_argument("--strict", action="store_true", help="return non-zero when any schema or hash chain check fails")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    report, _ = audit_run(
        args.run_id, args.data_run_id, formula_run_id=args.formula_run_id, sr_root=SR_ROOT,
    )
    print(json.dumps(report["summary"], sort_keys=True))
    print(f"audit report: results/runs/{args.run_id}/audit_report.json")
    return 1 if args.strict and report["strict_failure"] else 0


if __name__ == "__main__":
    raise SystemExit(main())
