#!/usr/bin/env python3
"""
emitter.py - Receipt Emitter with Merkle Chain Integration

Combines Trail + Seal + CDR anchor into complete receipts with:
- Automatic chain hash linking
- Merkle tree proof generation
- Persistence to disk
- Integration with CDRStateManager

Usage:
    from helix_cdc.receipts import ReceiptEmitter
    from helix_cdc.cdr_state_manager import CDRStateManager

    cdr = CDRStateManager()
    emitter = ReceiptEmitter(cdr_manager=cdr, output_dir="receipts/run_001/")

    # Emit receipts as inference proceeds
    receipt = emitter.emit(
        step_id="step_5",
        trail=trail,
        inputs={"prompt": "Hello"},
        outputs={"token": 42},
    )

    # Verify any step with O(log n) proof
    proof = emitter.get_proof(step_index=5)
    assert emitter.verify_step(proof, emitter.get_root())

    # Finalize run and write summary
    emitter.finalize()
"""

from __future__ import annotations
import json
import time
import hashlib
from pathlib import Path
from dataclasses import dataclass, field
from typing import Dict, List, Any, Optional, TYPE_CHECKING

from .inference_receipt_v1 import (
    InferenceReceipt,
    Trail,
    CDRAnchor,
    MerkleProof,
    canonicalize,
)
from .merkle import MerkleTree

if TYPE_CHECKING:
    from helix_cdc.cdr_state_manager import CDRStateManager


@dataclass
class EmitterStats:
    """Statistics for the receipt emitter."""
    receipts_emitted: int = 0
    total_bytes_written: int = 0
    first_timestamp: str = ""
    last_timestamp: str = ""
    chain_length: int = 0


