"""
Sanna middleware — runtime enforcement decorator for AI agent pipelines.

@sanna_observe wraps agent functions, captures I/O, runs coherence checks
driven by constitution invariants, and enforces policy (halt/warn/log) per
check. Every execution produces a reasoning receipt.

v0.6.0: The constitution is the control plane. Invariants drive which checks
run and at what enforcement level.
"""

import functools
import inspect
import json
import logging
import time
import uuid
import warnings
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, List, Optional, Sequence

from .receipt import (
    generate_receipt,
    check_c1_context_contradiction,
    check_c2_unmarked_inference,
    check_c3_false_certainty,
    check_c4_conflict_collapse,
    check_c5_premature_compression,
    CheckResult,
    ConstitutionProvenance,
    Enforcement,
    HaltEvent,
    SannaReceipt,
    TOOL_VERSION,
    SPEC_VERSION,
    CHECKS_VERSION,
    select_final_answer,
    extract_context,
    extract_query,
    FinalAnswerProvenance,
)
from .hashing import hash_text, hash_obj, EMPTY_HASH

logger = logging.getLogger("sanna.middleware")

# Check function registry keyed by check ID
_CHECK_FUNCTIONS = {
    "C1": check_c1_context_contradiction,
    "C2": check_c2_unmarked_inference,
    "C3": check_c3_false_certainty,
    "C4": check_c4_conflict_collapse,
    "C5": check_c5_premature_compression,
}

# Parameter names to auto-detect for context and query
_CONTEXT_PARAM_NAMES = {"context", "retrieved_context", "documents", "retrieved_docs"}
_QUERY_PARAM_NAMES = {"query", "prompt", "input", "user_input", "question"}
_OUTPUT_PARAM_NAMES = {"response", "output", "answer", "result"}


# =============================================================================
# PUBLIC TYPES
# =============================================================================

class SannaHaltError(Exception):
    """Raised when reasoning checks fail and enforcement_level='halt'."""

    def __init__(self, message: str, receipt: dict):
        super().__init__(message)
        self.receipt = receipt


class SannaResult:
    """Wrapper around agent output with receipt attached."""

    def __init__(self, output: Any, receipt: dict):
        self.output = output
        self.receipt = receipt

    @property
    def status(self) -> str:
        return self.receipt.get("status", "UNKNOWN")

    @property
    def passed(self) -> bool:
        return self.status == "PASS"

    def __repr__(self) -> str:
        return f"SannaResult(status={self.status!r}, output={self.output!r})"


# =============================================================================
# INPUT MAPPING
# =============================================================================

def _resolve_inputs(
    func,
    args: tuple,
    kwargs: dict,
    context_param: Optional[str],
    query_param: Optional[str],
) -> dict:
    """
    Resolve context and query from function arguments.

    Precedence:
    1. Explicit mapping via context_param/query_param
    2. Convention-based parameter name matching
    3. Single dict argument: look for keys inside it

    Returns dict with ``context`` (str), ``query`` (str), and
    ``structured_context`` (list or None).
    """
    sig = inspect.signature(func)
    bound = sig.bind(*args, **kwargs)
    bound.apply_defaults()
    all_args = dict(bound.arguments)

    context = ""
    query = ""
    structured_context = None
    raw_context = None  # preserve raw value for structured extraction

    # --- Explicit mapping ---
    if context_param and context_param in all_args:
        raw_context = all_args[context_param]
        context = _to_str(raw_context)
    if query_param and query_param in all_args:
        query = _to_str(all_args[query_param])

    # --- Convention-based ---
    if not context:
        for name in _CONTEXT_PARAM_NAMES:
            if name in all_args:
                raw_context = all_args[name]
                context = _to_str(raw_context)
                break

    if not query:
        for name in _QUERY_PARAM_NAMES:
            if name in all_args:
                query = _to_str(all_args[name])
                break

    # --- Single dict argument fallback ---
    if (not context or not query) and len(all_args) == 1:
        single_val = next(iter(all_args.values()))
        if isinstance(single_val, dict):
            if not context:
                for name in _CONTEXT_PARAM_NAMES:
                    if name in single_val:
                        raw_context = single_val[name]
                        context = _to_str(raw_context)
                        break
            if not query:
                for name in _QUERY_PARAM_NAMES:
                    if name in single_val:
                        query = _to_str(single_val[name])
                        break

    # Extract structured context if raw value is a list of source-annotated dicts
    if raw_context is not None:
        structured_context = _extract_structured_context(raw_context)

    return {"context": context, "query": query, "structured_context": structured_context}


def _to_str(val: Any) -> str:
    """Coerce a value to string for check inputs.

    When ``val`` is a list of dicts with ``"text"`` keys (structured
    context), extracts the text portions for a clean string representation.
    """
    if val is None:
        return ""
    if isinstance(val, str):
        return val
    if isinstance(val, list):
        parts = []
        for item in val:
            if isinstance(item, dict) and "text" in item:
                parts.append(item["text"])
            else:
                parts.append(str(item))
        return "\n".join(parts)
    return str(val)


def _extract_structured_context(val: Any) -> Optional[list]:
    """Extract structured context if val is a list of source-annotated dicts.

    Returns the list if every item is a dict with at least a ``"text"`` key.
    Returns None otherwise (including for plain strings).
    """
    if not isinstance(val, list) or not val:
        return None
    if all(isinstance(item, dict) and "text" in item for item in val):
        return val
    return None


