"""
Receipt Chain helpers for staging DIP → Prophecy/Refusal commits behind a feature flag.

The collector accumulates receipts emitted during a pipeline run and writes them
in order once the run succeeds, emitting a ReceiptChainReceipt manifest that
captures the lineage.
"""

from __future__ import annotations

import hashlib
import logging
import os
import time
import uuid
from datetime import datetime, timedelta, timezone
from functools import lru_cache
from pathlib import Path
from typing import Callable, List, Optional

from config.paths import DATA_ROOT
from receipts.compat.pyd import BaseModel, ConfigDictLike, Field
from receipts.base import BaseReceipt
from receipts.index import lookup_receipt
from receipts.models import DIPReceipt
from tools.lib import metrics as observability_metrics

logger = logging.getLogger(__name__)

__all__ = ["ReceiptChain", "ReceiptChainReceipt", "ReceiptChainFile"]


class ReceiptChainFile(BaseModel):
    """Manifest entry describing a persisted receipt."""

    model_config = ConfigDictLike(extra="forbid", frozen=True)

    path: str = Field(description="Relative path to the persisted receipt JSON payload")
    receipt_id: str = Field(description="Receipt identifier")
    receipt_type: str = Field(description="Concrete receipt type (ProphecyStone, RefusalStone, ...)")
    receipt_hash: str = Field(description="JCS SHA-256 hash for the receipt payload")


class ReceiptChainReceipt(BaseReceipt):
    """Manifest receipt describing an atomic chain commit."""

    receipt_type: str = "ReceiptChainReceipt"

    chain_id: str = Field(description="Unique identifier for the chain commit")
    dip_receipt_id: str = Field(description="Root DIP receipt id")
    prophecy_receipt_id: Optional[str] = Field(default=None, description="Prophecy receipt id if emitted")
    refusal_receipt_id: Optional[str] = Field(default=None, description="Refusal receipt id if emitted")
    receipts: List[ReceiptChainFile] = Field(default_factory=list, description="Persisted receipts in commit order")


