#!/usr/bin/env python3
"""
inference_receipt_v1.py - Explainable Receipt Schema with Cryptographic Commits

Two layers:
1. Trail (the WHY): ISA trace, routing trace, decision evidence, guardrails
2. Seal (tamper-evident): SHA hashes, Merkle proofs, chain commits

Design principle: "SHA = the wax seal, not the diary."
"""

from __future__ import annotations
import json
import hashlib
import time
import os
import platform
import subprocess
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Any, Optional, Tuple

# Schema version - bump on breaking changes
SCHEMA_VERSION = "inference_receipt_v1.0.0"


# =============================================================================
# CANONICALIZATION
# =============================================================================

def canonicalize(obj: Any) -> str:
    """
    Canonicalize object to deterministic JSON string.

    - Sorted keys
    - No whitespace
    - Consistent float formatting
    - Handles dataclasses via asdict
    """
    def serialize(o: Any) -> Any:
        if hasattr(o, "__dataclass_fields__"):
            return {k: serialize(v) for k, v in asdict(o).items()}
        elif isinstance(o, dict):
            return {k: serialize(v) for k, v in sorted(o.items())}
        elif isinstance(o, (list, tuple)):
            return [serialize(v) for v in o]
        elif isinstance(o, float):
            # Consistent float formatting (avoid 1.0 vs 1.0000000000000001)
            if o == int(o):
                return int(o)
            return round(o, 10)
        else:
            return o

    return json.dumps(serialize(obj), sort_keys=True, separators=(',', ':'))


def sha256_canonical(obj: Any) -> str:
    """Compute SHA256 of canonicalized object."""
    return hashlib.sha256(canonicalize(obj).encode()).hexdigest()


# =============================================================================
# TRAIL: THE WHY
# =============================================================================

@dataclass
class ISAOp:
    """Single ISA operation executed."""
    op: str              # Opcode name (e.g., "LOAD", "DOT", "CDR_ABSORB")
    op_hex: str          # Hex opcode (e.g., "0x01", "0x06", "0xC1")
    args: List[Any]      # Operation arguments
    cost_cycles: int     # CPU cycles consumed
    cost_flops: int = 0  # Floating point operations
    cost_bytes: int = 0  # Memory bytes accessed


@dataclass
class ISATrace:
    """ISA execution trace for a step."""
    ops_executed: List[ISAOp]
    ops_summary: Dict[str, int] = field(default_factory=lambda: {
        "total_ops": 0,
        "total_cycles": 0,
        "total_flops": 0,
        "total_bytes": 0,
    })

    def __post_init__(self):
        if self.ops_executed and self.ops_summary["total_ops"] == 0:
            self.ops_summary = {
                "total_ops": len(self.ops_executed),
                "total_cycles": sum(op.cost_cycles for op in self.ops_executed),
                "total_flops": sum(op.cost_flops for op in self.ops_executed),
                "total_bytes": sum(op.cost_bytes for op in self.ops_executed),
            }


@dataclass
class SeInput:
    """Symbolic Entropy input components."""
    H_entropy: float    # Shannon entropy (0-8 bits)
    C_coherence: float  # Contextual coherence (0-1)
    D_depth: float      # Dimensional depth


@dataclass
class Routing:
    """Se-based routing decision trace."""
    se_input: SeInput
    se_computed: float         # H * C * D
    route_decision: str        # "CPU", "GPU", "QPU", "core_0", etc.
    route_reason: str          # Human-readable explanation

    def __post_init__(self):
        if self.se_computed == 0 and self.se_input:
            self.se_computed = (
                self.se_input.H_entropy *
                self.se_input.C_coherence *
                self.se_input.D_depth
            )


@dataclass
class TopKEntry:
    """Single entry in top-k decision alternatives."""
    id: int              # Token ID or choice index
    logit: float         # Raw logit value
    prob: float          # Probability after softmax
    rank: int            # Rank in sorted order (1 = highest)
    label: str = ""      # Optional human-readable label


@dataclass
class DecisionMargins:
    """Decision margin statistics."""
    top1_gap: float         # logit[0] - logit[1]
    teacher_gap: float      # logit[teacher_choice] - logit[model_choice] (if different)
    entropy: float          # Entropy of probability distribution
    kl_divergence: float    # KL vs reference (if available)


