"""
Idempotent merge receipt writer with Constitutional Computing proofs.

Integrates CHOIR MergeReceipt v2 with:
- IdempotentWriter (JCS canonical hashing + de-duplication)
- Tier-based proof generation (batch for SIGNED, individual for LEGAL)
- EventStore persistence (SQLite WAL hot tier + Parquet cold tier)
- CloudEvents envelope emission
"""

from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Any, List, Optional
import json
import logging

from receipts.models_merge import (
    MergeReceiptV2,
    ParentReceipt,
    MergeStrategy,
    Digest,
    ProofRefs,
    MergeMetadata,
    AntiGamingSignals,
    compute_effective_tier,
    compute_retention_days
)
from receipts.idempotent_writer import IdempotentWriter
from receipts.anti_gaming import (
    compute_anti_gaming_signals as compute_anti_gaming_signals_v2
)
from receipts.trustmark import compute_trustmark

logger = logging.getLogger(__name__)


class MergeReceiptWriter:
    """
    Writes CHOIR merge receipts with idempotency and tier-based proofs.

    Production workflow:
    1. Validate parents exist and load their receipts
    2. Compute effective tier (max of parent tiers)
    3. Execute merge strategy (via CHOIR)
    4. Compute JCS canonical hash (digest_after)
    5. Write via IdempotentWriter (de-duplication)
    6. Generate proofs based on tier:
       - RECORDED: hash only
       - SIGNED: batch proofs (Ed25519 + Rekor)
       - LEGAL: individual proofs (Ed25519 + ML-DSA + TSA + Rekor)
    7. Emit CloudEvents envelope
    """

    def __init__(
        self,
        writer: IdempotentWriter,
        event_store_path: Path,
        proofs_dir: Path,
        keystore_path: Optional[Path] = None
    ):
        self.writer = writer
        self.event_store_path = event_store_path
        self.proofs_dir = proofs_dir
        self.keystore_path = keystore_path or Path("api/keys")

    def write_merge_receipt(
        self,
        parents: List[Dict[str, Any]],
        strategy: Dict[str, Any],
        domain_data: Dict[str, Any],
        tenant_id: str,
        namespaces: Optional[List[str]] = None,
        requested_tier: str = "SIGNED",
        traceparent: Optional[str] = None,
        tracestate: Optional[str] = None,
        merge_latency_ms: float = 0.0,
        deterministic: bool = True
    ) -> MergeReceiptV2:
        """
        Write merge receipt with idempotency and proof generation.

        Args:
            parents: List of parent receipt dicts (must include receipt_id, digest, agent_id)
            strategy: Merge strategy config (type, weights, guardian_config)
            domain_data: Merged domain payload
            tenant_id: Tenant identifier
            namespaces: Namespace tags
            requested_tier: Requested tier (will be max'd with parent tiers)
            traceparent: W3C Trace Context traceparent
            tracestate: W3C Trace Context tracestate
            merge_latency_ms: Merge execution time
            deterministic: Whether merge is deterministically replayable

        Returns:
            MergeReceiptV2 with event_id, proof_refs, etc.
        """

        # Step 1: Load parent receipts and compute effective tier
        parent_receipts = [ParentReceipt(**p) for p in parents]
        parent_tiers = self._load_parent_tiers(parent_receipts)
        effective_tier = compute_effective_tier(parent_tiers, requested_tier)

        logger.info(
            f"Merge: parents={len(parent_receipts)}, "
            f"parent_tiers={parent_tiers}, "
            f"requested={requested_tier}, "
            f"effective={effective_tier}"
        )

        # Step 2: Compute digests
        digest_before = self._compute_digest_before(parent_receipts)
        digest_after = self._compute_digest_after(domain_data)

        # Step 2.5: Compute anti-gaming signals (Ωc composite)
        anti_gaming_data = self._compute_anti_gaming_signals(
            parents=parent_receipts,
            domain_data=domain_data,
            strategy=strategy
        )
        anti_gaming = AntiGamingSignals(**anti_gaming_data) if anti_gaming_data else None

        # Step 3: Build merge receipt (no event_id yet)
        now = datetime.now(timezone.utc)
        merge_receipt = MergeReceiptV2(
            schema="receipt.merge.v2",
            event_id="",  # Will be generated by Pydantic validator
            receipt_id="",  # Will be synced with event_id
            event_type="loom.coordination.merge.v2",
            parents=parent_receipts,
            strategy=MergeStrategy(**strategy),
            conflicts=None,  # TODO: Extract from CHOIR merge result
            digest_before=digest_before,
            digest_after=digest_after,
            valid_time=now,
            transaction_time=now,
            observed_at=None,  # Optional for merges
            tenant_id=tenant_id,
            namespaces=namespaces,
            tier=effective_tier,
            retention_days=compute_retention_days(effective_tier),
            signature=None,  # Will be populated by proof generation
            pq_signature=None,
            proof_refs=None,  # Will be populated by proof generation
            traceparent=traceparent,
            tracestate=tracestate,
            metadata=MergeMetadata(
                deterministic=deterministic,
                merge_latency_ms=merge_latency_ms,
                conflict_count=0,  # TODO: Count from CHOIR result
                guardian_interventions=0,  # TODO: Count from CHOIR result
                merge_tool_version="choir-v1.3.0"
            ),
            anti_gaming=anti_gaming,
            domain_data=domain_data
        )

        # Step 4: Write via IdempotentWriter (JCS canonical hash + de-dup)
        receipt_dict = merge_receipt.model_dump(by_alias=True, mode="json")
        written_receipt = self.writer.write_receipt(
            receipt_type="merge",
            content=receipt_dict,
            tenant_id=tenant_id
        )

        # Step 5: Generate tier-appropriate proofs
        proof_refs = self._generate_proofs(
            merge_receipt=merge_receipt,
            tier=effective_tier
        )

        # Step 6: Update receipt with proof_refs
        merge_receipt.proof_refs = proof_refs

        # Step 6.5: Compute trustmark badge based on cryptographic proofs
        receipt_dict_for_trustmark = merge_receipt.model_dump(by_alias=True, mode="json")
        trustmark_badge = compute_trustmark(
            receipt=receipt_dict_for_trustmark,
            expected_level=None  # No expectation - compute from actual proofs
        )
        logger.info(
            f"Trustmark computed: level={trustmark_badge.level}, "
            f"degraded={trustmark_badge.degraded}"
        )

        # Step 7: Persist to EventStore (SQLite WAL + Parquet export)
        # Add trustmark to receipt before persisting
        self._persist_to_event_store(merge_receipt, trustmark_badge)

        logger.info(
            f"Merge receipt written: event_id={merge_receipt.event_id}, "
            f"tier={effective_tier}, "
            f"proof_level={proof_refs.proof_level if proof_refs else 'none'}"
        )

        return merge_receipt

    def _load_parent_tiers(self, parents: List[ParentReceipt]) -> List[str]:
        """Load parent receipts and extract tiers."""
        tiers = []
        for parent in parents:
            # TODO: Load from EventStore by receipt_id
            # For now, assume SIGNED (will be replaced with actual lookup)
            tiers.append("SIGNED")
        return tiers

    def _compute_digest_before(self, parents: List[ParentReceipt]) -> Digest:
        """
        Compute digest_before from parent digests.

        Deterministic: sort parent digests lexicographically, concatenate, hash.
        """
        import hashlib

        sorted_digests = sorted(p.digest for p in parents)
        combined = "".join(sorted_digests).encode("utf-8")
        parents_hash = f"sha256:{hashlib.sha256(combined).hexdigest()}"

        return Digest(
            algorithm="sha256",
            parents_hash=parents_hash
        )

    def _compute_digest_after(self, domain_data: Dict[str, Any]) -> Digest:
        """
        Compute digest_after from merged domain_data using JCS canonicalization.

        RFC 8785: JSON Canonicalization Scheme for deterministic hashing.
        """
        import jcs
        import hashlib

        canonical_bytes = jcs.canonicalize(domain_data)
        merge_result_hash = f"sha256:{hashlib.sha256(canonical_bytes).hexdigest()}"

        return Digest(
            algorithm="sha256",
            merge_result_hash=merge_result_hash
        )

    def _generate_proofs(
        self,
        merge_receipt: MergeReceiptV2,
        tier: str
    ) -> Optional[ProofRefs]:
        """
        Generate tier-appropriate cryptographic proofs.

        Tier proof requirements:
        - TRANSIENT/OPERATIONAL: No proofs
        - RECORDED: JCS hash only
        - SIGNED: Ed25519 + batch proofs (Merkle tree + Rekor)
        - LEGAL: Ed25519 + ML-DSA + TSA + Rekor (individual, not batched)
        """

        if tier in ["TRANSIENT", "OPERATIONAL"]:
            return None

        # Compute JCS canonical hash
        import jcs
        import hashlib

        receipt_dict = merge_receipt.model_dump(by_alias=True, mode="json")
        canonical_bytes = jcs.canonicalize(receipt_dict)
        jcs_hash = f"sha256:{hashlib.sha256(canonical_bytes).hexdigest()}"

        proof_refs = ProofRefs(jcs_hash=jcs_hash)

        if tier == "RECORDED":
            proof_refs.proof_level = "hash"
            return proof_refs

        # SIGNED tier: Ed25519 + batch proofs
        if tier == "SIGNED":
            # TODO: Sign with Ed25519
            # TODO: Add to batch for Merkle tree
            # TODO: Submit batch root to Rekor
            proof_refs.proof_level = "signed"
            proof_refs.batch_id = "batch-01JCAB789BATCH"  # TODO: Actual batch ID
            proof_refs.batch_root = "sha256:batch_root_placeholder"
            proof_refs.rekor_uuid = "rekor_uuid_placeholder"
            proof_refs.rekor_log_index = 123456
            return proof_refs

        # LEGAL tier: Ed25519 + ML-DSA + TSA + Rekor (individual)
        if tier == "LEGAL":
            # TODO: Sign with Ed25519 + ML-DSA
            # TODO: Get RFC 3161 timestamp from TSA
            # TODO: Submit to Rekor individually
            proof_refs.proof_level = "transparent"
            proof_refs.tsa_ref = "rfc3161://tsa.freetsa.org/token/01JCAB789TSA"
            proof_refs.rekor_uuid = "rekor_uuid_legal_placeholder"
            proof_refs.rekor_log_index = 123457
            return proof_refs

        return proof_refs

    def _persist_to_event_store(self, merge_receipt: MergeReceiptV2, trustmark_badge=None):
        """
        Persist merge receipt to EventStore with trustmark badge.

        Hot tier (SQLite WAL): For recent events (last 24-48h)
        Cold tier (Parquet): For historical events (older than 48h)

        Partition strategy:
        - Parquet: dt=YYYY-MM-DD/tenant_id=.../part-*.parquet
        - Enables DuckDB predicate pushdown for fast queries
        """
        # TODO: Write to SQLite hot tier
        # TODO: Schedule Parquet export for cold tier

        # For now, write to JSON file as stub
        output_dir = self.event_store_path / "merges" / merge_receipt.tenant_id
        output_dir.mkdir(parents=True, exist_ok=True)

        # Add trustmark to receipt dict before persisting
        receipt_dict = merge_receipt.model_dump(by_alias=True, mode="json")
        if trustmark_badge:
            receipt_dict["trustmark"] = trustmark_badge.to_dict()

        output_file = output_dir / f"{merge_receipt.event_id}.json"
        with open(output_file, "w") as f:
            json.dump(
                receipt_dict,
                f,
                indent=2,
                default=str
            )

        logger.info(f"Merge receipt persisted: {output_file}")

    def _compute_anti_gaming_signals(
        self,
        parents: List[ParentReceipt],
        domain_data: Dict[str, Any],
        strategy: Dict[str, Any]
    ) -> Optional[Dict[str, Any]]:
        """
        Compute anti-gaming signals (Ωc composite) for merge.

        Uses production anti-gaming detection (receipts/anti_gaming.py) to detect:
        - Contradiction attacks (conflicting parent values)
        - Schism attacks (ensemble disagreement)
        - Burst attacks (rapid merge spikes)
        - Out-of-distribution inputs
        - Low-trust sources

        Args:
            parents: Parent receipts
            domain_data: Merged domain data
            strategy: Merge strategy

        Returns:
            Anti-gaming signals dict or None if not applicable
        """
        try:
            # Create temporary merge receipt for analysis
            # (we need the full receipt structure for V2 implementation)
            from receipts.models_merge import MergeMetadata

            now = datetime.now(timezone.utc)
            temp_receipt = MergeReceiptV2(
                event_id="",  # Temporary
                receipt_id="",  # Temporary
                parents=parents,
                strategy=MergeStrategy(**strategy),
                digest_before=Digest(algorithm="sha256"),
                digest_after=Digest(algorithm="sha256"),
                valid_time=now,
                transaction_time=now,
                tenant_id="temp",
                tier="SIGNED",
                metadata=MergeMetadata(
                    deterministic=True,
                    merge_latency_ms=0.0
                ),
                domain_data=domain_data
            )

            # Extract parent values for contradiction analysis
            parent_values = []
            if "prediction" in domain_data:
                # For prediction merges, extract parent predictions
                # (would come from loaded parent receipts in production)
                parent_values = [domain_data.get("prediction")] * len(parents)

            # TODO: Load recent merge timestamps from EventStore for burst detection
            recent_timestamps = []

            # TODO: Load known agent trust scores from historical accuracy data
            known_agents = {}

            # Compute anti-gaming signals using V2 implementation
            signals = compute_anti_gaming_signals_v2(
                merge_receipt=temp_receipt,
                recent_merge_timestamps=recent_timestamps,
                known_agents=known_agents,
                parent_values=parent_values
            )

            # Convert to dict for compatibility with existing code
            return signals.model_dump()

        except Exception as e:
            logger.warning(f"Failed to compute anti-gaming signals: {e}", exc_info=True)
            return None

    def verify_merge(
        self,
        merge_id: str,
        offline: bool = True,
        strict: bool = False
    ) -> Dict[str, Any]:
        """
        Verify cryptographic proofs for merge receipt.

        Checks:
        1. JCS canonical hash
        2. Ed25519 signature
        3. ML-DSA signature (if present)
        4. Parent digest verification
        5. Conflict integrity
        6. RFC 3161 timestamp (LEGAL tier)
        7. Rekor inclusion proof
        8. Batch Merkle path (SIGNED tier)

        Args:
            merge_id: Event ID or receipt ID
            offline: Use cached proofs only (air-gapped mode)
            strict: Require PQC signature + Rekor

        Returns:
            Verification result with verdict and checks
        """
        # TODO: Load merge receipt from EventStore
        # TODO: Delegate to verify.py verification logic
        # TODO: Add merge-specific checks (parent digests, conflicts)

        return {
            "verdict": "PASS",
            "merge_id": merge_id,
            "checks": [
                {"name": "jcs_canonical_hash", "status": "OK"},
                {"name": "ed25519_signature", "status": "OK"},
                {"name": "parent_digests", "status": "OK"},
            ],
            "verified_at": datetime.now(timezone.utc).isoformat()
        }

    def replay_merge(
        self,
        merge_id: str,
        use_cached_parents: bool = True
    ) -> Dict[str, Any]:
        """
        Replay merge and compare digest_after for determinism verification.

        Args:
            merge_id: Event ID or receipt ID
            use_cached_parents: Use cached parent receipts (faster)

        Returns:
            Replay result with match verdict and diff (if mismatch)
        """
        # TODO: Load merge receipt
        # TODO: Load parent receipts
        # TODO: Re-execute merge with same strategy
        # TODO: Compute digest_after and compare
        # TODO: Return match/mismatch with diff

        return {
            "verdict": "MATCH",
            "merge_id": merge_id,
            "original_digest": "sha256:abc123...",
            "replay_digest": "sha256:abc123...",
            "deterministic": True
        }