class ReceiptChain:
    """
    Collects receipts during pipeline execution and persists them after the run finishes.

    The collector ensures receipts are written in DIP → Prophecy → Refusal order
    and emits a ReceiptChainReceipt manifest that captures the lineage.
    """

    _TRUTHY = {"1", "true", "t", "yes", "y", "on"}

    def __init__(
        self,
        domain: str,
        *,
        tenant: Optional[str] = None,
        root: Optional[Path] = None,
        writer: Optional[Callable[[BaseReceipt], None]] = None,
    ) -> None:
        self.domain = domain
        self.tenant = tenant
        self.root = Path(root or DATA_ROOT)
        if writer is None:
            from receipts.writer import write_receipt

            self._writer = write_receipt
        else:
            self._writer = writer
        self.chain_id = f"{domain}-{uuid.uuid4().hex}"

        self._dip: Optional[DIPReceipt] = None
        self._dip_existing = False
        self._prophecy: Optional[BaseReceipt] = None
        self._prophecy_flushed = False  # P0 fix: track flush state
        self._refusal: Optional[BaseReceipt] = None
        self._refusal_flushed = False  # P0 fix: track flush state
        self._committed = False
        self._manifest_path: Optional[Path] = None
        
        # Pre-generate manifest_receipt_id for guardian context injection (P0 fix)
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S%f")
        self.manifest_receipt_id = f"CHAIN_{self.chain_id}_{timestamp}"

        # Clean up stale staging artifacts on init
        self._gc_staging_artifacts()

    @staticmethod
    @lru_cache(maxsize=1)
    def feature_enabled() -> bool:
        """Check if receipt chain is enabled via env flag (cached)."""
        v = (os.getenv("LOOM_RECEIPT_CHAIN_ENABLED", "0") or "").strip().lower()
        return v in ReceiptChain._TRUTHY

    def stage_dip(self, dip: DIPReceipt, *, existing: bool = False) -> None:
        """Record the DIP that anchors the chain."""
        self._dip = dip
        self._dip_existing = existing
        self._prepare_receipt(dip, allow_orphan=True)
        related = dict(getattr(dip, "related_ids", {}) or {})
        related.setdefault("receipt_chain_id", self.chain_id)
        self._assign(dip, "related_ids", related)

    def stage(self, receipt: BaseReceipt) -> None:
        """Capture prophecy/refusal receipts; fall back to the real writer otherwise."""
        receipt_type = getattr(receipt, "receipt_type", receipt.__class__.__name__)
        if receipt_type == "ProphecyStone":
            self._prophecy = receipt
            self._tag_with_chain(receipt)
            self._prepare_receipt(receipt)
            return
        if receipt_type == "RefusalStone":
            self._refusal = receipt
            self._tag_with_chain(receipt)
            self._prepare_receipt(receipt)
            return
        self._writer(receipt)

    def commit(self) -> Optional[ReceiptChainReceipt]:
        """
        Persist staged receipts and emit the manifest (idempotent).

        If a manifest already exists and matches the staged content,
        this is a no-op. If it exists but differs, an error is logged
        and the existing manifest is returned.
        """
        commit_start = time.time()
        stage = "initialization"
        reason = "unknown"

        try:
            if self._committed:
                raise RuntimeError("ReceiptChain.commit() called more than once")
            self._committed = True

            if self._dip is None:
                # Nothing to persist – fall back to normal flow.
                return None

            if self._prophecy is None and self._refusal is None:
                # Only a DIP was staged; flush it via the regular writer.
                stage = "dip_only_persist"
                if not self._dip_existing:
                    self._writer(self._dip)
                    self._dip_existing = True
                return None

            stage = "dip_persist"
            entries: List[ReceiptChainFile] = []
            dip_entry = self._persist_and_entry(self._dip, existing=self._dip_existing)
            if not self._dip_existing:
                self._dip_existing = True  # P0 fix: mark as flushed after persist
            entries.append(dip_entry)

            stage = "prophecy_persist"
            prophecy_id: Optional[str] = None
            if self._prophecy is not None:
                entries.append(self._persist_and_entry(self._prophecy, existing=self._prophecy_flushed))  # P0 fix: check flushed state
                self._prophecy_flushed = True  # P0 fix: mark as flushed after persist
                prophecy_id = self._prophecy.receipt_id

            stage = "refusal_persist"
            refusal_id: Optional[str] = None
            if self._refusal is not None:
                entries.append(self._persist_and_entry(self._refusal, existing=self._refusal_flushed))  # P0 fix: check flushed state
                self._refusal_flushed = True  # P0 fix: mark as flushed after persist
                refusal_id = self._refusal.receipt_id

            stage = "manifest_build"
            manifest_related = {
                "dip_receipt_id": self._dip.receipt_id,
                "receipt_chain_id": self.chain_id,
            }
            if prophecy_id:
                manifest_related["prophecy_receipt_id"] = prophecy_id
            if refusal_id:
                manifest_related["refusal_receipt_id"] = refusal_id

            manifest = ReceiptChainReceipt(
                receipt_id=self.manifest_receipt_id,  # Use pre-generated ID (P0 fix)
                domain=self.domain,
                mode=getattr(self._dip, "mode", "LIVE"),
                tenant=self.tenant or getattr(self._dip, "tenant", None),
                chain_id=self.chain_id,
                dip_receipt_id=self._dip.receipt_id,
                prophecy_receipt_id=prophecy_id,
                refusal_receipt_id=refusal_id,
                receipts=entries,
                allow_orphan=True,
                related_ids=manifest_related,
            )

            stage = "manifest_idempotency_check"
            # Idempotent commit: check if manifest already exists
            existing_manifest = self._check_existing_manifest(manifest)
            if existing_manifest is not None:
                return existing_manifest

            stage = "manifest_write"
            self._writer(manifest)
            return manifest

        except Exception as exc:
            # P0 #3: Emit commit failure telemetry
            duration = time.time() - commit_start
            reason = exc.__class__.__name__
            
            # Extract more specific reason from exception message
            exc_msg = str(exc).lower()
            if "dip" in exc_msg:
                reason = "persist_dip_failed"
            elif "prophecy" in exc_msg:
                reason = "persist_prophecy_failed"
            elif "refusal" in exc_msg:
                reason = "persist_refusal_failed"
            elif "manifest" in exc_msg:
                reason = "manifest_write_failed"
            elif "idempotency" in exc_msg:
                reason = "idempotency_check_failed"
            
            logger.error(
                "ReceiptChain.commit() failed",
                extra={
                    "chain_id": self.chain_id,
                    "dip_receipt_id": self._dip.receipt_id if self._dip else None,
                    "domain": self.domain,
                    "stage": stage,
                    "reason": reason,
                    "duration_seconds": duration,
                    "exception": str(exc),
                },
            )
            
            observability_metrics.record_receipt_chain_commit_failure(
                domain=self.domain,
                reason=reason,
                stage=stage,
                duration_seconds=duration,
            )
            
            raise

    def flush(self) -> None:
        """Persist staged receipts without emitting a manifest (used on error paths)."""
        # P0 FIX: Idempotent flush with tracking flags
        if self._dip and not self._dip_existing:
            self._writer(self._dip)
            self._dip_existing = True
        if self._prophecy and not self._prophecy_flushed:
            self._writer(self._prophecy)
            self._prophecy_flushed = True
        if self._refusal and not self._refusal_flushed:
            self._writer(self._refusal)
            self._refusal_flushed = True

    # ------------------------------------------------------------------ helpers

    def _prepare_receipt(self, receipt: BaseReceipt, *, allow_orphan: bool = False) -> None:
        now = datetime.now(timezone.utc)
        if getattr(receipt, "system_time", None) is None:
            self._assign(receipt, "system_time", now)
        if allow_orphan and not getattr(receipt, "allow_orphan", False):
            self._assign(receipt, "allow_orphan", True)
        try:
            receipt_hash = receipt.compute_hash()
        except Exception:
            receipt_hash = None
        if receipt_hash:
            self._assign(receipt, "receipt_hash", receipt_hash)

    def _persist_and_entry(self, receipt: BaseReceipt, *, existing: bool = False) -> ReceiptChainFile:
        if not existing:
            self._writer(receipt)
        entry = lookup_receipt(receipt.receipt_id, root=self.root)
        if entry is None:
            raise FileNotFoundError(f"receipt '{receipt.receipt_id}' not indexed")
        try:
            relative = entry.path.relative_to(self.root)
            path_value = str(relative)
        except ValueError:
            path_value = str(entry.path)
        receipt_hash = getattr(receipt, "receipt_hash", None) or entry.jcs_sha256
        return ReceiptChainFile(
            path=path_value,
            receipt_id=receipt.receipt_id,
            receipt_type=getattr(receipt, "receipt_type", receipt.__class__.__name__),
            receipt_hash=str(receipt_hash),
        )

    def _tag_with_chain(self, receipt: BaseReceipt) -> None:
        related = dict(getattr(receipt, "related_ids", {}) or {})
        related.setdefault("receipt_chain_id", self.chain_id)
        if self._dip is not None:
            related.setdefault("dip_receipt_id", self._dip.receipt_id)
        self._assign(receipt, "related_ids", related)

    def _make_receipt_id(self) -> str:
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S%f")
        return f"CHAIN_{self.chain_id}_{timestamp}"

    @staticmethod
    def _assign(receipt: BaseReceipt, field: str, value) -> None:  # type: ignore[no-untyped-def]
        object.__setattr__(receipt, field, value)

    def _gc_staging_artifacts(self) -> None:
        """
        Remove stale staging artifacts older than 24 hours.

        This prevents accumulation of orphaned staging files from
        crashed processes or incomplete chains.
        """
        try:
            chain_temp_dir = self.root / "receipts" / "staging"
            if not chain_temp_dir.exists():
                return

            cutoff_time = time.time() - (24 * 3600)  # 24 hours ago
            removed_count = 0

            for staging_file in chain_temp_dir.rglob("*.staging"):
                try:
                    if staging_file.stat().st_mtime < cutoff_time:
                        staging_file.unlink()
                        removed_count += 1
                except OSError:
                    # Ignore errors on individual files (permissions, etc.)
                    continue

            if removed_count > 0:
                logger.info(
                    "Cleaned up stale staging artifacts",
                    extra={"count": removed_count, "domain": self.domain},
                )
        except Exception as exc:
            # Non-critical; log and continue
            logger.debug("Failed to GC staging artifacts: %s", exc)

    def _check_existing_manifest(
        self, manifest: ReceiptChainReceipt
    ) -> Optional[ReceiptChainReceipt]:
        """
        Check if an identical manifest already exists (idempotent commit).

        Returns:
            Existing manifest if found and hash matches, None otherwise.
            Logs an error if manifest exists but differs.
        """
        try:
            entry = lookup_receipt(manifest.receipt_id, root=self.root)
            if entry is None:
                return None

            # Manifest exists; compute hash to check if it's identical
            existing_hash = self._compute_manifest_hash(entry.path)
            new_hash = self._compute_manifest_hash_from_model(manifest)

            if existing_hash == new_hash:
                logger.info(
                    "Manifest already exists with matching hash (idempotent no-op)",
                    extra={
                        "receipt_id": manifest.receipt_id,
                        "chain_id": self.chain_id,
                    },
                )
                # Return the existing manifest (re-construct from file)
                import json
                from receipts.models import load_receipt

                payload = json.loads(entry.path.read_text(encoding="utf-8"))
                return load_receipt(payload)  # type: ignore[return-value]

            # Hash mismatch – this shouldn't happen in normal operation
            logger.error(
                "Manifest exists but hash differs (potential double-commit)",
                extra={
                    "receipt_id": manifest.receipt_id,
                    "chain_id": self.chain_id,
                    "existing_hash": existing_hash,
                    "new_hash": new_hash,
                },
            )
            return None  # Allow the new commit to proceed (caller will overwrite)

        except Exception as exc:
            logger.debug("Failed to check existing manifest: %s", exc)
            return None

    def _compute_manifest_hash(self, path: Path) -> str:
        """Compute SHA-256 hash of a manifest JSON file."""
        import json

        payload = json.loads(path.read_text(encoding="utf-8"))
        # Hash the items list (canonical representation)
        items = payload.get("receipts", [])
        canonical = json.dumps(items, separators=(",", ":"), sort_keys=True)
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()

    def _compute_manifest_hash_from_model(self, manifest: ReceiptChainReceipt) -> str:
        """Compute SHA-256 hash of a manifest model."""
        import json

        items = [entry.model_dump(mode="python") for entry in manifest.receipts]
        canonical = json.dumps(items, separators=(",", ":"), sort_keys=True)
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