@dataclass
class DecisionEvidence:
    """Evidence for a decision point."""
    task_type: str                     # "token_selection", "tensor_decode", "route_selection"
    decision_made: Any                 # The actual choice (token ID, route, etc.)
    top_k: List[TopKEntry]             # Top-k alternatives with scores
    margins: DecisionMargins
    confidence_classification: str     # "near_tie", "clear_winner", "uncertain"

    @classmethod
    def classify_confidence(cls, top1_gap: float, entropy: float) -> str:
        """Classify decision confidence based on margins."""
        if top1_gap < 0.5 or entropy > 2.0:
            return "near_tie"
        elif top1_gap > 3.0 and entropy < 1.0:
            return "clear_winner"
        else:
            return "moderate"


@dataclass
class GateCheck:
    """Single gate check result."""
    gate_id: str          # "impl_pinned", "fidelity_threshold", etc.
    passed: bool
    threshold: Any        # Expected value/threshold
    actual: Any           # Measured value
    description: str = "" # Human-readable description


@dataclass
class Guardrails:
    """Guardrail gate checks and verdict."""
    gates_checked: List[GateCheck]
    verdict: str                      # "PASS", "FAIL", "ACCEPTABLE_WITH_TAIL_RISK"
    tail_flags: List[str] = field(default_factory=list)  # ["p99_gap=4.72>2.0", ...]

    def compute_verdict(self) -> str:
        """Compute verdict from gate checks."""
        failed = [g for g in self.gates_checked if not g.passed]
        if not failed:
            return "PASS"

        # Check if all failures are "acceptable" (soft gates)
        hard_failures = [g for g in failed if not g.gate_id.startswith("soft_")]
        if hard_failures:
            return "FAIL"

        return "ACCEPTABLE_WITH_TAIL_RISK"


@dataclass
class Trail:
    """
    The "WHY" layer - explains the decision-making process.

    Human-readable audit trail of what happened and why.
    """
    isa_trace: Optional[ISATrace] = None
    routing: Optional[Routing] = None
    decision_evidence: Optional[DecisionEvidence] = None
    guardrails: Optional[Guardrails] = None


# =============================================================================
# SEAL: TAMPER-EVIDENT
# =============================================================================

@dataclass
class MerkleProof:
    """Merkle proof for random-access verification."""
    leaf_index: int
    leaf_hash: str                           # SHA256 of the leaf
    siblings: List[Dict[str, Any]]           # [{hash: str, position: "left"|"right"}]
    root_hash: str


@dataclass
class Seal:
    """
    Tamper-evident cryptographic seal.

    The "wax seal" that proves integrity.
    """
    trail_hash: str         # SHA256(canonicalized trail)
    state_hash: str         # SHA256(CDR sponge state)
    inputs_hash: str        # SHA256(canonicalized inputs)
    outputs_hash: str       # SHA256(canonicalized outputs)
    merkle_proof: Optional[MerkleProof] = None
    chain_hash: str = ""    # SHA256(parent_chain || trail_hash || state_hash)

    def compute_chain_hash(self, parent_chain_hash: str) -> str:
        """Compute chain hash linking to parent."""
        preimage = f"{parent_chain_hash}:{self.trail_hash}:{self.state_hash}"
        self.chain_hash = hashlib.sha256(preimage.encode()).hexdigest()
        return self.chain_hash


# =============================================================================
# CDR ANCHOR
# =============================================================================

@dataclass
class CDRAnchor:
    """
    CDR (Chained Deterministic Regeneration) anchor point.

    Enables reversible rollback to this state.
    """
    checkpoint_id: str
    cdr_step_count: int
    cdr_state_hash: str
    reversible: bool = True    # Can we rollback past this point?
    message_count: int = 0     # Messages absorbed since last checkpoint


# =============================================================================
# PROVENANCE
# =============================================================================

