"""Hardened formal-run and three-realization contracts for OID v2.

The implementation deliberately delegates canonical publication, hashing, artifact
validation, disk checks, and the repository-wide CFD lease to CCD's shared
``run_contract`` module.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from CCD_analysis.utils.run_contract import (
    ContractError,
    IntegrityError,
    array_sha256,
    deterministic_json_bytes,
    load_complete_manifest,
    REPOSITORY_CFD_LEASE_PATH,
    validate_artifact,
)

REPO_ROOT = Path(__file__).resolve().parents[3]
SHARED_CFD_LEASE = REPOSITORY_CFD_LEASE_PATH
OID_FORMAL_SCHEMA_ID = "oid-formal-run/v2"
REALIZATION_SCHEMA_ID = "oid-three-realization/v1"
PHYSICAL_GPU = 2
LOGICAL_GPU = 0
FORMAL_REALIZATION_COUNT = 3
CANONICAL_ROLES = ("controlled", "zero_baseline", "baseline_twin")
TELEMETRY_FIELDS = ("sensors", "forces", "lattice_steps", "sample_ids", "times")
BRANCH_SCHEMA_ID = "oid-paired-branch/v1"
PAIRED_COLLECTION_SCHEMA_ID = "oid-paired-collection/v1"


def _require_text(record: Mapping[str, Any], key: str) -> str:
    value = record.get(key)
    if not isinstance(value, str) or not value.strip():
        raise ContractError(f"{key} must be a non-empty string")
    return value


def validate_realization_contract(realizations: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
    """Validate three independent realizations and their canonical branch roles.

    Each realization has its own independently initialized parent/checkpoint. Its
    controlled, zero-action baseline, and baseline-twin branches descend from that
    same checkpoint. Analysis sample intervals cannot overlap between realizations.
    """
    if len(realizations) != FORMAL_REALIZATION_COUNT:
        raise ContractError("new formal OID runs require exactly three realizations")
    ids: set[str] = set()
    parents: set[str] = set()
    checkpoints: set[str] = set()
    intervals: list[tuple[int, int, str]] = []
    canonical: list[dict[str, Any]] = []
    for raw in realizations:
        record = dict(raw)
        realization_id = _require_text(record, "realization_id")
        parent_id = _require_text(record, "parent_id")
        checkpoint_sha256 = _require_text(record, "checkpoint_sha256")
        initialization = record.get("initialization")
        phase = record.get("phase")
        if not isinstance(initialization, Mapping) or not initialization:
            raise ContractError("initialization must explicitly describe the independent parent")
        if not isinstance(phase, Mapping) or not phase:
            raise ContractError("phase must explicitly describe phase origin/selection")
        if realization_id in ids or parent_id in parents or checkpoint_sha256 in checkpoints:
            raise ContractError("realization_id, parent_id, and checkpoint must be unique across realizations")
        ids.add(realization_id); parents.add(parent_id); checkpoints.add(checkpoint_sha256)
        sample_range = record.get("sample_range")
        if (not isinstance(sample_range, Sequence) or isinstance(sample_range, (str, bytes))
                or len(sample_range) != 2):
            raise ContractError("sample_range must be [start, stop)")
        start, stop = sample_range
        if type(start) is not int or type(stop) is not int or start < 0 or stop <= start:
            raise ContractError("sample_range must be a non-empty non-negative half-open interval")
        for other_start, other_stop, other_id in intervals:
            if max(start, other_start) < min(stop, other_stop):
                raise ContractError(f"realization sample ranges overlap: {realization_id} and {other_id}")
        intervals.append((start, stop, realization_id))
        branches = record.get("branches")
        if not isinstance(branches, Mapping) or set(branches) != set(CANONICAL_ROLES):
            raise ContractError(f"branches must have exactly the canonical roles {CANONICAL_ROLES}")
        normalized_branches: dict[str, Any] = {}
        for role in CANONICAL_ROLES:
            branch = branches[role]
            if not isinstance(branch, Mapping):
                raise ContractError(f"branch {role} must be an object")
            if branch.get("role") != role:
                raise ContractError(f"branch role mismatch for {role}")
            if branch.get("parent_id") != parent_id or branch.get("checkpoint_sha256") != checkpoint_sha256:
                raise ContractError(f"branch {role} does not descend from its realization checkpoint")
            normalized_branches[role] = dict(branch)
        record["branches"] = normalized_branches
        canonical.append(record)
    return {"schema_id": REALIZATION_SCHEMA_ID, "count": len(canonical), "realizations": canonical}


def telemetry_hashes(data: Mapping[str, Any]) -> dict[str, str]:
    """Require field-level telemetry and return identity hashes for every array."""
    missing = set(TELEMETRY_FIELDS).difference(data)
    if missing:
        raise ContractError(f"telemetry missing required fields: {sorted(missing)}")
    arrays = {key: np.asarray(data[key]) for key in TELEMETRY_FIELDS}
    count = len(arrays["lattice_steps"])
    if count == 0 or any(array.ndim == 0 or len(array) != count for array in arrays.values()):
        raise ContractError("all field-level telemetry arrays must share a non-empty leading dimension")
    steps = arrays["lattice_steps"]
    sample_ids = arrays["sample_ids"]
    times = arrays["times"]
    if steps.dtype != np.dtype("int64") or np.any(steps <= 0) or np.any(np.diff(steps) <= 0):
        raise ContractError("lattice_steps must be positive strictly increasing int64")
    if sample_ids.dtype != np.dtype("int64") or np.any(np.diff(sample_ids) <= 0):
        raise ContractError("sample_ids must be strictly increasing int64")
    if times.dtype.kind not in "fc" or np.any(np.diff(times) <= 0):
        raise ContractError("times must be strictly increasing floating point")
    if not all(np.isfinite(array).all() for array in arrays.values() if array.dtype.kind in "fc"):
        raise ContractError("telemetry arrays must be finite")
    return {key: array_sha256(value) for key, value in arrays.items()}


def exact_branch_gates(branches: Mapping[str, Mapping[str, Any]]) -> dict[str, Any]:
    """Fail closed unless canonical branches are exact same-parent timeline twins."""
    if set(branches) != set(CANONICAL_ROLES):
        raise ContractError(f"branches must have exactly the canonical roles {CANONICAL_ROLES}")
    normalized = {role: dict(branches[role]) for role in CANONICAL_ROLES}
    for role, branch in normalized.items():
        if branch.get("schema_id") != BRANCH_SCHEMA_ID or branch.get("role") != role:
            raise ContractError(f"invalid canonical branch manifest for {role}")
    identity_keys = ("parent_id", "checkpoint_sha256", "timeline_sha256", "sample_ids_sha256", "times_sha256")
    reference = normalized["controlled"]
    for role, branch in normalized.items():
        for key in identity_keys:
            if not isinstance(branch.get(key), str) or branch[key] != reference.get(key):
                raise ContractError(f"branch {role} fails exact twin gate for {key}")
    if normalized["controlled"].get("action_kind") != "controlled":
        raise ContractError("controlled branch must declare controlled actions")
    for role in ("zero_baseline", "baseline_twin"):
        if normalized[role].get("action_kind") != "zero":
            raise ContractError(f"branch {role} must declare zero actions")
    zero = normalized["zero_baseline"]
    twin = normalized["baseline_twin"]
    reproducibility = {
        "byte_identical_fields": zero.get("fields_sha256") == twin.get("fields_sha256"),
        "byte_identical_telemetry": zero.get("telemetry_sha256") == twin.get("telemetry_sha256"),
    }
    # GPU floating-point runs need not be byte deterministic. Exact timelines and
    # lineage are hard gates; numerical twin-noise thresholds are evaluated from
    # the arrays by validation code and recorded before promotion.
    return {"schema_id": PAIRED_COLLECTION_SCHEMA_ID, "roles": list(CANONICAL_ROLES),
            "parent_id": reference["parent_id"], "checkpoint_sha256": reference["checkpoint_sha256"],
            "timeline_sha256": reference["timeline_sha256"], "twin_reproducibility": reproducibility,
            "requires_numerical_twin_gate": not all(reproducibility.values()), "branches": normalized}


def load_formal_run(run_dir: str | Path, *, allow_legacy_realization0: bool = False) -> dict[str, Any]:
    """Load a manifest-bound COMPLETE formal run, failing closed on any artifact.

    Existing pre-v2 formal directories are accepted only when explicitly requested;
    they are labeled legacy realization-0 evidence and are never upgraded in place.
    """
    directory = Path(run_dir)
    complete = directory / "COMPLETE.json"
    if not complete.exists():
        if allow_legacy_realization0 and directory.name.startswith("formal-") and (directory / "manifest.json").is_file():
            manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8"))
            return {"evidence_status": "legacy-realization-0", "realization_id": "0", "manifest": manifest}
        # Shared loader supplies strict state/canonical JSON diagnostics.
        load_complete_manifest(directory)
        raise AssertionError("unreachable")
    manifest = load_complete_manifest(directory)
    if manifest.get("oid_schema_id") != OID_FORMAL_SCHEMA_ID:
        raise IntegrityError("missing hardened OID formal schema")
    realization_contract = validate_realization_contract(manifest.get("realization_contract", {}).get("realizations", []))
    if manifest.get("realization_contract") != realization_contract:
        raise IntegrityError("realization contract is not canonical")
    artifacts = manifest.get("artifacts")
    if not isinstance(artifacts, Mapping) or not artifacts:
        raise IntegrityError("formal manifest must bind at least one artifact")
    for relative_name, metadata in artifacts.items():
        relative = Path(relative_name)
        if relative.is_absolute() or ".." in relative.parts or not isinstance(metadata, Mapping):
            raise IntegrityError("invalid manifest artifact entry")
        validate_artifact(directory / relative, metadata)
    timeline_hashes = manifest.get("timeline_hashes")
    array_hashes = manifest.get("array_hashes")
    if not isinstance(timeline_hashes, Mapping) or not timeline_hashes:
        raise IntegrityError("formal manifest must bind timeline hashes")
    if not isinstance(array_hashes, Mapping) or not array_hashes:
        raise IntegrityError("formal manifest must bind array hashes")
    loaded_npz: dict[str, dict[str, np.ndarray]] = {}
    for artifact_name in set(timeline_hashes) | {
            reference.rsplit(":", 1)[0] for reference in array_hashes if isinstance(reference, str) and ":" in reference}:
        if artifact_name not in artifacts or Path(artifact_name).suffix != ".npz":
            raise IntegrityError(f"hash references an unbound NPZ artifact: {artifact_name}")
        with np.load(directory / artifact_name, allow_pickle=False) as archive:
            loaded_npz[artifact_name] = {key: archive[key].copy() for key in archive.files}
    for artifact_name, expected in timeline_hashes.items():
        arrays = loaded_npz[artifact_name]
        if "lattice_steps" not in arrays or array_sha256(arrays["lattice_steps"]) != expected:
            raise IntegrityError(f"timeline hash mismatch: {artifact_name}")
    for reference, expected in array_hashes.items():
        if not isinstance(reference, str) or ":" not in reference:
            raise IntegrityError("array hash keys must be '<artifact>:<array>'")
        artifact_name, array_name = reference.rsplit(":", 1)
        arrays = loaded_npz.get(artifact_name, {})
        if array_name not in arrays or array_sha256(arrays[array_name]) != expected:
            raise IntegrityError(f"array hash mismatch: {reference}")
    # Canonical encoding is already checked by load_complete_manifest; retain this
    # explicit assertion to keep the OID loader's contract obvious.
    if complete.read_bytes() != deterministic_json_bytes(manifest):
        raise IntegrityError("formal manifest is not canonical JSON")
    return manifest