# =============================================================================
# TRACE CONSTRUCTION
# =============================================================================

def build_trace_data(
    *,
    correlation_id: str = "",
    query: str,
    context: str,
    output: str,
) -> dict:
    """
    Build a trace_data dict in the shape generate_receipt() expects.

    Constructs a minimal trace with a retrieval span (for context/query)
    and a trace-level output (for the final answer).
    """
    # FIX-33: correlation_id must not contain pipe (used as fingerprint delimiter)
    if correlation_id and "|" in correlation_id:
        raise ValueError("correlation_id must not contain the pipe character '|'")
    cid = correlation_id
    return {
        "correlation_id": cid,
        "name": "sanna_observe",
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "input": {"query": query},
        "output": {"final_answer": output},
        "metadata": {},
        "observations": [
            {
                "id": f"obs-retrieval-{cid}",
                "name": "retrieval",
                "type": "SPAN",
                "input": {"query": query},
                "output": {"context": context},
                "metadata": {},
                "start_time": None,
                "end_time": None,
            }
        ],
    }


# Backward-compatible alias
_build_trace_data = build_trace_data


def _namespace_extensions(extensions: Optional[dict]) -> dict:
    """Convert flat extension keys to reverse-domain notation (v0.13.0).

    ``"middleware"`` → ``"com.sanna.middleware"``
    ``"gateway"`` → ``"com.sanna.gateway"``

    Keys already using reverse-domain notation are passed through.
    """
    if not extensions:
        return {}
    result = {}
    for key, value in extensions.items():
        if key == "middleware":
            result["com.sanna.middleware"] = value
        elif key == "gateway":
            gw = result.get("com.sanna.gateway", {})
            gw.update(value)
            result["com.sanna.gateway"] = gw
        elif "." in key:
            # Already namespaced — merge if key already exists (e.g.
            # from a prior "gateway" flat key that was promoted).
            existing = result.get(key)
            if existing and isinstance(existing, dict) and isinstance(value, dict):
                existing.update(value)
            else:
                result[key] = value
        else:
            # Unknown flat key — namespace under com.sanna
            result[f"com.sanna.{key}"] = value
    return result


# =============================================================================
# SOURCE TRUST HELPERS
# =============================================================================

_VALID_TIERS = frozenset({"tier_1", "tier_2", "tier_3", "untrusted", "unclassified"})


def _normalize_tier(tier: str) -> str:
    """Normalize a tier string to a canonical value.

    Handles case insensitivity, underscores/hyphens/spaces, and returns
    ``"unclassified"`` for unrecognized values.
    """
    normalized = tier.strip().lower().replace("-", "_").replace(" ", "_")
    if normalized in _VALID_TIERS:
        return normalized
    return "unclassified"


def _resolve_source_tiers(
    structured_context: list,
    trusted_sources,
) -> list:
    """Resolve trust tiers for each source in structured context.

    Looks up each source name in the constitution's ``trusted_sources``
    mapping. Sources not found in any tier default to ``"unclassified"``.

    Returns the structured context list with ``tier`` resolved on each item.
    """
    if trusted_sources is None:
        # No trusted_sources in constitution → use explicit tier or unclassified
        return [
            {**item, "tier": _normalize_tier(item["tier"]) if item.get("tier") else "unclassified"}
            for item in structured_context
        ]

    # Build reverse lookup: source_name → tier
    tier_map: dict[str, str] = {}
    for tier_name in ("tier_1", "tier_2", "tier_3", "untrusted"):
        for source in getattr(trusted_sources, tier_name, []):
            tier_map[source] = tier_name

    resolved = []
    for item in structured_context:
        source_name = item.get("source", "unknown")
        # Explicit tier in the context item takes precedence
        explicit_tier = item.get("tier")
        if explicit_tier:
            tier = _normalize_tier(explicit_tier)
        else:
            tier = tier_map.get(source_name, "unclassified")
        resolved.append({**item, "tier": tier})
    return resolved


def _build_source_trust_evaluations(
    structured_context: list,
) -> list:
    """Build source_trust_evaluations records for the receipt.

    Deduplicates by source name. Each record documents what tier
    a source was classified as during this trace.
    """
    evaluations = []
    seen: set[str] = set()
    timestamp = datetime.now(timezone.utc).isoformat()

    for item in structured_context:
        source_name = item.get("source", "unknown")
        if source_name in seen:
            continue
        seen.add(source_name)

        tier = item.get("tier", "unclassified")
        evaluations.append({
            "source_name": source_name,
            "trust_tier": tier,
            "evaluated_at": timestamp,
            "verification_flag": tier == "tier_2",
            "context_used": tier in ("tier_1", "tier_2", "tier_3"),
        })
    return evaluations


# =============================================================================
# CONSTITUTION-DRIVEN RECEIPT GENERATION
# =============================================================================