@dataclass
class Provenance:
    """Execution provenance metadata."""
    git_commit: Optional[str] = None
    impl_sha256: str = ""
    cpu_flags: str = ""
    python_version: str = ""
    platform: str = ""
    machine: str = ""
    hostname: str = ""

    @classmethod
    def from_environment(cls) -> "Provenance":
        """Capture provenance from current environment."""
        git_commit = None
        try:
            result = subprocess.run(
                ["git", "rev-parse", "HEAD"],
                capture_output=True, text=True, timeout=2
            )
            if result.returncode == 0:
                git_commit = result.stdout.strip()
        except Exception:
            pass

        # Get CPU flags
        cpu_flags = "unknown"
        try:
            if os.path.exists("/proc/cpuinfo"):
                with open("/proc/cpuinfo") as f:
                    for line in f:
                        if line.startswith("flags"):
                            flags = line.split(":")[1].strip()
                            key_flags = []
                            for flag in ["avx2", "avx", "sse4_2", "sse4_1", "neon"]:
                                if flag in flags:
                                    key_flags.append(flag)
                            cpu_flags = ",".join(key_flags) if key_flags else "none"
                            break
        except Exception:
            pass

        return cls(
            git_commit=git_commit,
            cpu_flags=cpu_flags,
            python_version=platform.python_version(),
            platform=platform.system().lower(),
            machine=platform.machine(),
            hostname=platform.node()[:16],  # Truncate for privacy
        )


# =============================================================================
# INFERENCE RECEIPT v1
# =============================================================================

