"""
Lineage Governor for Loom Receipt Infrastructure.

Enforces lineage rules with policy, not caller discipline.

The LineageGovernor ensures:
1. Parent receipts exist in the vault before being referenced
2. Temporal ordering (parent.transaction_time <= child.observed_at)
3. Chain depth limits to prevent infinite chains
4. Merkle proofs for parent→child relationships
5. Attestation level propagation (child can't exceed weakest parent)

Usage:
    from receipts.lineage_governor import get_lineage_governor
    
    governor = get_lineage_governor()
    
    # Attach parents with validation
    governor.attach_parents(child_receipt, parent_ids=["rcpt-ABC", "rcpt-DEF"])
    
    # Validate entire chain
    result = governor.validate_chain("rcpt-XYZ")
    if not result.valid:
        print(f"Chain broken: {result.violations}")
"""

from __future__ import annotations

import hashlib
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, Dict, List, Optional, Protocol, Set
from enum import Enum

from loom.proof_config import AttestationLevel
from receipts.policy_registry import LineagePolicy

logger = logging.getLogger(__name__)


# =============================================================================
# Protocols
# =============================================================================

class ReceiptVault(Protocol):
    """Protocol for receipt storage."""
    
    def exists(self, receipt_id: str) -> bool:
        """Check if receipt exists."""
        ...
    
    def get(self, receipt_id: str) -> Optional[Dict[str, Any]]:
        """Get receipt by ID."""
        ...
    
    def get_many(self, receipt_ids: List[str]) -> Dict[str, Dict[str, Any]]:
        """Get multiple receipts. Returns {id: receipt} for found receipts."""
        ...


# =============================================================================
# Errors
# =============================================================================

class LineageError(Exception):
    """Base error for lineage validation."""
    pass


class ParentNotFoundError(LineageError):
    """Parent receipt not found in vault."""
    pass


class TemporalViolationError(LineageError):
    """Temporal ordering violation in lineage."""
    pass


class ChainDepthError(LineageError):
    """Chain exceeds maximum depth."""
    pass


class AttestationDowngradeError(LineageError):
    """Child attestation exceeds parent attestation."""
    pass


class CircularReferenceError(LineageError):
    """Circular reference detected in lineage."""
    pass


# =============================================================================
# Lineage Proof
# =============================================================================

@dataclass
class LineageProof:
    """
    Proof of parent→child relationship.
    
    Contains Merkle-style hash linking parent to child.
    """
    child_id: str
    parent_ids: List[str]
    parent_hashes: Dict[str, str]  # parent_id → hash
    combined_hash: str  # Hash of all parent hashes
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "child_id": self.child_id,
            "parent_ids": self.parent_ids,
            "parent_hashes": self.parent_hashes,
            "combined_hash": self.combined_hash,
            "timestamp": self.timestamp.isoformat(),
        }


@dataclass
class ChainValidation:
    """Result of chain validation."""
    valid: bool
    chain_length: int
    violations: List[str] = field(default_factory=list)
    receipts_checked: List[str] = field(default_factory=list)
    weakest_attestation: Optional[AttestationLevel] = None
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "valid": self.valid,
            "chain_length": self.chain_length,
            "violations": self.violations,
            "receipts_checked": self.receipts_checked,
            "weakest_attestation": self.weakest_attestation.value if self.weakest_attestation else None,
        }


# =============================================================================
# Simple In-Memory Vault (for testing)
# =============================================================================

class InMemoryVault:
    """Simple in-memory vault for testing."""
    
    def __init__(self):
        self._receipts: Dict[str, Dict[str, Any]] = {}
    
    def store(self, receipt: Dict[str, Any]) -> None:
        receipt_id = receipt.get("receipt_id", "")
        if receipt_id:
            self._receipts[receipt_id] = receipt
    
    def exists(self, receipt_id: str) -> bool:
        return receipt_id in self._receipts
    
    def get(self, receipt_id: str) -> Optional[Dict[str, Any]]:
        return self._receipts.get(receipt_id)
    
    def get_many(self, receipt_ids: List[str]) -> Dict[str, Dict[str, Any]]:
        return {rid: self._receipts[rid] for rid in receipt_ids if rid in self._receipts}


# =============================================================================
# Lineage Governor
# =============================================================================