def _generate_constitution_receipt(
    trace_data: dict,
    check_configs: list,
    custom_records: list,
    constitution_ref: Optional[dict],
    constitution_version: str,
    extensions: Optional[dict] = None,
    enforcement: Optional[HaltEvent] = None,
    authority_decisions: Optional[list] = None,
    escalation_events: Optional[list] = None,
    source_trust_evaluations: Optional[list] = None,
    structured_context: Optional[list] = None,
    error_policy: str = "fail_closed",
) -> dict:
    """Generate a receipt using constitution-driven check configs.

    Runs only the checks specified by check_configs, at their enforcement
    levels, and includes custom invariant records as NOT_CHECKED.
    """
    final_answer, answer_provenance = select_final_answer(trace_data)
    context = extract_context(trace_data)
    query_text = extract_query(trace_data)

    # Run configured checks
    check_results = []
    for cfg in check_configs:
        try:
            # C1 gets structured context for source-aware evaluation
            if structured_context and cfg.check_id == "sanna.context_contradiction":
                result = cfg.check_fn(
                    context, final_answer,
                    enforcement=cfg.enforcement_level,
                    structured_context=structured_context,
                )
            else:
                result = cfg.check_fn(context, final_answer, enforcement=cfg.enforcement_level)
        except Exception as exc:
            if cfg.source == "custom_evaluator":
                if error_policy == "fail_closed":
                    # Fail-closed: treat error as a real failure
                    check_results.append({
                        "check_id": cfg.check_id,
                        "name": "Custom Invariant",
                        "passed": False,
                        "severity": cfg.enforcement_level if cfg.enforcement_level in ("critical", "high", "medium", "low") else "critical",
                        "evidence": None,
                        "details": f"Evaluator error (fail_closed): {exc}",
                        "triggered_by": cfg.triggered_by,
                        "enforcement_level": cfg.enforcement_level,
                        "constitution_version": constitution_version,
                        "check_impl": cfg.check_impl or None,
                        "replayable": False,
                        "status": "FAILED",
                    })
                else:
                    # Fail-open: exclude from counts via ERRORED
                    check_results.append({
                        "check_id": cfg.check_id,
                        "name": "Custom Invariant",
                        "passed": True,
                        "severity": "info",
                        "evidence": None,
                        "details": f"Evaluator error (fail_open): {exc}",
                        "triggered_by": cfg.triggered_by,
                        "enforcement_level": cfg.enforcement_level,
                        "constitution_version": constitution_version,
                        "check_impl": cfg.check_impl or None,
                        "replayable": False,
                        "status": "ERRORED",
                    })
                continue
            raise

        entry = {
            "check_id": cfg.check_id,
            "name": result.name,
            "passed": result.passed,
            "severity": result.severity,
            "evidence": result.evidence,
            "details": result.details,
            "triggered_by": cfg.triggered_by,
            "enforcement_level": cfg.enforcement_level,
            "constitution_version": constitution_version,
            "check_impl": cfg.check_impl or None,
            "replayable": cfg.source != "custom_evaluator",
        }
        check_results.append(entry)

    # Add custom invariants as NOT_CHECKED
    for custom in custom_records:
        check_results.append({
            "check_id": custom.invariant_id,
            "name": "Custom Invariant",
            "passed": True,  # NOT_CHECKED counts as not-failed
            "severity": "info",
            "evidence": None,
            "details": custom.rule,
            "triggered_by": custom.invariant_id,
            "enforcement_level": custom.enforcement,
            "constitution_version": constitution_version,
            "status": custom.status,
            "reason": custom.reason,
            "check_impl": None,
            "replayable": False,
        })

    # Count pass/fail (NOT_CHECKED and ERRORED don't count as failures)
    _NON_EVALUATED = ("NOT_CHECKED", "ERRORED")
    standard_checks = [c for c in check_results if c.get("status") not in _NON_EVALUATED]
    not_evaluated = [c for c in check_results if c.get("status") in _NON_EVALUATED]
    passed = sum(1 for c in standard_checks if c["passed"])
    failed = len(standard_checks) - passed

    # Determine status from standard checks only
    # Severity hierarchy: critical/high → FAIL, warning/medium/low → WARN
    _FAIL_SEVERITIES = {"critical", "high"}
    _WARN_SEVERITIES = {"warning", "medium", "low"}
    critical_fails = sum(1 for c in standard_checks if not c["passed"] and c["severity"] in _FAIL_SEVERITIES)
    warning_fails = sum(1 for c in standard_checks if not c["passed"] and c["severity"] in _WARN_SEVERITIES)

    if critical_fails > 0:
        status = "FAIL"
    elif warning_fails > 0:
        status = "WARN"
    elif not_evaluated:
        status = "PARTIAL"
    else:
        status = "PASS"

    # Build evaluation_coverage
    total_invariants = len(check_results)
    evaluated_count = len(standard_checks)
    not_checked_count = len(not_evaluated)
    coverage_basis_points = (evaluated_count * 10000) // total_invariants if total_invariants > 0 else 10000
    evaluation_coverage = {
        "total_invariants": total_invariants,
        "evaluated": evaluated_count,
        "not_checked": not_checked_count,
        "coverage_basis_points": coverage_basis_points,
    }

    inputs = {"query": query_text if query_text else None, "context": context if context else None}
    outputs = {"response": final_answer if final_answer else None}

    context_hash = hash_obj(inputs)
    output_hash = hash_obj(outputs)

    # Strip mutable constitution_approval before hashing — approval can change
    # without invalidating the fingerprint (verified separately).
    if constitution_ref:
        _cref = {k: v for k, v in constitution_ref.items() if k != "constitution_approval"}
        constitution_hash_val = hash_obj(_cref)
    else:
        constitution_hash_val = EMPTY_HASH

    # Build enforcement object
    enforcement_dict = None
    if enforcement is not None:
        enforcement_obj_dict = asdict(enforcement)
        enforcement_dict = {
            "action": "halted" if enforcement.halted else "allowed",
            "reason": enforcement.reason,
            "failed_checks": enforcement.failed_checks,
            "enforcement_mode": enforcement.enforcement_mode,
            "timestamp": enforcement.timestamp,
        }
    else:
        enforcement_obj_dict = None
    enforcement_hash_val = hash_obj(enforcement_dict) if enforcement_dict else EMPTY_HASH

    # Build fingerprint from check results (include all enforcement and impl fields)
    checks_fingerprint_data = [
        {
            "check_id": c["check_id"],
            "passed": c["passed"],
            "severity": c["severity"],
            "evidence": c["evidence"],
            "triggered_by": c.get("triggered_by"),
            "enforcement_level": c.get("enforcement_level"),
            "check_impl": c.get("check_impl"),
            "replayable": c.get("replayable"),
        }
        for c in check_results
    ]
    checks_hash = hash_obj(checks_fingerprint_data)
    coverage_hash = hash_obj(evaluation_coverage)

    # Unified fingerprint formula (v0.13.0) — always 12 pipe-delimited fields
    correlation_id = trace_data.get("correlation_id", "")
    authority_hash = hash_obj(authority_decisions) if authority_decisions else EMPTY_HASH
    escalation_hash = hash_obj(escalation_events) if escalation_events else EMPTY_HASH
    trust_hash = hash_obj(source_trust_evaluations) if source_trust_evaluations else EMPTY_HASH

    # Namespace extensions for v0.13.0
    namespaced_ext = _namespace_extensions(extensions) if extensions else {}
    extensions_hash = hash_obj(namespaced_ext) if namespaced_ext else EMPTY_HASH

    fingerprint_input = f"{correlation_id}|{context_hash}|{output_hash}|{CHECKS_VERSION}|{checks_hash}|{constitution_hash_val}|{enforcement_hash_val}|{coverage_hash}|{authority_hash}|{escalation_hash}|{trust_hash}|{extensions_hash}"

    full_fp = hash_text(fingerprint_input)
    receipt_fingerprint = hash_text(fingerprint_input, truncate=16)

    receipt_dict = {
        "spec_version": SPEC_VERSION,
        "tool_version": TOOL_VERSION,
        "checks_version": CHECKS_VERSION,
        "receipt_id": str(uuid.uuid4()),
        "receipt_fingerprint": receipt_fingerprint,
        "full_fingerprint": full_fp,
        "correlation_id": correlation_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "inputs": inputs,
        "outputs": outputs,
        "context_hash": context_hash,
        "output_hash": output_hash,
        "checks": check_results,
        "checks_passed": passed,
        "checks_failed": failed,
        "status": status,
        "evaluation_coverage": evaluation_coverage,
        "constitution_ref": constitution_ref,
    }

    # Enforcement object
    if enforcement_dict:
        receipt_dict["enforcement"] = enforcement_dict

    # Authority enforcement sections (only included when present)
    if authority_decisions:
        receipt_dict["authority_decisions"] = authority_decisions
    if escalation_events:
        receipt_dict["escalation_events"] = escalation_events
    if source_trust_evaluations:
        receipt_dict["source_trust_evaluations"] = source_trust_evaluations

    receipt_dict["extensions"] = namespaced_ext

    return receipt_dict