@dataclass
class InferenceReceipt:
    """
    Complete inference receipt with explainable trail and cryptographic seal.

    Structure:
    - step_id: Unique identifier for this step
    - timestamp_utc: When this step occurred
    - parent_step_id: Link to previous step (for chaining)
    - trail: The WHY (ISA trace, routing, decisions, guardrails)
    - seal: Tamper-evident hashes and proofs
    - cdr_anchor: Reversible checkpoint anchor
    - provenance: Execution environment metadata
    """
    step_id: str
    timestamp_utc: str
    parent_step_id: Optional[str]

    trail: Trail
    seal: Seal
    cdr_anchor: Optional[CDRAnchor]
    provenance: Provenance

    schema_version: str = SCHEMA_VERSION

    @classmethod
    def create(
        cls,
        step_id: str,
        trail: Trail,
        inputs: Any,
        outputs: Any,
        cdr_state_hash: str,
        parent_step_id: Optional[str] = None,
        parent_chain_hash: str = "",
        cdr_anchor: Optional[CDRAnchor] = None,
        merkle_proof: Optional[MerkleProof] = None,
    ) -> "InferenceReceipt":
        """
        Factory method to create a complete receipt with computed seals.

        Args:
            step_id: Unique step identifier
            trail: Trail object with ISA/routing/decision traces
            inputs: Input data (will be canonicalized and hashed)
            outputs: Output data (will be canonicalized and hashed)
            cdr_state_hash: Current CDR sponge state hash
            parent_step_id: Previous step ID (for linking)
            parent_chain_hash: Previous chain hash (for chaining)
            cdr_anchor: Optional CDR checkpoint anchor
            merkle_proof: Optional Merkle proof

        Returns:
            Complete InferenceReceipt
        """
        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

        # Compute seal hashes
        trail_hash = sha256_canonical(trail)
        inputs_hash = sha256_canonical(inputs) if inputs else hashlib.sha256(b"").hexdigest()
        outputs_hash = sha256_canonical(outputs) if outputs else hashlib.sha256(b"").hexdigest()

        seal = Seal(
            trail_hash=trail_hash,
            state_hash=cdr_state_hash,
            inputs_hash=inputs_hash,
            outputs_hash=outputs_hash,
            merkle_proof=merkle_proof,
        )

        # Compute chain hash
        if parent_chain_hash:
            seal.compute_chain_hash(parent_chain_hash)
        else:
            # Genesis receipt - chain hash is just trail + state
            seal.chain_hash = hashlib.sha256(
                f"{trail_hash}:{cdr_state_hash}".encode()
            ).hexdigest()

        return cls(
            step_id=step_id,
            timestamp_utc=timestamp,
            parent_step_id=parent_step_id,
            trail=trail,
            seal=seal,
            cdr_anchor=cdr_anchor,
            provenance=Provenance.from_environment(),
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to serializable dict."""
        return asdict(self)

    def to_json(self, indent: int = 2) -> str:
        """Serialize to JSON."""
        return json.dumps(self.to_dict(), indent=indent, sort_keys=True, default=str)

    def to_canonical_json(self) -> str:
        """Serialize to canonical (deterministic) JSON."""
        return canonicalize(self)

    def receipt_hash(self) -> str:
        """Compute hash of the complete receipt."""
        return hashlib.sha256(self.to_canonical_json().encode()).hexdigest()

    @classmethod
    def from_json(cls, json_str: str) -> "InferenceReceipt":
        """Deserialize from JSON string."""
        data = json.loads(json_str)
        return cls.from_dict(data)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "InferenceReceipt":
        """Reconstruct from dict representation."""
        # Reconstruct nested dataclasses
        trail_data = data.get("trail", {})
        trail = Trail(
            isa_trace=_reconstruct_isa_trace(trail_data.get("isa_trace")),
            routing=_reconstruct_routing(trail_data.get("routing")),
            decision_evidence=_reconstruct_decision_evidence(trail_data.get("decision_evidence")),
            guardrails=_reconstruct_guardrails(trail_data.get("guardrails")),
        )

        seal_data = data.get("seal", {})
        merkle_proof_data = seal_data.get("merkle_proof")
        seal = Seal(
            trail_hash=seal_data.get("trail_hash", ""),
            state_hash=seal_data.get("state_hash", ""),
            inputs_hash=seal_data.get("inputs_hash", ""),
            outputs_hash=seal_data.get("outputs_hash", ""),
            merkle_proof=MerkleProof(**merkle_proof_data) if merkle_proof_data else None,
            chain_hash=seal_data.get("chain_hash", ""),
        )

        cdr_anchor_data = data.get("cdr_anchor")
        cdr_anchor = CDRAnchor(**cdr_anchor_data) if cdr_anchor_data else None

        prov_data = data.get("provenance", {})
        provenance = Provenance(**prov_data)

        return cls(
            step_id=data.get("step_id", ""),
            timestamp_utc=data.get("timestamp_utc", ""),
            parent_step_id=data.get("parent_step_id"),
            trail=trail,
            seal=seal,
            cdr_anchor=cdr_anchor,
            provenance=provenance,
            schema_version=data.get("schema_version", SCHEMA_VERSION),
        )

    def verify_chain(self, parent_receipt: Optional["InferenceReceipt"]) -> bool:
        """Verify chain hash links to parent receipt."""
        if parent_receipt is None:
            # Genesis receipt - no parent to verify
            expected = hashlib.sha256(
                f"{self.seal.trail_hash}:{self.seal.state_hash}".encode()
            ).hexdigest()
            return self.seal.chain_hash == expected

        # Verify chain link
        expected = hashlib.sha256(
            f"{parent_receipt.seal.chain_hash}:{self.seal.trail_hash}:{self.seal.state_hash}".encode()
        ).hexdigest()
        return self.seal.chain_hash == expected

    def verify_trail_hash(self) -> bool:
        """Verify trail hash matches current trail."""
        computed = sha256_canonical(self.trail)
        return computed == self.seal.trail_hash


# =============================================================================
# RECONSTRUCTION HELPERS
# =============================================================================

def _reconstruct_isa_trace(data: Optional[Dict]) -> Optional[ISATrace]:
    if not data:
        return None
    ops = [ISAOp(**op) for op in data.get("ops_executed", [])]
    return ISATrace(ops_executed=ops, ops_summary=data.get("ops_summary", {}))


def _reconstruct_routing(data: Optional[Dict]) -> Optional[Routing]:
    if not data:
        return None
    se_input_data = data.get("se_input", {})
    se_input = SeInput(**se_input_data) if se_input_data else SeInput(0, 0, 0)
    return Routing(
        se_input=se_input,
        se_computed=data.get("se_computed", 0),
        route_decision=data.get("route_decision", ""),
        route_reason=data.get("route_reason", ""),
    )


def _reconstruct_decision_evidence(data: Optional[Dict]) -> Optional[DecisionEvidence]:
    if not data:
        return None
    top_k = [TopKEntry(**e) for e in data.get("top_k", [])]
    margins_data = data.get("margins", {})
    margins = DecisionMargins(**margins_data) if margins_data else DecisionMargins(0, 0, 0, 0)
    return DecisionEvidence(
        task_type=data.get("task_type", ""),
        decision_made=data.get("decision_made"),
        top_k=top_k,
        margins=margins,
        confidence_classification=data.get("confidence_classification", "unknown"),
    )


def _reconstruct_guardrails(data: Optional[Dict]) -> Optional[Guardrails]:
    if not data:
        return None
    gates = [GateCheck(**g) for g in data.get("gates_checked", [])]
    return Guardrails(
        gates_checked=gates,
        verdict=data.get("verdict", "UNKNOWN"),
        tail_flags=data.get("tail_flags", []),
    )


# =============================================================================
# DEMO
# =============================================================================

if __name__ == "__main__":
    print("InferenceReceipt v1 Demo")
    print("=" * 60)

    # Create sample ISA trace
    isa_trace = ISATrace(ops_executed=[
        ISAOp(op="LOAD", op_hex="0x01", args=[0, 100], cost_cycles=2),
        ISAOp(op="DOT", op_hex="0x06", args=[0, 1, 2], cost_cycles=100, cost_flops=4096),
        ISAOp(op="CDR_ABSORB", op_hex="0xC1", args=[64], cost_cycles=50),
    ])

    # Create sample routing
    routing = Routing(
        se_input=SeInput(H_entropy=5.33, C_coherence=0.87, D_depth=8),
        se_computed=0,  # Will be computed in __post_init__
        route_decision="core_0",
        route_reason="Low Se score (37.1) routes to primary core"
    )

    # Create sample decision evidence
    decision = DecisionEvidence(
        task_type="token_selection",
        decision_made=1234,
        top_k=[
            TopKEntry(id=1234, logit=8.5, prob=0.45, rank=1, label="the"),
            TopKEntry(id=5678, logit=7.9, prob=0.32, rank=2, label="a"),
            TopKEntry(id=9012, logit=6.2, prob=0.12, rank=3, label="an"),
        ],
        margins=DecisionMargins(top1_gap=0.6, teacher_gap=0.0, entropy=1.2, kl_divergence=0.05),
        confidence_classification="moderate"
    )

    # Create guardrails
    guardrails = Guardrails(
        gates_checked=[
            GateCheck("impl_pinned", True, "helix64_v1", "helix64_v1", "Impl SHA256 matches"),
            GateCheck("fidelity_threshold", True, 0.001, 0.0004, "Max error < threshold"),
            GateCheck("soft_p99_gap", False, 2.0, 4.72, "P99 logit gap within tolerance"),
        ],
        verdict="ACCEPTABLE_WITH_TAIL_RISK",
        tail_flags=["p99_gap=4.72>2.0"],
    )

    # Create trail
    trail = Trail(
        isa_trace=isa_trace,
        routing=routing,
        decision_evidence=decision,
        guardrails=guardrails,
    )

    # Create receipt
    receipt = InferenceReceipt.create(
        step_id="step_812_5",
        trail=trail,
        inputs={"prompt": "What is 2+2?"},
        outputs={"token": 1234, "text": "4"},
        cdr_state_hash="abc123def456",
        parent_step_id="step_812_4",
        parent_chain_hash="parent_chain_hash_here",
    )

    print("\nGenerated Receipt:")
    print("-" * 40)
    print(receipt.to_json())

    print(f"\nReceipt Hash: {receipt.receipt_hash()[:32]}...")
    print(f"Trail Hash:   {receipt.seal.trail_hash[:32]}...")
    print(f"Chain Hash:   {receipt.seal.chain_hash[:32]}...")

    # Verify roundtrip
    json_str = receipt.to_json()
    restored = InferenceReceipt.from_json(json_str)
    print(f"\nRoundtrip OK: {receipt.receipt_hash() == restored.receipt_hash()}")
    print(f"Trail hash verified: {receipt.verify_trail_hash()}")