class LineageGovernor:
    """
    Governs lineage with policy, not caller discipline.
    
    All parent attachments must go through this class for validation.
    """
    
    def __init__(
        self,
        vault: ReceiptVault,
        default_policy: Optional[LineagePolicy] = None,
    ):
        self.vault = vault
        self.default_policy = default_policy or LineagePolicy()
        
        # Cache for performance
        self._hash_cache: Dict[str, str] = {}
    
    def attach_parents(
        self,
        child: Dict[str, Any],
        parent_ids: List[str],
        policy: Optional[LineagePolicy] = None,
    ) -> LineageProof:
        """
        Attach parents to a child receipt with full validation.
        
        Args:
            child: Child receipt dictionary
            parent_ids: List of parent receipt IDs
            policy: Lineage policy to enforce (defaults to self.default_policy)
        
        Returns:
            LineageProof containing the relationship proof
            
        Raises:
            ParentNotFoundError: If parent not in vault and policy requires it
            TemporalViolationError: If temporal ordering violated
            AttestationDowngradeError: If child attestation exceeds parents
        """
        policy = policy or self.default_policy
        child_id = child.get("receipt_id", "")
        
        # Get all parents
        parents = self.vault.get_many(parent_ids)
        
        # Check parent existence
        if policy.require_parent_in_vault:
            missing = set(parent_ids) - set(parents.keys())
            if missing:
                raise ParentNotFoundError(
                    f"Parents not found in vault: {missing}. "
                    f"Policy requires parent_in_vault=True"
                )
        
        # Check temporal ordering
        if policy.enforce_temporal_order:
            child_observed = self._get_timestamp(child, "observed_at") or self._get_timestamp(child, "created_at")
            if child_observed:
                for pid, parent in parents.items():
                    parent_tx = self._get_timestamp(parent, "transaction_time") or self._get_timestamp(parent, "created_at")
                    if parent_tx and parent_tx > child_observed:
                        raise TemporalViolationError(
                            f"Parent {pid} transaction_time ({parent_tx}) > "
                            f"child observed_at ({child_observed})"
                        )
        
        # Check attestation levels (child can't claim higher than weakest parent)
        child_attestation = self._get_attestation_level(child)
        for pid, parent in parents.items():
            parent_attestation = self._get_attestation_level(parent)
            if child_attestation > parent_attestation:
                raise AttestationDowngradeError(
                    f"Child attestation {child_attestation.value} exceeds "
                    f"parent {pid} attestation {parent_attestation.value}"
                )
        
        # Build lineage proof
        parent_hashes = {}
        for pid, parent in parents.items():
            parent_hashes[pid] = self._compute_receipt_hash(parent)
        
        # Combined hash of all parent hashes (sorted for determinism)
        sorted_hashes = [parent_hashes[pid] for pid in sorted(parent_hashes.keys())]
        combined = hashlib.sha256("|".join(sorted_hashes).encode()).hexdigest()
        
        proof = LineageProof(
            child_id=child_id,
            parent_ids=list(parent_ids),
            parent_hashes=parent_hashes,
            combined_hash=f"sha256:{combined}",
        )
        
        # Update child with lineage metadata
        child["_lineage"] = proof.to_dict()
        child["_lineage_validated"] = True
        child["parents"] = parent_ids
        
        logger.debug(f"Attached {len(parent_ids)} parents to {child_id}")
        
        return proof
    
    def validate_parents(
        self,
        child_id: str,
        parent_ids: List[str],
        policy: Optional[LineagePolicy] = None,
    ) -> List[str]:
        """
        Validate parent relationships without modifying receipt.
        
        Returns list of violations (empty = valid).
        """
        policy = policy or self.default_policy
        violations = []
        
        # Get parents
        parents = self.vault.get_many(parent_ids)
        
        # Check existence
        if policy.require_parent_in_vault:
            missing = set(parent_ids) - set(parents.keys())
            if missing:
                violations.append(f"Parents not found: {missing}")
        
        # We can't check temporal ordering without the child receipt
        # This is a lightweight validation for policy_writer
        
        return violations
    
    def validate_chain(
        self,
        leaf_id: str,
        policy: Optional[LineagePolicy] = None,
        max_depth: Optional[int] = None,
    ) -> ChainValidation:
        """
        Validate entire chain from leaf back to roots.
        
        Args:
            leaf_id: Receipt ID to start from
            policy: Lineage policy to enforce
            max_depth: Maximum chain depth to traverse
        
        Returns:
            ChainValidation with results
        """
        policy = policy or self.default_policy
        max_depth = max_depth or policy.max_chain_depth
        
        violations = []
        checked = []
        visited: Set[str] = set()
        weakest_attestation = AttestationLevel.L5_COURT_GRADE
        
        def traverse(receipt_id: str, depth: int) -> None:
            nonlocal weakest_attestation
            
            if depth > max_depth:
                violations.append(f"Chain depth {depth} exceeds max {max_depth}")
                return
            
            if receipt_id in visited:
                violations.append(f"Circular reference detected at {receipt_id}")
                return
            
            visited.add(receipt_id)
            checked.append(receipt_id)
            
            receipt = self.vault.get(receipt_id)
            if receipt is None:
                if policy.require_parent_in_vault:
                    violations.append(f"Receipt {receipt_id} not found in vault")
                return
            
            # Track weakest attestation
            attestation = self._get_attestation_level(receipt)
            if attestation < weakest_attestation:
                weakest_attestation = attestation
            
            # Get parents and recurse
            parents = receipt.get("parents", [])
            if isinstance(parents, list):
                for parent_id in parents:
                    traverse(parent_id, depth + 1)
        
        # Start traversal
        traverse(leaf_id, 0)
        
        return ChainValidation(
            valid=len(violations) == 0,
            chain_length=len(checked),
            violations=violations,
            receipts_checked=checked,
            weakest_attestation=weakest_attestation if checked else None,
        )
    
    def compute_chain_attestation(self, leaf_id: str) -> AttestationLevel:
        """
        Compute the effective attestation level for a chain.
        
        The chain's attestation is the minimum of all receipts in the chain.
        """
        result = self.validate_chain(leaf_id)
        return result.weakest_attestation or AttestationLevel.L0_DEV
    
    def _get_timestamp(self, receipt: Dict[str, Any], field: str) -> Optional[datetime]:
        """Extract timestamp from receipt."""
        value = receipt.get(field)
        if value is None:
            return None
        if isinstance(value, datetime):
            return value
        if isinstance(value, str):
            try:
                # Try ISO format
                return datetime.fromisoformat(value.replace("Z", "+00:00"))
            except ValueError:
                return None
        return None
    
    def _get_attestation_level(self, receipt: Dict[str, Any]) -> AttestationLevel:
        """Extract attestation level from receipt."""
        attestation = receipt.get("_attestation", {})
        level_str = attestation.get("attestation_level", "L0_DEV")
        try:
            return AttestationLevel(level_str)
        except ValueError:
            return AttestationLevel.L0_DEV
    
    def _compute_receipt_hash(self, receipt: Dict[str, Any]) -> str:
        """Compute hash of receipt for lineage proof."""
        receipt_id = receipt.get("receipt_id", "")
        
        # Check cache
        if receipt_id in self._hash_cache:
            return self._hash_cache[receipt_id]
        
        # Use existing hash if present
        existing = receipt.get("jcs_hash") or receipt.get("_jcs_hash")
        if existing:
            self._hash_cache[receipt_id] = existing
            return existing
        
        # Compute from receipt ID + timestamp (lightweight)
        ts = receipt.get("created_at", "")
        data = f"{receipt_id}:{ts}"
        computed = f"sha256:{hashlib.sha256(data.encode()).hexdigest()}"
        
        self._hash_cache[receipt_id] = computed
        return computed