def _generate_no_invariants_receipt(
    trace_data: dict,
    constitution_ref: Optional[dict],
    extensions: Optional[dict] = None,
) -> dict:
    """Generate a receipt for a constitution with no invariants.

    No checks run. The receipt documents that no invariants were defined.
    Uses the unified fingerprint formula (v0.13.0).
    """
    final_answer, _answer_provenance = select_final_answer(trace_data)
    context = extract_context(trace_data)
    query_text = extract_query(trace_data)

    inputs = {"query": query_text if query_text else None, "context": context if context else None}
    outputs = {"response": final_answer if final_answer else None}

    context_hash = hash_obj(inputs)
    output_hash = hash_obj(outputs)

    # Strip mutable constitution_approval before hashing (parity with verify.py).
    if constitution_ref:
        _cref = {k: v for k, v in constitution_ref.items() if k != "constitution_approval"}
        constitution_hash_val = hash_obj(_cref)
    else:
        constitution_hash_val = EMPTY_HASH

    # Unified fingerprint formula — always 12 pipe-delimited fields
    correlation_id = trace_data.get("correlation_id", "")
    checks_hash = hash_obj([])

    # Namespace extensions for v0.13.0
    namespaced_ext = _namespace_extensions(extensions) if extensions else {}
    extensions_hash = hash_obj(namespaced_ext) if namespaced_ext else EMPTY_HASH

    fingerprint_input = f"{correlation_id}|{context_hash}|{output_hash}|{CHECKS_VERSION}|{checks_hash}|{constitution_hash_val}|{EMPTY_HASH}|{EMPTY_HASH}|{EMPTY_HASH}|{EMPTY_HASH}|{EMPTY_HASH}|{extensions_hash}"

    full_fp = hash_text(fingerprint_input)
    receipt_fingerprint = hash_text(fingerprint_input, truncate=16)

    receipt_dict = {
        "spec_version": SPEC_VERSION,
        "tool_version": TOOL_VERSION,
        "checks_version": CHECKS_VERSION,
        "receipt_id": str(uuid.uuid4()),
        "receipt_fingerprint": receipt_fingerprint,
        "full_fingerprint": full_fp,
        "correlation_id": correlation_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "inputs": inputs,
        "outputs": outputs,
        "context_hash": context_hash,
        "output_hash": output_hash,
        "checks": [],
        "checks_passed": 0,
        "checks_failed": 0,
        "status": "PASS",
        "constitution_ref": constitution_ref,
    }

    receipt_dict["extensions"] = namespaced_ext

    return receipt_dict


