execution_metadata.py

"""Canonical, optional runtime transparency metadata for model packages.

The legacy training tuple and required model-manifest fields intentionally stay
unchanged.  New producers place this validated sidecar under the optional
``execution_metadata`` manifest key; consumers must continue to accept packages
where that key is absent.
"""

from __future__ import annotations

import json
import os
import re
import tempfile
import zipfile
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from typing import Any, Mapping, Sequence

from ui_contracts import (
    DecisionSource,
    DecisionTrace,
    ExperimentReflection,
    FeatureExtractionReport,
    ResearchPlan,
)


EXECUTION_METADATA_KEY = "execution_metadata"
EXECUTION_METADATA_SCHEMA_VERSION = 1
APPROVED_RESEARCH_KEY = "approved_research"
TRUSTED_EXTRACTOR_KEY = "trusted_extractor"
EXPERIMENT_REFLECTION_KEY = "experiment_reflection"
_EXTRACTOR_ID_PATTERN = re.compile(r"^[a-z][a-z0-9_]{2,63}$")
_EXTRACTOR_VERSION_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._+-]{0,63}$")
_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")
_TRUSTED_EXTRACTOR_FIELDS = {
    "schema_version",
    "extractor_id",
    "version",
    "configuration",
    "bundle_hash",
    "required_assets",
}
_TRUSTED_ASSET_FIELDS = {"asset_id", "sha256", "size_bytes"}
_APPROVED_RESEARCH_FIELDS = {
    "schema_version",
    "research_plan",
    "approved_candidate_id",
    "pipeline_config",
    "evidence_ids",
    "approval_trace",
}
_PIPELINE_CONFIG_FIELDS = {
    "schema_version",
    "candidate_id",
    "task_type",
    "data_modality",
    "feature_selection_tool",
    "transform_tool",
    "ml_tool",
    "image_backend",
    "image_backbone",
    "extractor_configuration",
}


def _mapping(value: Any, name: str) -> Mapping[str, Any]:
    if not isinstance(value, Mapping):
        raise ValueError(f"{name} must be a mapping")
    return value


def make_feature_extraction_report(
    *,
    report_id: str,
    modality: str,
    method: str,
    input_shape: Sequence[int] = (),
    output_shape: Sequence[int] = (),
    features_per_sample: int | None = None,
    input_feature_count: int | None = None,
    output_feature_count: int | None = None,
    selected_features: Sequence[str] = (),
    generated_features: Sequence[str] = (),
    dropped_features: Sequence[str] = (),
    parameters: Mapping[str, Any] | None = None,
    warnings: Sequence[str] = (),
    provenance: Mapping[str, Any] | None = None,
    extensions: Mapping[str, Any] | None = None,
    status: str = "completed",
) -> FeatureExtractionReport:
    """Build one schema-validated report from actual producer values."""

    return FeatureExtractionReport.from_mapping(
        {
            "schema_version": 1,
            "report_id": report_id,
            "status": status,
            "modality": modality,
            "method": method,
            "input_shape": list(input_shape),
            "output_shape": list(output_shape),
            "features_per_sample": features_per_sample,
            "input_feature_count": input_feature_count,
            "output_feature_count": output_feature_count,
            "selected_features": list(selected_features),
            "generated_features": list(generated_features),
            "dropped_features": list(dropped_features),
            "parameters": dict(parameters or {}),
            "warnings": list(warnings),
            "provenance": {"phase": "execution", **dict(provenance or {})},
            "extensions": dict(extensions or {}),
        }
    )


def make_decision_trace(
    *,
    trace_id: str,
    decision_type: str,
    stage: str,
    source: str,
    selected: Any,
    options: Sequence[Any] = (),
    rationale: str,
    parameters: Mapping[str, Any] | None = None,
    actor: str = "runtime",
    context_summary: str = "",
    evidence_ids: Sequence[str] = (),
    warnings: Sequence[str] = (),
    provenance: Mapping[str, Any] | None = None,
    extensions: Mapping[str, Any] | None = None,
    timestamp_utc: str | None = None,
) -> DecisionTrace:
    """Build a reviewable trace while rejecting hidden internal reasoning."""

    return DecisionTrace.from_mapping(
        {
            "schema_version": 1,
            "trace_id": trace_id,
            "decision_type": decision_type,
            "stage": stage,
            "decision_source": source,
            "actor": actor,
            "context_summary": context_summary,
            "options": list(options),
            "selected": selected,
            "rationale": rationale,
            "parameters": dict(parameters or {}),
            "evidence_ids": list(evidence_ids),
            "warnings": list(warnings),
            "provenance": {"phase": "execution", **dict(provenance or {})},
            "timestamp_utc": timestamp_utc
            or datetime.now(timezone.utc).isoformat(),
            "extensions": dict(extensions or {}),
        }
    )