class ReceiptEmitter:
    """
    Receipt emitter with Merkle chain integration.

    Features:
    - Automatic chain hash linking between receipts
    - Merkle tree for O(log n) proof generation
    - Integration with CDRStateManager for anchors
    - Persistence to disk with atomic writes
    """

    def __init__(
        self,
        cdr_manager: Optional["CDRStateManager"] = None,
        output_dir: Optional[Path] = None,
        merkle_hash_algo: str = "sha256",
    ):
        """
        Initialize receipt emitter.

        Args:
            cdr_manager: Optional CDRStateManager for state anchoring
            output_dir: Optional directory for persisting receipts
            merkle_hash_algo: Hash algorithm for Merkle tree
        """
        self.cdr_manager = cdr_manager
        self.output_dir = Path(output_dir) if output_dir else None
        self.merkle_tree = MerkleTree(hash_algo=merkle_hash_algo)

        # Chain state
        self.receipts: List[InferenceReceipt] = []
        self.parent_chain_hash: str = ""
        self.parent_step_id: Optional[str] = None

        # Stats
        self.stats = EmitterStats()

        # Ensure output directory exists
        if self.output_dir:
            self.output_dir.mkdir(parents=True, exist_ok=True)

    def emit(
        self,
        step_id: str,
        trail: Trail,
        inputs: Any,
        outputs: Any,
        cdr_checkpoint_id: Optional[str] = None,
    ) -> InferenceReceipt:
        """
        Emit a new receipt and add to the chain.

        Args:
            step_id: Unique step identifier
            trail: Trail object with ISA/routing/decision traces
            inputs: Input data for this step
            outputs: Output data from this step
            cdr_checkpoint_id: Optional checkpoint ID for CDR anchor

        Returns:
            Complete InferenceReceipt
        """
        # Get CDR state
        cdr_state_hash = "uninitialized"
        cdr_anchor = None

        if self.cdr_manager is not None:
            cdr_state_hash = self.cdr_manager.cdr_state_hash()

            if cdr_checkpoint_id:
                cdr_anchor = CDRAnchor(
                    checkpoint_id=cdr_checkpoint_id,
                    cdr_step_count=self.cdr_manager.cdr_state.step_count,
                    cdr_state_hash=cdr_state_hash,
                    reversible=True,
                    message_count=len(self.cdr_manager.message_history),
                )
        elif cdr_checkpoint_id:
            # No CDR manager but checkpoint requested - create minimal anchor
            cdr_anchor = CDRAnchor(
                checkpoint_id=cdr_checkpoint_id,
                cdr_step_count=len(self.receipts),
                cdr_state_hash="no_cdr_manager",
                reversible=False,
            )

        # Create receipt (without Merkle proof - will be added at finalize)
        receipt = InferenceReceipt.create(
            step_id=step_id,
            trail=trail,
            inputs=inputs,
            outputs=outputs,
            cdr_state_hash=cdr_state_hash,
            parent_step_id=self.parent_step_id,
            parent_chain_hash=self.parent_chain_hash,
            cdr_anchor=cdr_anchor,
        )

        # Add to Merkle tree (proof generated at finalize)
        receipt_hash = receipt.receipt_hash()
        self.merkle_tree.add_leaf(receipt_hash)

        # Update chain state
        self.receipts.append(receipt)
        self.parent_chain_hash = receipt.seal.chain_hash
        self.parent_step_id = step_id

        # Update stats
        self._update_stats(receipt)

        # Note: Receipts are persisted at finalize() with their Merkle proofs

        return receipt

    def emit_with_isa_trace(
        self,
        step_id: str,
        ops_executed: List[Dict[str, Any]],
        inputs: Any,
        outputs: Any,
        routing: Optional[Dict[str, Any]] = None,
        decision: Optional[Dict[str, Any]] = None,
        guardrails: Optional[Dict[str, Any]] = None,
    ) -> InferenceReceipt:
        """
        Convenience method to emit receipt with ISA trace from dict data.

        Args:
            step_id: Step identifier
            ops_executed: List of operation dicts [{op, op_hex, args, cost_cycles, ...}]
            inputs: Input data
            outputs: Output data
            routing: Optional routing dict
            decision: Optional decision evidence dict
            guardrails: Optional guardrails dict

        Returns:
            InferenceReceipt
        """
        from .inference_receipt_v1 import (
            ISAOp, ISATrace, Routing, SeInput,
            DecisionEvidence, TopKEntry, DecisionMargins,
            Guardrails, GateCheck,
        )

        # Build ISA trace
        isa_trace = None
        if ops_executed:
            ops = [ISAOp(**op) for op in ops_executed]
            isa_trace = ISATrace(ops_executed=ops)

        # Build routing
        routing_obj = None
        if routing:
            se_input_data = routing.get("se_input", {})
            routing_obj = Routing(
                se_input=SeInput(**se_input_data) if se_input_data else SeInput(0, 0, 0),
                se_computed=routing.get("se_computed", 0),
                route_decision=routing.get("route_decision", ""),
                route_reason=routing.get("route_reason", ""),
            )

        # Build decision evidence
        decision_obj = None
        if decision:
            top_k = [TopKEntry(**e) for e in decision.get("top_k", [])]
            margins_data = decision.get("margins", {})
            decision_obj = DecisionEvidence(
                task_type=decision.get("task_type", "unknown"),
                decision_made=decision.get("decision_made"),
                top_k=top_k,
                margins=DecisionMargins(**margins_data) if margins_data else DecisionMargins(0, 0, 0, 0),
                confidence_classification=decision.get("confidence_classification", "unknown"),
            )

        # Build guardrails
        guardrails_obj = None
        if guardrails:
            gates = [GateCheck(**g) for g in guardrails.get("gates_checked", [])]
            guardrails_obj = Guardrails(
                gates_checked=gates,
                verdict=guardrails.get("verdict", "UNKNOWN"),
                tail_flags=guardrails.get("tail_flags", []),
            )

        trail = Trail(
            isa_trace=isa_trace,
            routing=routing_obj,
            decision_evidence=decision_obj,
            guardrails=guardrails_obj,
        )

        return self.emit(step_id=step_id, trail=trail, inputs=inputs, outputs=outputs)

    def get_proof(self, step_index: int) -> Dict[str, Any]:
        """
        Get Merkle proof for a step.

        Args:
            step_index: Index of the step (0-based)

        Returns:
            Proof dict with leaf_hash, siblings, root_hash
        """
        return self.merkle_tree.get_proof(step_index)

    def get_root(self) -> str:
        """Get current Merkle root."""
        return self.merkle_tree.get_root()

    def verify_step(
        self,
        proof: Dict[str, Any],
        expected_root: Optional[str] = None
    ) -> bool:
        """
        Verify a step using its Merkle proof.

        Args:
            proof: Proof dict from get_proof()
            expected_root: Expected Merkle root (uses current if None)

        Returns:
            True if proof is valid
        """
        root = expected_root or self.get_root()
        return MerkleTree.verify_proof(
            proof["leaf_hash"],
            proof["siblings"],
            root,
            hash_algo=self.merkle_tree.hash_algo,
        )

    def verify_chain(self) -> Dict[str, Any]:
        """
        Verify the complete receipt chain.

        Returns:
            {
                "valid": bool,
                "steps_verified": int,
                "first_failure": Optional[int],
                "issues": [str, ...]
            }
        """
        issues = []
        first_failure = None

        for i, receipt in enumerate(self.receipts):
            parent = self.receipts[i - 1] if i > 0 else None

            # Verify chain link
            if not receipt.verify_chain(parent):
                issues.append(f"Step {i}: Chain hash verification failed")
                if first_failure is None:
                    first_failure = i

            # Verify trail hash
            if not receipt.verify_trail_hash():
                issues.append(f"Step {i}: Trail hash verification failed")
                if first_failure is None:
                    first_failure = i

        return {
            "valid": len(issues) == 0,
            "steps_verified": len(self.receipts),
            "first_failure": first_failure,
            "issues": issues,
        }

    def finalize(self) -> Dict[str, Any]:
        """
        Finalize the receipt chain and write summary.

        Generates Merkle proofs for all receipts (against the final root)
        and persists them to disk.

        Returns:
            Summary dict with root, stats, and verification result
        """
        # Generate Merkle proofs for all receipts (against final root)
        final_root = self.get_root()
        for i, receipt in enumerate(self.receipts):
            proof_data = self.merkle_tree.get_proof(i)
            receipt.seal.merkle_proof = MerkleProof(
                leaf_index=proof_data["leaf_index"],
                leaf_hash=proof_data["leaf_hash"],
                siblings=proof_data["siblings"],
                root_hash=final_root,  # Use final root, not intermediate
            )

        verification = self.verify_chain()

        summary = {
            "schema_version": "receipt_emitter_v1.0.0",
            "finalized_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "merkle_root": final_root,
            "receipt_count": len(self.receipts),
            "chain_valid": verification["valid"],
            "verification": verification,
            "stats": {
                "receipts_emitted": self.stats.receipts_emitted,
                "total_bytes_written": self.stats.total_bytes_written,
                "first_timestamp": self.stats.first_timestamp,
                "last_timestamp": self.stats.last_timestamp,
                "chain_length": self.stats.chain_length,
            },
        }

        if self.output_dir:
            # Write all receipts with their Merkle proofs
            for i, receipt in enumerate(self.receipts):
                self._persist_receipt(receipt, i)

            # Write summary
            summary_path = self.output_dir / "RECEIPT_SUMMARY.json"
            with open(summary_path, "w") as f:
                json.dump(summary, f, indent=2)

            # Write Merkle root for easy access
            root_path = self.output_dir / "merkle_root.txt"
            with open(root_path, "w") as f:
                f.write(final_root + "\n")

            # Write leaf index for lookup
            index_path = self.output_dir / "receipt_index.json"
            index = {
                "leaves": [
                    {
                        "index": i,
                        "step_id": r.step_id,
                        "receipt_hash": r.receipt_hash(),
                        "chain_hash": r.seal.chain_hash[:16] + "...",
                    }
                    for i, r in enumerate(self.receipts)
                ],
                "merkle_root": final_root,
            }
            with open(index_path, "w") as f:
                json.dump(index, f, indent=2)

        return summary

    def _update_stats(self, receipt: InferenceReceipt) -> None:
        """Update emitter statistics."""
        self.stats.receipts_emitted += 1
        self.stats.chain_length = len(self.receipts)
        self.stats.last_timestamp = receipt.timestamp_utc

        if not self.stats.first_timestamp:
            self.stats.first_timestamp = receipt.timestamp_utc

    def _persist_receipt(self, receipt: InferenceReceipt, index: int) -> None:
        """Persist receipt to disk with atomic write."""
        if not self.output_dir:
            return

        # Write receipt
        filename = f"step_{index:06d}_{receipt.step_id}.json"
        receipt_path = self.output_dir / filename
        tmp_path = receipt_path.with_suffix(".tmp")

        content = receipt.to_json()
        with open(tmp_path, "w") as f:
            f.write(content)

        # Atomic rename
        tmp_path.rename(receipt_path)

        self.stats.total_bytes_written += len(content)

    def get_receipt(self, step_index: int) -> Optional[InferenceReceipt]:
        """Get receipt by step index."""
        if 0 <= step_index < len(self.receipts):
            return self.receipts[step_index]
        return None

    def get_receipt_by_id(self, step_id: str) -> Optional[InferenceReceipt]:
        """Get receipt by step ID."""
        for receipt in self.receipts:
            if receipt.step_id == step_id:
                return receipt
        return None