# =============================================================================
# FILE OUTPUT
# =============================================================================

def _write_receipt(receipt: dict, receipt_dir: str) -> Path:
    """Write receipt JSON to the configured directory (atomic)."""
    from .utils.safe_io import atomic_write_text_sync, ensure_secure_dir

    dir_path = Path(receipt_dir)
    ensure_secure_dir(dir_path)

    cid = receipt.get('correlation_id', 'unknown')
    filename = f"{receipt['receipt_id']}_{cid}.json"
    filepath = dir_path / filename

    atomic_write_text_sync(filepath, json.dumps(receipt, indent=2))

    logger.debug("Receipt written to %s", filepath)
    return filepath


# =============================================================================
# PRE-EXECUTION REASONING GATE
# =============================================================================

# Shared executor for sync-path reasoning gate (avoids creating per-call pools)
import concurrent.futures as _cf
_SYNC_EXECUTOR = _cf.ThreadPoolExecutor(max_workers=4)


def _run_reasoning_gate(
    constitution,
    justification: str,
    kwargs: dict,
    func_name: str = "",
) -> dict | None:
    """Run reasoning pipeline synchronously before function execution.

    Returns a dict with ``passed`` (bool) and ``failure_reason`` (str)
    if reasoning checks ran, or *None* if reasoning is not configured
    or the pipeline is unavailable.
    """
    import asyncio

    try:
        from .reasoning.pipeline import ReasoningPipeline
    except ImportError:
        return None

    try:
        pipeline = ReasoningPipeline(constitution)
        if not pipeline.enabled:
            return None

        # Run the async pipeline in a sync context.
        # Detect whether we're already inside an event loop to avoid
        # "cannot call asyncio.run() while another loop is running".
        coro = pipeline.evaluate(
            tool_name=func_name,
            args=kwargs,
            enforcement_level="must_escalate",
        )
        try:
            asyncio.get_running_loop()
            # Already inside an event loop — offload to shared thread pool.
            evaluation = _SYNC_EXECUTOR.submit(asyncio.run, coro).result(timeout=30)
        except RuntimeError:
            # No running loop — safe to use asyncio.run()
            evaluation = asyncio.run(coro)

        return {
            "passed": evaluation.passed,
            "failure_reason": evaluation.failure_reason,
            "overall_score": evaluation.overall_score,
            "assurance": evaluation.assurance,
        }
    except Exception as e:
        logger.warning("Pre-execution reasoning gate failed: %s", e)
        return None


async def _run_reasoning_gate_async(
    constitution,
    justification: str,
    kwargs: dict,
    func_name: str = "",
) -> dict | None:
    """Run reasoning pipeline asynchronously — no thread pool, no blocking.

    Directly awaits the pipeline coroutine. Use this from async code paths
    to avoid blocking the event loop.
    """
    try:
        from .reasoning.pipeline import ReasoningPipeline
    except ImportError:
        return None

    try:
        pipeline = ReasoningPipeline(constitution)
        if not pipeline.enabled:
            return None

        evaluation = await pipeline.evaluate(
            tool_name=func_name,
            args=kwargs,
            enforcement_level="must_escalate",
        )

        return {
            "passed": evaluation.passed,
            "failure_reason": evaluation.failure_reason,
            "overall_score": evaluation.overall_score,
            "assurance": evaluation.assurance,
        }
    except Exception as e:
        logger.warning("Pre-execution reasoning gate failed: %s", e)
        return None


# =============================================================================
# THE DECORATOR
# =============================================================================