# =============================================================================
# Global Governor Access
# =============================================================================

_global_governor: Optional[LineageGovernor] = None


def get_lineage_governor(vault: Optional[ReceiptVault] = None) -> LineageGovernor:
    """
    Get the global lineage governor.
    
    Creates governor with in-memory vault on first call if none configured.
    """
    global _global_governor
    
    if _global_governor is None:
        if vault is None:
            vault = InMemoryVault()
        _global_governor = LineageGovernor(vault=vault)
    
    return _global_governor


def configure_lineage_governor(governor: LineageGovernor) -> None:
    """Explicitly set the global lineage governor."""
    global _global_governor
    _global_governor = governor


def reset_lineage_governor() -> None:
    """Reset global governor (for testing)."""
    global _global_governor
    _global_governor = None


# =============================================================================
# Convenience Functions
# =============================================================================

def attach_parents(
    child: Dict[str, Any],
    parent_ids: List[str],
    policy: Optional[LineagePolicy] = None,
) -> LineageProof:
    """
    Convenience function to attach parents using global governor.
    
    Usage:
        from receipts.lineage_governor import attach_parents
        
        proof = attach_parents(my_receipt, ["parent-1", "parent-2"])
    """
    return get_lineage_governor().attach_parents(child, parent_ids, policy)


def validate_chain(receipt_id: str) -> ChainValidation:
    """
    Convenience function to validate chain using global governor.
    
    Usage:
        from receipts.lineage_governor import validate_chain
        
        result = validate_chain("rcpt-XYZ")
        if not result.valid:
            print(f"Issues: {result.violations}")
    """
    return get_lineage_governor().validate_chain(receipt_id)


def is_chain_valid(receipt_id: str) -> bool:
    """Quick check if chain is valid."""
    return validate_chain(receipt_id).valid