# =============================================================================
# STANDALONE VERIFICATION
# =============================================================================

def verify_receipt_standalone(
    receipt_path: str,
    merkle_root: str,
    hash_algo: str = "sha256"
) -> Dict[str, Any]:
    """
    Verify a single receipt without loading the full chain.

    Args:
        receipt_path: Path to receipt JSON file
        merkle_root: Expected Merkle root
        hash_algo: Hash algorithm

    Returns:
        Verification result dict
    """
    with open(receipt_path) as f:
        receipt = InferenceReceipt.from_json(f.read())

    # Verify trail hash
    trail_valid = receipt.verify_trail_hash()

    # Verify Merkle proof
    merkle_valid = False
    if receipt.seal.merkle_proof:
        merkle_valid = MerkleTree.verify_proof(
            receipt.seal.merkle_proof.leaf_hash,
            receipt.seal.merkle_proof.siblings,
            merkle_root,
            hash_algo=hash_algo,
        )

    return {
        "step_id": receipt.step_id,
        "trail_hash_valid": trail_valid,
        "merkle_proof_valid": merkle_valid,
        "overall_valid": trail_valid and merkle_valid,
        "receipt_hash": receipt.receipt_hash()[:32] + "...",
    }


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

if __name__ == "__main__":
    import tempfile
    from .inference_receipt_v1 import (
        Trail, ISATrace, ISAOp, Routing, SeInput,
        DecisionEvidence, TopKEntry, DecisionMargins,
        Guardrails, GateCheck,
    )

    print("Receipt Emitter Demo")
    print("=" * 60)

    # Create emitter with temp directory
    with tempfile.TemporaryDirectory() as tmpdir:
        emitter = ReceiptEmitter(output_dir=Path(tmpdir))

        print(f"\n1. Emitting 5 receipts to {tmpdir}")

        for i in range(5):
            # Create sample trail
            trail = Trail(
                isa_trace=ISATrace(ops_executed=[
                    ISAOp(op="LOAD", op_hex="0x01", args=[i], cost_cycles=2),
                    ISAOp(op="ADD", op_hex="0x03", args=[i, 1], cost_cycles=1),
                ]),
                decision_evidence=DecisionEvidence(
                    task_type="token_selection",
                    decision_made=100 + i,
                    top_k=[
                        TopKEntry(id=100+i, logit=5.0, prob=0.8, rank=1),
                        TopKEntry(id=200+i, logit=3.0, prob=0.2, rank=2),
                    ],
                    margins=DecisionMargins(top1_gap=2.0, teacher_gap=0, entropy=0.5, kl_divergence=0.01),
                    confidence_classification="clear_winner",
                ),
            )

            receipt = emitter.emit(
                step_id=f"step_{i}",
                trail=trail,
                inputs={"step": i},
                outputs={"result": i * 2},
            )

            print(f"   Step {i}: chain_hash={receipt.seal.chain_hash[:16]}...")

        print(f"\n2. Merkle root: {emitter.get_root()[:32]}...")

        # Verify chain
        print("\n3. Chain verification:")
        verification = emitter.verify_chain()
        print(f"   Valid: {verification['valid']}")
        print(f"   Steps verified: {verification['steps_verified']}")

        # Get and verify proof for step 2
        print("\n4. Merkle proof for step 2:")
        proof = emitter.get_proof(2)
        print(f"   Leaf hash: {proof['leaf_hash'][:32]}...")
        print(f"   Siblings: {len(proof['siblings'])}")
        is_valid = emitter.verify_step(proof)
        print(f"   Verification: {'PASS' if is_valid else 'FAIL'}")

        # Finalize and show summary
        print("\n5. Finalizing:")
        summary = emitter.finalize()
        print(f"   Receipts: {summary['receipt_count']}")
        print(f"   Chain valid: {summary['chain_valid']}")
        print(f"   Root: {summary['merkle_root'][:32]}...")

        # List output files
        print("\n6. Output files:")
        for f in sorted(Path(tmpdir).iterdir()):
            print(f"   {f.name}")