def sanna_observe(
    _func=None,
    *,
    receipt_dir: Optional[str] = None,
    store=None,
    context_param: Optional[str] = None,
    query_param: Optional[str] = None,
    constitution_path: Optional[str] = None,
    constitution_public_key_path: Optional[str] = None,
    private_key_path: Optional[str] = None,
    identity_provider_keys: Optional[dict[str, str]] = None,
    require_constitution_sig: bool = True,
    error_policy: str = "fail_closed",
    strict: bool = True,
    # Legacy parameters — ignored when constitution has invariants
    on_violation: str = "halt",
    checks: Optional[List[str]] = None,
    halt_on: Optional[List[str]] = None,
    constitution: Optional[ConstitutionProvenance] = None,
):
    """
    Decorator that wraps an agent function with Sanna coherence checks.

    v0.6.0: The constitution is the control plane. When a constitution with
    invariants is provided, the invariants drive which checks run and at what
    enforcement level. Each check enforces independently based on its invariant.

    Args:
        constitution_path: Path to a Sanna constitution YAML/JSON file.
            This is the primary way to configure checks. The constitution's
            invariants drive which checks run and how they enforce.
        constitution_public_key_path: Path to the Ed25519 public key for
            cryptographic verification of the constitution signature.
        receipt_dir: Directory to write receipt JSON. None to skip.
        store: Optional ReceiptStore instance or db_path string. When
            provided, receipts are auto-saved after generation. Store
            failures are logged but never break receipt generation.
        context_param: Explicit name of the context parameter.
        query_param: Explicit name of the query parameter.
        identity_provider_keys: Optional mapping of public_key_id to path to
            public key file for identity claim verification. When not provided:
            claims with signatures get status "no_key" (signature present but
            no key to verify against); claims without signatures get status
            "unverified".
        require_constitution_sig: If True (default), require cryptographic
            verification of the constitution signature. When True, a public
            key must be provided and the signature must verify. Set to False
            for local development only.
        strict: If True (default), validate constitution against JSON schema
            on load. Catches typos like ``invariant:`` instead of ``invariants:``.
        on_violation: Legacy. Used when no constitution invariants are present.
        checks: Legacy. Used when no constitution invariants are present.
        halt_on: Legacy. Used when no constitution invariants are present.
        constitution: Legacy ConstitutionProvenance for governance tracking.

    Returns:
        SannaResult wrapping the function output and receipt, or raises
        SannaHaltError if policy dictates halt.
    """
    if on_violation not in ("halt", "warn", "log"):
        raise ValueError(f"on_violation must be 'halt', 'warn', or 'log', got {on_violation!r}")
    if error_policy not in ("fail_closed", "fail_open"):
        raise ValueError(f"error_policy must be 'fail_closed' or 'fail_open', got {error_policy!r}")

    # Resolve store at decoration time
    _store_instance = None
    if store is not None:
        if isinstance(store, str):
            from .store import ReceiptStore as _ReceiptStore
            _store_instance = _ReceiptStore(store)
        else:
            _store_instance = store

    # Decoration-time constitution loading
    loaded_constitution = None
    constitution_ref_override = None
    check_configs = None
    custom_records = None

    # Track signature verification status for receipt
    _sig_verified = None  # None until we know

    if constitution_path is not None:
        import os as _os
        from .constitution import load_constitution as _load_constitution, constitution_to_receipt_ref, SannaConstitutionError
        from .enforcement import configure_checks
        from .utils.crypto_validation import is_valid_signature_structure

        loaded_constitution = _load_constitution(constitution_path, validate=strict)

        # Resolve public key path: explicit parameter > env var (with key_id check)
        _effective_pub_key = constitution_public_key_path
        if not _effective_pub_key:
            _env_key = _os.environ.get("SANNA_CONSTITUTION_PUBLIC_KEY")
            if _env_key and _os.path.isfile(_env_key):
                # Only use env var key if its key_id matches the constitution's signature key_id
                try:
                    from .crypto import load_public_key, compute_key_id
                    _env_pub = load_public_key(_env_key)
                    _env_key_id = compute_key_id(_env_pub)
                    _const_sig = loaded_constitution.provenance.signature if loaded_constitution.provenance else None
                    if _const_sig and getattr(_const_sig, 'key_id', None) == _env_key_id:
                        _effective_pub_key = _env_key
                except Exception:
                    pass  # Env var key unreadable — ignore
        if not loaded_constitution.policy_hash:
            raise SannaConstitutionError(
                f"Constitution is not signed (no policy hash): {constitution_path}. "
                f"Run: sanna-sign-constitution {constitution_path}"
            )
        # Reject constitutions that are hashed but not Ed25519-signed (always enforced)
        _sig = loaded_constitution.provenance.signature if loaded_constitution.provenance else None
        _has_structural_sig = is_valid_signature_structure(_sig)
        if not _has_structural_sig:
            raise SannaConstitutionError(
                f"Constitution signature is missing or malformed: {constitution_path}. "
                f"Sign with: sanna-sign-constitution {constitution_path} --private-key <key>"
            )

        if require_constitution_sig:
            # Strict mode: require public key and cryptographic verification
            if not _effective_pub_key:
                raise SannaConstitutionError(
                    f"Constitution has signature but no public key configured to "
                    f"verify it. Provide constitution_public_key_path, set "
                    f"SANNA_CONSTITUTION_PUBLIC_KEY env var, or set "
                    f"require_constitution_sig=False for local development."
                )
            # Cryptographic verification
            from .crypto import verify_constitution_full
            if not verify_constitution_full(loaded_constitution, _effective_pub_key):
                raise SannaConstitutionError(
                    f"Constitution signature verification failed. "
                    f"The constitution may have been tampered with. "
                    f"Public key: {_effective_pub_key}"
                )
            _sig_verified = True
        else:
            # Permissive mode: verify if we can, warn otherwise
            if _effective_pub_key:
                from .crypto import verify_constitution_full
                if not verify_constitution_full(loaded_constitution, _effective_pub_key):
                    raise SannaConstitutionError(
                        f"Constitution signature verification failed. "
                        f"The constitution may have been tampered with. "
                        f"Public key: {_effective_pub_key}"
                    )
                _sig_verified = True
            else:
                logger.warning(
                    "Constitution signature present but no public key configured "
                    "to verify it: %s", constitution_path,
                )
                _sig_verified = False

        constitution_ref_override = constitution_to_receipt_ref(loaded_constitution)
        # Add signature_verified to constitution_ref
        if _sig_verified is not None:
            constitution_ref_override["signature_verified"] = _sig_verified

        # Configure checks from invariants
        check_configs, custom_records = configure_checks(loaded_constitution)

        logger.info(
            "Constitution loaded from %s (hash=%s, invariants=%d, checks=%d)",
            constitution_path,
            loaded_constitution.policy_hash[:16],
            len(loaded_constitution.invariants),
            len(check_configs),
        )

    def decorator(func):
        def _pre_execution_check(kwargs):
            """Run pre-execution reasoning gate (sync). Returns (reasoning_result, should_halt, halt_msg)."""
            reasoning_pre_result = None
            if loaded_constitution and loaded_constitution.reasoning:
                justification = kwargs.get("_justification", "")
                if justification and isinstance(justification, str):
                    reasoning_pre_result = _run_reasoning_gate(
                        loaded_constitution, justification, kwargs,
                        func_name=func.__name__,
                    )
                    if reasoning_pre_result and not reasoning_pre_result.get("passed", True):
                        enforcement = loaded_constitution.reasoning.on_missing_justification
                        auto_deny = loaded_constitution.reasoning.auto_deny_on_reasoning_failure
                        if enforcement == "block" or auto_deny:
                            return reasoning_pre_result, True, (
                                f"Reasoning checks failed before execution: "
                                f"{reasoning_pre_result.get('failure_reason', 'unknown')}"
                            )
            return reasoning_pre_result, False, ""

        async def _pre_execution_check_async(kwargs):
            """Run pre-execution reasoning gate (async). Directly awaits the pipeline."""
            reasoning_pre_result = None
            if loaded_constitution and loaded_constitution.reasoning:
                justification = kwargs.get("_justification", "")
                if justification and isinstance(justification, str):
                    reasoning_pre_result = await _run_reasoning_gate_async(
                        loaded_constitution, justification, kwargs,
                        func_name=func.__name__,
                    )
                    if reasoning_pre_result and not reasoning_pre_result.get("passed", True):
                        enforcement = loaded_constitution.reasoning.on_missing_justification
                        auto_deny = loaded_constitution.reasoning.auto_deny_on_reasoning_failure
                        if enforcement == "block" or auto_deny:
                            return reasoning_pre_result, True, (
                                f"Reasoning checks failed before execution: "
                                f"{reasoning_pre_result.get('failure_reason', 'unknown')}"
                            )
            return reasoning_pre_result, False, ""

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            correlation_id = f"sanna-{uuid.uuid4().hex[:12]}"

            # 1. Capture inputs
            resolved = _resolve_inputs(func, args, kwargs, context_param, query_param)

            # 2. Pre-execution reasoning gate
            reasoning_pre_result, should_halt, halt_msg = _pre_execution_check(kwargs)
            if should_halt:
                raise SannaHaltError(halt_msg, receipt=reasoning_pre_result)

            # 3. Execute the wrapped function
            start_ms = time.monotonic_ns() // 1_000_000
            result = func(*args, **kwargs)
            end_ms = time.monotonic_ns() // 1_000_000
            execution_time_ms = end_ms - start_ms

            # Post-execution governance (shared with async path)
            return _post_execution_governance(
                result, resolved, correlation_id, execution_time_ms, func,
            )

        if inspect.iscoroutinefunction(func):
            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                correlation_id = f"sanna-{uuid.uuid4().hex[:12]}"

                # 1. Capture inputs
                resolved = _resolve_inputs(func, args, kwargs, context_param, query_param)

                # 2. Pre-execution reasoning gate (async — no event loop blocking)
                reasoning_pre_result, should_halt, halt_msg = await _pre_execution_check_async(kwargs)
                if should_halt:
                    raise SannaHaltError(halt_msg, receipt=reasoning_pre_result)

                # 3. Execute the wrapped async function
                start_ms = time.monotonic_ns() // 1_000_000
                result = await func(*args, **kwargs)
                end_ms = time.monotonic_ns() // 1_000_000
                execution_time_ms = end_ms - start_ms

                # Post-execution governance (same as sync)
                return _post_execution_governance(
                    result, resolved, correlation_id, execution_time_ms, func,
                )

            return async_wrapper

        return wrapper

    def _post_execution_governance(result, resolved, correlation_id, execution_time_ms, func):
        """Post-execution governance logic shared by sync and async wrappers."""
        # 4. Capture output
        output_str = _to_str(result)

        # 4. Build trace
        trace_data = build_trace_data(
            correlation_id=correlation_id,
            query=resolved["query"],
            context=resolved["context"],
            output=output_str,
        )

        enforcement_decision = "PASSED"

        extensions = {
            "middleware": {
                "decorator": "@sanna_observe",
                "on_violation": on_violation,
                "enforcement_decision": enforcement_decision,
                "function_name": func.__name__,
                "execution_time_ms": execution_time_ms,
            }
        }

        # 4b. Resolve structured context and source tiers
        raw_structured = resolved.get("structured_context")
        resolved_structured = None
        source_trust_evals = None
        if raw_structured and loaded_constitution:
            resolved_structured = _resolve_source_tiers(
                raw_structured,
                loaded_constitution.trusted_sources,
            )
            source_trust_evals = _build_source_trust_evaluations(resolved_structured)

        # 5. Generate receipt — constitution-driven or legacy
        if check_configs is not None:
            constitution_version = loaded_constitution.schema_version if loaded_constitution else ""

            if not check_configs and not custom_records:
                receipt = _generate_no_invariants_receipt(
                    trace_data,
                    constitution_ref=constitution_ref_override,
                    extensions=extensions,
                )
            else:
                receipt = _generate_constitution_receipt(
                    trace_data,
                    check_configs=check_configs,
                    custom_records=custom_records,
                    constitution_ref=constitution_ref_override,
                    constitution_version=constitution_version,
                    extensions=extensions,
                    source_trust_evaluations=source_trust_evals,
                    structured_context=resolved_structured,
                    error_policy=error_policy,
                )

                halt_checks = []
                warn_checks = []
                log_checks = []

                for check in receipt.get("checks", []):
                    if check.get("status") == "NOT_CHECKED":
                        continue
                    if not check.get("passed"):
                        level = check.get("enforcement_level", "log")
                        if level == "halt":
                            halt_checks.append(check)
                        elif level == "warn":
                            warn_checks.append(check)
                        else:
                            log_checks.append(check)

                if halt_checks:
                    enforcement_decision = "HALTED"
                elif warn_checks:
                    enforcement_decision = "WARNED"
                elif log_checks:
                    enforcement_decision = "LOGGED"
                else:
                    enforcement_decision = "PASSED"

                if enforcement_decision != "PASSED":
                    extensions["middleware"]["enforcement_decision"] = enforcement_decision

                    enforcement_obj = None
                    if enforcement_decision == "HALTED":
                        failed_ids = [c["check_id"] for c in halt_checks]
                        enforcement_obj = HaltEvent(
                            halted=True,
                            reason=f"Coherence check failed: {', '.join(failed_ids)}",
                            failed_checks=failed_ids,
                            timestamp=datetime.now(timezone.utc).isoformat(),
                            enforcement_mode="halt",
                        )

                    receipt = _generate_constitution_receipt(
                        trace_data,
                        check_configs=check_configs,
                        custom_records=custom_records,
                        constitution_ref=constitution_ref_override,
                        constitution_version=constitution_version,
                        extensions=extensions,
                        enforcement=enforcement_obj,
                        source_trust_evaluations=source_trust_evals,
                        structured_context=resolved_structured,
                        error_policy=error_policy,
                    )

        else:
            receipt = _generate_no_invariants_receipt(
                trace_data,
                constitution_ref=asdict(constitution) if constitution else None,
                extensions=extensions,
            )
            enforcement_decision = "PASSED"

        # 5a. Identity verification (NOT in fingerprint — added post-hash)
        if loaded_constitution and loaded_constitution.identity.identity_claims:
            from .constitution import verify_identity_claims as _verify_claims
            iv_summary = _verify_claims(
                loaded_constitution.identity,
                provider_keys=identity_provider_keys,
            )
            receipt["identity_verification"] = {
                "total_claims": iv_summary.total_claims,
                "verified": iv_summary.verified_count,
                "failed": iv_summary.failed_count,
                "unverified": iv_summary.unverified_count,
                "all_verified": iv_summary.all_verified,
                "claims": [
                    {
                        "provider": r.claim.provider,
                        "claim_type": r.claim.claim_type,
                        "credential_id": r.claim.credential_id,
                        "status": r.status,
                    }
                    for r in iv_summary.results
                ],
            }

        # 5b. Sign receipt if private key provided
        if private_key_path is not None:
            from .crypto import sign_receipt as _sign_receipt
            receipt = _sign_receipt(receipt, private_key_path)

        # 6. Write receipt to disk if configured
        if receipt_dir is not None:
            _write_receipt(receipt, receipt_dir)

        # 6b. Save to store if configured
        if _store_instance is not None:
            try:
                _store_instance.save(receipt)
            except Exception:
                logger.warning("Failed to save receipt to store", exc_info=True)

        # 7. Apply enforcement
        if enforcement_decision == "HALTED":
            failed = [c for c in receipt["checks"] if not c.get("passed") and c.get("enforcement_level") == "halt"]
            names = ", ".join(f"{c['check_id']} ({c['name']})" for c in failed)
            raise SannaHaltError(
                f"Sanna coherence check failed: {names}",
                receipt=receipt,
            )

        if enforcement_decision == "WARNED":
            warned = [c for c in receipt["checks"] if not c.get("passed") and c.get("enforcement_level") == "warn"]
            names = ", ".join(f"{c['check_id']} ({c['name']})" for c in warned)
            warnings.warn(
                f"Sanna coherence warning: {names} — status={receipt.get('status', 'UNKNOWN')}",
                stacklevel=2,
            )

        return SannaResult(output=result, receipt=receipt)

    # Support both @sanna_observe and @sanna_observe(...)
    if _func is not None:
        return decorator(_func)
    return decorator


# =============================================================================
# PUBLIC ALIASES FOR INTERNAL HELPERS
# =============================================================================
# These functions were originally private (_-prefixed) but are used by
# sanna.gateway and may be useful for third-party integrations that need
# to generate receipts outside the @sanna_observe decorator.

build_trace_data = _build_trace_data
generate_constitution_receipt = _generate_constitution_receipt