def make_execution_metadata(
    feature_report: FeatureExtractionReport | Mapping[str, Any] | None = None,
    decision_traces: Sequence[DecisionTrace | Mapping[str, Any]] = (),
    *,
    approved_research: Mapping[str, Any] | None = None,
    trusted_extractor: Mapping[str, Any] | None = None,
    experiment_reflection: ExperimentReflection | Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Return the canonical JSON-safe execution metadata container."""

    report = (
        None
        if feature_report is None
        else feature_report
        if isinstance(feature_report, FeatureExtractionReport)
        else FeatureExtractionReport.from_mapping(
            _mapping(feature_report, "feature_extraction_report")
        )
    )
    traces = [
        trace
        if isinstance(trace, DecisionTrace)
        else DecisionTrace.from_mapping(_mapping(trace, "decision trace"))
        for trace in decision_traces
    ]
    approved = (
        normalize_approved_research(approved_research)
        if approved_research
        else {}
    )
    extractor_binding = (
        normalize_trusted_extractor_binding(trusted_extractor)
        if trusted_extractor
        else {}
    )
    reflection = (
        None
        if experiment_reflection is None
        else experiment_reflection
        if isinstance(experiment_reflection, ExperimentReflection)
        else ExperimentReflection.from_mapping(
            _mapping(experiment_reflection, EXPERIMENT_REFLECTION_KEY)
        )
    )
    if approved:
        approval_trace = DecisionTrace.from_mapping(approved["approval_trace"])
        if all(trace.trace_id != approval_trace.trace_id for trace in traces):
            traces.insert(0, approval_trace)
    if reflection is not None and reflection.decision_trace is not None:
        if all(
            trace.trace_id != reflection.decision_trace.trace_id for trace in traces
        ):
            traces.append(reflection.decision_trace)
    if (
        report is None
        and not traces
        and not approved
        and not extractor_binding
        and reflection is None
    ):
        return {}
    result = {
        "schema_version": EXECUTION_METADATA_SCHEMA_VERSION,
        "feature_extraction_report": report.to_dict() if report else None,
        "decision_traces": [trace.to_dict() for trace in traces],
    }
    if approved:
        result[APPROVED_RESEARCH_KEY] = approved
    if extractor_binding:
        result[TRUSTED_EXTRACTOR_KEY] = extractor_binding
    if reflection is not None:
        result[EXPERIMENT_REFLECTION_KEY] = reflection.to_dict()
    return result


def normalize_trusted_extractor_binding(value: Mapping[str, Any]) -> dict[str, Any]:
    """Validate a package-safe binding without treating it as executable authority."""

    data = _mapping(value, TRUSTED_EXTRACTOR_KEY)
    if set(data) != _TRUSTED_EXTRACTOR_FIELDS or data.get("schema_version") != 1:
        raise ValueError("trusted_extractor has an unsupported shape")
    extractor_id = str(data.get("extractor_id") or "").strip().lower()
    version = str(data.get("version") or "").strip()
    bundle_hash = str(data.get("bundle_hash") or "").strip().lower()
    if not _EXTRACTOR_ID_PATTERN.fullmatch(extractor_id):
        raise ValueError("trusted_extractor.extractor_id is invalid")
    if not _EXTRACTOR_VERSION_PATTERN.fullmatch(version):
        raise ValueError("trusted_extractor.version is invalid")
    if not _SHA256_PATTERN.fullmatch(bundle_hash):
        raise ValueError("trusted_extractor.bundle_hash is invalid")
    configuration = data.get("configuration")
    if not isinstance(configuration, Mapping) or not all(
        isinstance(key, str) for key in configuration
    ):
        raise ValueError("trusted_extractor.configuration must be a mapping")
    try:
        configuration_copy = json.loads(
            json.dumps(dict(configuration), allow_nan=False)
        )
    except (TypeError, ValueError, OverflowError, RecursionError) as exc:
        raise ValueError("trusted_extractor.configuration must be JSON-safe") from exc
    if len(json.dumps(configuration_copy).encode("utf-8")) > 32_768:
        raise ValueError("trusted_extractor.configuration is too large")
    raw_assets = data.get("required_assets")
    if not isinstance(raw_assets, (list, tuple)) or len(raw_assets) > 64:
        raise ValueError("trusted_extractor.required_assets must be a bounded array")
    assets: list[dict[str, Any]] = []
    asset_ids: set[str] = set()
    for raw in raw_assets:
        asset = _mapping(raw, "trusted extractor asset")
        if set(asset) != _TRUSTED_ASSET_FIELDS:
            raise ValueError("trusted extractor asset has an unsupported shape")
        asset_id = str(asset.get("asset_id") or "").strip()
        digest = str(asset.get("sha256") or "").strip().lower()
        size = asset.get("size_bytes")
        if (
            not asset_id
            or len(asset_id) > 128
            or asset_id in {".", ".."}
            or "/" in asset_id
            or "\\" in asset_id
            or any(ord(character) < 32 for character in asset_id)
            or asset_id in asset_ids
        ):
            raise ValueError("trusted extractor asset_id is invalid or duplicated")
        if not _SHA256_PATTERN.fullmatch(digest):
            raise ValueError("trusted extractor asset sha256 is invalid")
        if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
            raise ValueError("trusted extractor asset size_bytes must be positive")
        asset_ids.add(asset_id)
        assets.append({"asset_id": asset_id, "sha256": digest, "size_bytes": size})
    return {
        "schema_version": 1,
        "extractor_id": extractor_id,
        "version": version,
        "configuration": configuration_copy,
        "bundle_hash": bundle_hash,
        "required_assets": assets,
    }


def normalize_approved_research(value: Mapping[str, Any]) -> dict[str, Any]:
    """Validate the optional approved plan carried into trusted execution."""

    data = _mapping(value, APPROVED_RESEARCH_KEY)
    if set(data) - _APPROVED_RESEARCH_FIELDS:
        raise ValueError("approved_research contains unsupported fields")
    if data.get("schema_version", 1) != 1:
        raise ValueError("unsupported approved_research schema_version")
    plan = ResearchPlan.from_mapping(
        _mapping(data.get("research_plan"), "approved_research.research_plan")
    )
    if not 1 <= len(plan.candidates) <= 3:
        raise ValueError("approved research plan must contain one to three candidates")
    approved_id = str(data.get("approved_candidate_id") or "").strip()
    candidate = next(
        (item for item in plan.candidates if item.candidate_id == approved_id),
        None,
    )
    if candidate is None:
        raise ValueError("approved candidate is not present in the research plan")
    config = _mapping(
        data.get("pipeline_config"), "approved_research.pipeline_config"
    )
    if not all(isinstance(key, str) for key in config):
        raise ValueError("approved pipeline config keys must be strings")
    if set(config) - _PIPELINE_CONFIG_FIELDS:
        raise ValueError("approved pipeline config contains unsupported fields")
    if config.get("schema_version", 1) != 1:
        raise ValueError("unsupported approved pipeline config schema_version")
    if config.get("candidate_id") not in (None, approved_id):
        raise ValueError("approved pipeline config references another candidate")
    try:
        config_copy = json.loads(json.dumps(dict(config), allow_nan=False))
    except (TypeError, ValueError, RecursionError) as exc:
        raise ValueError("approved pipeline config must be JSON-safe") from exc
    evidence_ids = data.get("evidence_ids") or ()
    if not isinstance(evidence_ids, (list, tuple)) or not all(
        isinstance(item, str) and 0 < len(item) <= 200 for item in evidence_ids
    ):
        raise ValueError("approved research evidence_ids must be an array of strings")
    if len(evidence_ids) != len(set(evidence_ids)):
        raise ValueError("approved research evidence_ids must be unique")
    known_evidence = {
        str(item.get("evidence_id"))
        for item in plan.evidence
        if isinstance(item, Mapping) and item.get("evidence_id")
    }
    if not set(evidence_ids) <= known_evidence or not set(evidence_ids) <= set(
        candidate.evidence_ids
    ):
        raise ValueError("approved research references unknown candidate evidence")
    trace = DecisionTrace.from_mapping(
        _mapping(data.get("approval_trace"), "approved_research.approval_trace")
    )
    if trace.decision_source is not DecisionSource.RESEARCH_AGENT:
        raise ValueError("approval trace must come from the research agent")
    if trace.selected != approved_id:
        raise ValueError("approval trace does not match the approved candidate")
    return {
        "schema_version": 1,
        "research_plan": plan.to_dict(),
        "approved_candidate_id": approved_id,
        "pipeline_config": config_copy,
        "evidence_ids": list(evidence_ids),
        "approval_trace": trace.to_dict(),
    }


def normalize_execution_metadata(value: Mapping[str, Any]) -> dict[str, Any]:
    """Validate and canonicalize metadata received from a producer/package."""

    data = _mapping(value, EXECUTION_METADATA_KEY)
    if not data:
        return {}
    version = data.get("schema_version", 1)
    if version != EXECUTION_METADATA_SCHEMA_VERSION:
        raise ValueError(f"unsupported execution metadata schema_version {version!r}")
    traces = data.get("decision_traces", ())
    if not isinstance(traces, (list, tuple)):
        raise ValueError("execution_metadata.decision_traces must be an array")
    return make_execution_metadata(
        data.get("feature_extraction_report"),
        traces,
        approved_research=data.get(APPROVED_RESEARCH_KEY),
        trusted_extractor=data.get(TRUSTED_EXTRACTOR_KEY),
        experiment_reflection=data.get(EXPERIMENT_REFLECTION_KEY),
    )


def merge_execution_metadata(
    *values: Mapping[str, Any],
    feature_report: FeatureExtractionReport | Mapping[str, Any] | None = None,
    decision_traces: Sequence[DecisionTrace | Mapping[str, Any]] = (),
    approved_research: Mapping[str, Any] | None = None,
    trusted_extractor: Mapping[str, Any] | None = None,
    experiment_reflection: ExperimentReflection | Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Merge producer metadata without duplicating stable trace identifiers."""

    report = feature_report
    traces: list[DecisionTrace | Mapping[str, Any]] = []
    approved = approved_research
    reflection = experiment_reflection
    extractor_binding: Mapping[str, Any] | None = (
        normalize_trusted_extractor_binding(trusted_extractor)
        if trusted_extractor
        else None
    )
    for value in values:
        canonical = normalize_execution_metadata(value) if value else {}
        if canonical.get("feature_extraction_report") is not None:
            report = canonical["feature_extraction_report"]
        traces.extend(canonical.get("decision_traces", ()))
        if canonical.get(APPROVED_RESEARCH_KEY):
            approved = canonical[APPROVED_RESEARCH_KEY]
        if canonical.get(TRUSTED_EXTRACTOR_KEY):
            incoming_binding = canonical[TRUSTED_EXTRACTOR_KEY]
            if extractor_binding is not None and dict(extractor_binding) != dict(incoming_binding):
                raise ValueError("cannot merge conflicting trusted extractor bindings")
            extractor_binding = incoming_binding
        if canonical.get(EXPERIMENT_REFLECTION_KEY):
            reflection = canonical[EXPERIMENT_REFLECTION_KEY]
    traces.extend(decision_traces)

    deduplicated: dict[str, DecisionTrace] = {}
    anonymous: list[DecisionTrace] = []
    for raw in traces:
        trace = (
            raw
            if isinstance(raw, DecisionTrace)
            else DecisionTrace.from_mapping(_mapping(raw, "decision trace"))
        )
        if trace.trace_id:
            deduplicated[trace.trace_id] = trace
        else:
            anonymous.append(trace)
    return make_execution_metadata(
        report,
        [*deduplicated.values(), *anonymous],
        approved_research=approved,
        trusted_extractor=extractor_binding,
        experiment_reflection=reflection,
    )


def merge_approved_research_into_package(
    package_path: str | os.PathLike[str],
    approved_research: Mapping[str, Any],
) -> dict[str, Any]:
    """Atomically add approved research to one locally produced ZIP manifest."""

    approved = normalize_approved_research(approved_research)
    return _merge_optional_metadata_into_package(
        package_path, approved_research=approved
    )


def merge_experiment_reflection_into_package(
    package_path: str | os.PathLike[str],
    experiment_reflection: ExperimentReflection | Mapping[str, Any],
) -> dict[str, Any]:
    """Atomically persist one validated, optional reflection in a package."""

    reflection = (
        experiment_reflection
        if isinstance(experiment_reflection, ExperimentReflection)
        else ExperimentReflection.from_mapping(
            _mapping(experiment_reflection, EXPERIMENT_REFLECTION_KEY)
        )
    )
    return _merge_optional_metadata_into_package(
        package_path, experiment_reflection=reflection
    )


def _merge_optional_metadata_into_package(
    package_path: str | os.PathLike[str],
    *,
    approved_research: Mapping[str, Any] | None = None,
    experiment_reflection: ExperimentReflection | Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Rewrite exactly one trusted manifest while retaining all ZIP members."""

    source = Path(package_path)
    if not source.is_file():
        raise OSError("model package is unavailable")
    descriptor, temp_name = tempfile.mkstemp(
        prefix=f".{source.name}.", suffix=".tmp", dir=str(source.parent)
    )
    os.close(descriptor)
    temp_path = Path(temp_name)
    manifest_count = 0
    merged: dict[str, Any] = {}
    try:
        with zipfile.ZipFile(source, "r") as incoming, zipfile.ZipFile(
            temp_path, "w"
        ) as outgoing:
            for info in incoming.infolist():
                payload = incoming.read(info)
                # The trusted package contract designates exactly one
                # automl_model_manifest.json as the model manifest.  Image
                # artifacts may also contain a dataset-level manifest.json;
                # that file is not the package manifest and must remain
                # untouched.
                if (
                    PurePosixPath(info.filename).name
                    == "automl_model_manifest.json"
                ):
                    manifest_count += 1
                    if manifest_count > 1:
                        raise ValueError("model package contains multiple manifests")
                    try:
                        manifest = json.loads(payload.decode("utf-8"))
                    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
                        raise ValueError("model package manifest is invalid") from exc
                    if not isinstance(manifest, Mapping):
                        raise ValueError("model package manifest must be a mapping")
                    merged = merge_execution_metadata(
                        execution_metadata_from_manifest(manifest),
                        approved_research=approved_research,
                        experiment_reflection=experiment_reflection,
                    )
                    payload = json.dumps(
                        manifest_with_execution_metadata(manifest, merged),
                        indent=2,
                        sort_keys=True,
                    ).encode("utf-8")
                outgoing.writestr(info, payload)
        if manifest_count != 1:
            raise ValueError("model package manifest is missing")
        os.replace(temp_path, source)
    except Exception:
        try:
            temp_path.unlink()
        except OSError:
            pass
        raise
    return merged


def execution_metadata_from_manifest(manifest: Mapping[str, Any]) -> dict[str, Any]:
    """Read optional metadata; absence is the legacy-compatible empty value."""

    value = _mapping(manifest, "manifest").get(EXECUTION_METADATA_KEY)
    if value is None:
        return {}
    return normalize_execution_metadata(_mapping(value, EXECUTION_METADATA_KEY))


def trusted_extractor_from_manifest(manifest: Mapping[str, Any]) -> dict[str, Any]:
    """Read the optional package binding; legacy packages return an empty mapping."""

    metadata = execution_metadata_from_manifest(manifest)
    return dict(metadata.get(TRUSTED_EXTRACTOR_KEY) or {})


def manifest_with_execution_metadata(
    manifest: Mapping[str, Any], metadata: Mapping[str, Any]
) -> dict[str, Any]:
    """Return a manifest copy with validated metadata only when it is present."""

    result = dict(_mapping(manifest, "manifest"))
    canonical = normalize_execution_metadata(metadata) if metadata else {}
    if canonical:
        result[EXECUTION_METADATA_KEY] = canonical
    return result


__all__ = [
    "EXECUTION_METADATA_KEY",
    "EXECUTION_METADATA_SCHEMA_VERSION",
    "APPROVED_RESEARCH_KEY",
    "EXPERIMENT_REFLECTION_KEY",
    "TRUSTED_EXTRACTOR_KEY",
    "execution_metadata_from_manifest",
    "make_decision_trace",
    "make_execution_metadata",
    "make_feature_extraction_report",
    "manifest_with_execution_metadata",
    "merge_approved_research_into_package",
    "merge_experiment_reflection_into_package",
    "merge_execution_metadata",
    "normalize_execution_metadata",
    "normalize_approved_research",
    "normalize_trusted_extractor_binding",
    "trusted_extractor_from_manifest",
]