from __future__ import annotations

"""
EvidenceReceipt: atomic proof of sensory input at action time.
Hashes of raw inputs (audio, pixels, bio) become the "evidence" that
Guardian and auditors verify. Freshness gates latency; missing/stale
evidence blocks high-risk actions.
"""

import hashlib
import time
from dataclasses import dataclass
from typing import Dict, Literal, Optional, List
from .merkle import (
    compute_merkle_leaf_from_value,
    compute_merkle_root,
    verify_merkle_inclusion,
)


@dataclass(frozen=True)
class EvidenceReceipt:
    """Cryptographic proof of a multimodal sensory capture."""

    timestamp: float
    guardian_policy_version: str  # Version of Guardian policy used
    dna_hash: str  # Computed SHA-256 over all fields for integrity
    evidence_type: Literal["voice", "vision", "bio", "multimodal"]  # Generalized
    source: str  # e.g., "whisper_local", "screen_observer"
    metadata: Dict
    voice_hash: Optional[str] = None  # SHA-256 of voice input
    scene_hash: Optional[str] = None  # SHA-256 of vision input
    prediction_hash: Optional[str] = None  # SHA-256 of JEPA prediction
    freshness_ms: Optional[int] = None  # computed at attach time

    @classmethod
    def from_bytes(
        cls,
        evidence_type: Literal["voice", "vision", "bio"],
        blob: bytes,
        source: str,
        metadata: Optional[Dict] = None,
    ) -> "EvidenceReceipt":
        """Hash the raw bytes and mint an evidence receipt."""
        content_hash = hashlib.sha256(blob).hexdigest()
        timestamp = time.time()
        # For backward compatibility, set the appropriate hash
        if evidence_type == "voice":
            voice_hash = content_hash
            scene_hash = None
            prediction_hash = None
        elif evidence_type == "vision":
            voice_hash = None
            scene_hash = content_hash
            prediction_hash = None
        else:
            voice_hash = None
            scene_hash = None
            prediction_hash = None
        # Compute dna_hash
        dna_data = f"{voice_hash or ''}|{scene_hash or ''}|{prediction_hash or ''}|{timestamp}|{source}"
        dna_hash = hashlib.sha256(dna_data.encode()).hexdigest()
        return cls(
            timestamp=timestamp,
            guardian_policy_version="v1.0",  # Placeholder, update as needed
            dna_hash=dna_hash,
            evidence_type=evidence_type,
            source=source,
            metadata=metadata or {},
            voice_hash=voice_hash,
            scene_hash=scene_hash,
            prediction_hash=prediction_hash,
            freshness_ms=None,
        )

    @classmethod
    def unified(
        cls,
        voice_hash: Optional[str] = None,
        scene_hash: Optional[str] = None,
        prediction_hash: Optional[str] = None,
        source: str = "multimodal_observer",
        guardian_policy_version: str = "v1.0",
        metadata: Optional[Dict] = None,
    ) -> "EvidenceReceipt":
        """Create a unified multimodal evidence receipt."""
        timestamp = time.time()
        dna_data = f"{voice_hash or ''}|{scene_hash or ''}|{prediction_hash or ''}|{timestamp}|{guardian_policy_version}|{source}"
        dna_hash = hashlib.sha256(dna_data.encode()).hexdigest()
        return cls(
            timestamp=timestamp,
            guardian_policy_version=guardian_policy_version,
            dna_hash=dna_hash,
            evidence_type="multimodal",
            source=source,
            metadata=metadata or {},
            voice_hash=voice_hash,
            scene_hash=scene_hash,
            prediction_hash=prediction_hash,
            freshness_ms=None,
        )

    def is_fresh(self, max_age_ms: int = 500) -> bool:
        """Check if the capture is recent enough."""
        age_ms = (time.time() - self.timestamp) * 1000
        return age_ms < max_age_ms

    def verify_dna(self) -> bool:
        """Verify the DNA hash integrity."""
        dna_data = f"{self.voice_hash or ''}|{self.scene_hash or ''}|{self.prediction_hash or ''}|{self.timestamp}|{self.guardian_policy_version}|{self.source}"
        expected_dna = hashlib.sha256(dna_data.encode()).hexdigest()
        return self.dna_hash == expected_dna


def compute_merkle_leaf(receipt_or_hex) -> str:
    """Compute canonical merkle leaf from an EvidenceReceipt or a hex string value.

    Accepts either an EvidenceReceipt (uses its dna_hash) or a hex string
    that represents the underlying value to be hashed into a leaf.
    """
    if hasattr(receipt_or_hex, "dna_hash"):
        value_hex = receipt_or_hex.dna_hash
    elif isinstance(receipt_or_hex, str):
        value_hex = receipt_or_hex
    else:
        raise TypeError("compute_merkle_leaf expects EvidenceReceipt or hex string")

    return compute_merkle_leaf_from_value(value_hex)


def verify_receipt_dna(
    receipt: EvidenceReceipt,
    merkle_root: Optional[str] = None,
    merkle_proof: Optional[List[str]] = None,
    index: Optional[int] = None,
) -> bool:
    """High level verifier for receipts.

    This validates (1) the receipt's internal DNA hash matches its fields and
    (2) if a Merkle root and proof are provided, checks that the receipt's
    Merkle leaf is included in that root using the provided inclusion proof.
    """
    if not receipt.verify_dna():
        return False

    if merkle_root is not None:
        if merkle_proof is None or index is None:
            # Missing proof/index for Merkle inclusion check
            return False

        leaf = compute_merkle_leaf(receipt)
        return verify_merkle_inclusion(leaf, merkle_proof, merkle_root, index)

    return True
