"""
Receipt Wind Tunnel — Replay historical receipts through new policies.

For any change (Guardian rule, healer strategy, learning query, Prophecy model),
you don't just test forward—you replay months of receipts through it offline and see:

- How many decisions would have changed?
- Which invariants would have been violated or newly satisfied?
- How calibration/forecast scores compare to previous version?
- How many repairs/escalations it would have triggered or avoided?

Usage:
    harness = ReceiptReplayHarness(
        source=FileReceiptSource(".loom/receipts"),
        policies={"v1": OldGuardianPolicy(), "v2": NewGuardianPolicy()},
        invariants=[inv1, inv2, inv3],
    )
    report = harness.run(start_date, end_date)
    print(report.summary())
"""

import json
import os
from abc import ABC, abstractmethod
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Callable, Dict, Generic, Iterator, List, Optional, Protocol, TypeVar


# === Receipt Type ===

@dataclass
class Receipt:
    """Generic receipt for replay purposes."""
    id: str
    type: str  # decision, healing, escalation, prophecy, etc.
    timestamp: datetime
    data: Dict[str, Any]

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "Receipt":
        ts = d.get("timestamp") or d.get("time")
        if isinstance(ts, str):
            ts = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        elif ts is None:
            ts = datetime.now(timezone.utc)

        # Ensure timezone-aware
        if ts.tzinfo is None:
            ts = ts.replace(tzinfo=timezone.utc)

        return cls(
            id=d.get("id", ""),
            type=d.get("type", d.get("receipt_type", "unknown")),
            timestamp=ts,
            data=d,
        )


# === Receipt Source Protocol ===

class ReceiptSource(Protocol):
    """Protocol for receipt sources."""

    def stream(
        self,
        start: Optional[datetime] = None,
        end: Optional[datetime] = None,
        receipt_types: Optional[List[str]] = None,
    ) -> Iterator[Receipt]:
        """Stream receipts in chronological order."""
        ...


# === File-based Receipt Source ===

class FileReceiptSource:
    """Load receipts from .loom/receipts directory structure."""

    def __init__(self, base_path: str = ".loom/receipts"):
        self.base_path = Path(base_path)

    def stream(
        self,
        start: Optional[datetime] = None,
        end: Optional[datetime] = None,
        receipt_types: Optional[List[str]] = None,
    ) -> Iterator[Receipt]:
        """Stream receipts from files, sorted by timestamp."""
        receipts = []

        if not self.base_path.exists():
            return iter([])

        # Walk all JSON files
        for json_file in self.base_path.rglob("*.json"):
            try:
                with open(json_file) as f:
                    data = json.load(f)

                receipt = Receipt.from_dict(data)

                # Filter by date range
                if start and receipt.timestamp < start:
                    continue
                if end and receipt.timestamp > end:
                    continue

                # Filter by type
                if receipt_types and receipt.type not in receipt_types:
                    continue

                receipts.append(receipt)
            except (json.JSONDecodeError, KeyError, ValueError):
                continue  # Skip malformed receipts

        # Sort by timestamp
        receipts.sort(key=lambda r: r.timestamp)
        return iter(receipts)


# === Policy Protocol ===

T = TypeVar("T")


class Policy(Protocol[T]):
    """Protocol for policies that can be applied to receipts."""

    def evaluate(self, receipt: Receipt, context: Dict[str, Any]) -> T:
        """Evaluate the policy on a receipt. Returns policy's decision."""
        ...


# === Invariant Protocol ===

class Invariant(Protocol):
    """Protocol for invariants to check during replay."""

    @property
    def name(self) -> str:
        """Name of the invariant."""
        ...

    def check(self, receipt: Receipt, context: Dict[str, Any]) -> bool:
        """Check if invariant holds. Returns True if satisfied."""
        ...


# === Replay Results ===

@dataclass
class PolicyResult:
    """Result of applying a policy to a receipt."""
    receipt_id: str
    decision: Any
    timestamp: datetime


@dataclass
class InvariantViolation:
    """Record of an invariant violation."""
    invariant_name: str
    receipt_id: str
    timestamp: datetime
    context: Dict[str, Any] = field(default_factory=dict)


@dataclass
class PolicyDiff:
    """Difference between two policies on the same receipt."""
    receipt_id: str
    timestamp: datetime
    policy_a_decision: Any
    policy_b_decision: Any
    receipt_type: str


@dataclass
class ReplayMetrics:
    """Metrics collected during replay."""
    total_receipts: int = 0
    receipts_by_type: Dict[str, int] = field(default_factory=lambda: defaultdict(int))

    # Decision metrics
    decisions_changed: int = 0
    decisions_same: int = 0

    # Invariant metrics
    invariant_violations: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
    invariant_satisfied: Dict[str, int] = field(default_factory=lambda: defaultdict(int))

    # Healer-specific metrics (if applicable)
    repairs_triggered: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
    escalations_triggered: Dict[str, int] = field(default_factory=lambda: defaultdict(int))

    # Calibration (for prophecy/forecast policies)
    predictions: int = 0
    correct_predictions: int = 0

    @property
    def decision_change_rate(self) -> float:
        total = self.decisions_changed + self.decisions_same
        return self.decisions_changed / total if total > 0 else 0.0

    @property
    def calibration_accuracy(self) -> float:
        return self.correct_predictions / self.predictions if self.predictions > 0 else 0.0


@dataclass
class ReplayReport:
    """Complete report from a replay run."""
    start_time: datetime
    end_time: datetime
    replay_window_start: Optional[datetime]
    replay_window_end: Optional[datetime]

    # Policies compared
    policies: List[str]

    # Metrics per policy
    metrics: Dict[str, ReplayMetrics]

    # All diffs between policies (A vs B)
    diffs: List[PolicyDiff] = field(default_factory=list)

    # All invariant violations
    violations: List[InvariantViolation] = field(default_factory=list)

    def summary(self) -> str:
        """Generate human-readable summary."""
        lines = [
            "=" * 60,
            "RECEIPT REPLAY REPORT",
            "=" * 60,
            f"Replay window: {self.replay_window_start} to {self.replay_window_end}",
            f"Policies compared: {', '.join(self.policies)}",
            "",
        ]

        for policy_name, m in self.metrics.items():
            lines.append(f"--- {policy_name} ---")
            lines.append(f"  Total receipts: {m.total_receipts}")
            lines.append(f"  Receipts by type: {dict(m.receipts_by_type)}")
            lines.append(f"  Repairs triggered: {dict(m.repairs_triggered)}")
            lines.append(f"  Escalations triggered: {dict(m.escalations_triggered)}")
            if m.predictions > 0:
                lines.append(f"  Calibration: {m.calibration_accuracy:.1%} ({m.correct_predictions}/{m.predictions})")
            lines.append("")

        if len(self.policies) >= 2:
            p1, p2 = self.policies[0], self.policies[1]
            m1, m2 = self.metrics.get(p1), self.metrics.get(p2)
            if m1 and m2:
                lines.append("--- A/B COMPARISON ---")
                lines.append(f"  Decisions changed: {len(self.diffs)}")
                if m1.decision_change_rate > 0:
                    lines.append(f"  Change rate: {m1.decision_change_rate:.1%}")

                # Compare repairs/escalations
                total_repairs_1 = sum(m1.repairs_triggered.values())
                total_repairs_2 = sum(m2.repairs_triggered.values())
                total_esc_1 = sum(m1.escalations_triggered.values())
                total_esc_2 = sum(m2.escalations_triggered.values())

                if total_repairs_1 != total_repairs_2:
                    lines.append(f"  Repairs: {p1}={total_repairs_1}, {p2}={total_repairs_2} (Δ={total_repairs_2 - total_repairs_1})")
                if total_esc_1 != total_esc_2:
                    lines.append(f"  Escalations: {p1}={total_esc_1}, {p2}={total_esc_2} (Δ={total_esc_2 - total_esc_1})")
                lines.append("")

        if self.violations:
            lines.append("--- INVARIANT VIOLATIONS ---")
            violation_counts = defaultdict(int)
            for v in self.violations:
                violation_counts[v.invariant_name] += 1
            for inv_name, count in violation_counts.items():
                lines.append(f"  {inv_name}: {count} violations")
            lines.append("")

        lines.append("=" * 60)
        return "\n".join(lines)


# === Replay Harness ===

class ReceiptReplayHarness:
    """
    The wind tunnel: replay receipts through policies and measure impact.

    Usage:
        harness = ReceiptReplayHarness(
            source=FileReceiptSource(".loom/receipts"),
            policies={"v1": old_policy, "v2": new_policy},
            invariants=[inv1, inv2],
        )
        report = harness.run(start_date, end_date)
    """

    def __init__(
        self,
        source: ReceiptSource,
        policies: Optional[Dict[str, Policy]] = None,
        invariants: Optional[List[Invariant]] = None,
    ):
        self.source = source
        self.policies = policies or {}
        self.invariants = invariants or []

    def add_policy(self, name: str, policy: Policy) -> None:
        """Add a policy to compare."""
        self.policies[name] = policy

    def add_invariant(self, invariant: Invariant) -> None:
        """Add an invariant to check."""
        self.invariants.append(invariant)

    def run(
        self,
        start: Optional[datetime] = None,
        end: Optional[datetime] = None,
        receipt_types: Optional[List[str]] = None,
    ) -> ReplayReport:
        """Run the replay and return a report."""
        run_start = datetime.now(timezone.utc)

        # Initialize metrics for each policy
        metrics: Dict[str, ReplayMetrics] = {
            name: ReplayMetrics() for name in self.policies
        }

        diffs: List[PolicyDiff] = []
        violations: List[InvariantViolation] = []
        context: Dict[str, Any] = {}  # Shared context across receipts

        # Stream and process receipts
        for receipt in self.source.stream(start, end, receipt_types):
            # Update metrics
            for name in self.policies:
                metrics[name].total_receipts += 1
                metrics[name].receipts_by_type[receipt.type] += 1

            # Apply each policy
            results: Dict[str, Any] = {}
            for name, policy in self.policies.items():
                try:
                    results[name] = policy.evaluate(receipt, context)
                except Exception as e:
                    results[name] = f"ERROR: {e}"

            # Compare policies (A/B)
            policy_names = list(self.policies.keys())
            if len(policy_names) >= 2:
                a, b = policy_names[0], policy_names[1]
                if results.get(a) != results.get(b):
                    diffs.append(PolicyDiff(
                        receipt_id=receipt.id,
                        timestamp=receipt.timestamp,
                        policy_a_decision=results.get(a),
                        policy_b_decision=results.get(b),
                        receipt_type=receipt.type,
                    ))
                    for name in [a, b]:
                        metrics[name].decisions_changed += 1
                else:
                    for name in [a, b]:
                        metrics[name].decisions_same += 1

            # Check invariants
            for invariant in self.invariants:
                try:
                    if invariant.check(receipt, context):
                        for name in self.policies:
                            metrics[name].invariant_satisfied[invariant.name] += 1
                    else:
                        violations.append(InvariantViolation(
                            invariant_name=invariant.name,
                            receipt_id=receipt.id,
                            timestamp=receipt.timestamp,
                            context=dict(context),
                        ))
                        for name in self.policies:
                            metrics[name].invariant_violations[invariant.name] += 1
                except Exception:
                    pass  # Invariant check failed

            # Track healer-specific metrics
            if receipt.type in ("repair", "healing", "repair_request"):
                for name in self.policies:
                    action = receipt.data.get("repair_action", "unknown")
                    metrics[name].repairs_triggered[action] += 1
            elif receipt.type in ("escalation", "escalation_notice"):
                for name in self.policies:
                    reason = receipt.data.get("reason", "unknown")
                    metrics[name].escalations_triggered[reason] += 1

            # Track calibration for prophecy
            if receipt.type in ("prophecy", "prediction", "forecast"):
                for name, result in results.items():
                    if isinstance(result, dict) and "predicted" in result and "actual" in result:
                        metrics[name].predictions += 1
                        if result["predicted"] == result["actual"]:
                            metrics[name].correct_predictions += 1

            # Update context for next receipt
            context["last_receipt"] = receipt
            context["receipt_count"] = context.get("receipt_count", 0) + 1

        run_end = datetime.now(timezone.utc)

        return ReplayReport(
            start_time=run_start,
            end_time=run_end,
            replay_window_start=start,
            replay_window_end=end,
            policies=list(self.policies.keys()),
            metrics=metrics,
            diffs=diffs,
            violations=violations,
        )


# === Example Policies ===

class PassThroughPolicy:
    """Policy that just returns the receipt data (baseline)."""

    def evaluate(self, receipt: Receipt, context: Dict[str, Any]) -> Dict[str, Any]:
        return receipt.data


class HealerCorePolicy:
    """Policy that runs receipts through HealerCore."""

    def __init__(self, core=None):
        from organism.autonomous_healer.core import HealerCore
        self.core = core or HealerCore()

    def evaluate(self, receipt: Receipt, context: Dict[str, Any]) -> Optional[str]:
        # Feed detection receipts to core
        if receipt.type in ("detection", "anomaly"):
            from organism.autonomous_healer.invariants import DetectionEvidence, IndependenceGroup
            evidence = DetectionEvidence(
                source=receipt.data.get("source", "unknown"),
                group=IndependenceGroup(receipt.data.get("group", "external")),
                anomaly_type=receipt.data.get("anomaly_type", "unknown"),
                timestamp=receipt.timestamp,
                confidence=receipt.data.get("confidence", 0.8),
            )
            self.core.observe(evidence)

        # Tick and return decision
        result = self.core.tick()
        if result:
            return type(result).__name__
        return None


# === Helper: Replay Last N Days ===

def replay_last_n_days(
    days: int,
    policies: Dict[str, Policy],
    invariants: Optional[List[Invariant]] = None,
    receipts_path: str = ".loom/receipts",
) -> ReplayReport:
    """Convenience function to replay the last N days."""
    end = datetime.now(timezone.utc)
    start = end - timedelta(days=days)

    harness = ReceiptReplayHarness(
        source=FileReceiptSource(receipts_path),
        policies=policies,
        invariants=invariants or [],
    )
    return harness.run(start=start, end=end)


# === ReplayReportReceipt for downstream tooling ===

@dataclass
class ReplayReportReceipt:
    """
    Receipt emitted after a replay run for downstream consumption.

    This allows other systems (dashboards, alerting, learning) to react to
    replay results.
    """
    id: str
    timestamp: datetime
    replay_window_start: Optional[datetime]
    replay_window_end: Optional[datetime]
    policies_compared: List[str]
    total_receipts_replayed: int
    total_diffs: int
    total_violations: int
    decision_change_rate: float
    metrics_summary: Dict[str, Dict[str, Any]]
    passed: bool
    failure_reasons: List[str]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to JSON-serializable dict."""
        return {
            "id": self.id,
            "type": "replay_report",
            "timestamp": self.timestamp.isoformat(),
            "replay_window_start": self.replay_window_start.isoformat() if self.replay_window_start else None,
            "replay_window_end": self.replay_window_end.isoformat() if self.replay_window_end else None,
            "policies_compared": self.policies_compared,
            "total_receipts_replayed": self.total_receipts_replayed,
            "total_diffs": self.total_diffs,
            "total_violations": self.total_violations,
            "decision_change_rate": self.decision_change_rate,
            "metrics_summary": self.metrics_summary,
            "passed": self.passed,
            "failure_reasons": self.failure_reasons,
        }

    @classmethod
    def from_report(
        cls,
        report: ReplayReport,
        passed: bool = True,
        failure_reasons: Optional[List[str]] = None,
    ) -> "ReplayReportReceipt":
        """Create a ReplayReportReceipt from a ReplayReport."""
        import uuid

        # Calculate total receipts (from first policy's metrics)
        first_policy = report.policies[0] if report.policies else None
        total_receipts = report.metrics[first_policy].total_receipts if first_policy else 0

        # Calculate decision change rate
        change_rate = 0.0
        if first_policy and report.metrics[first_policy].decisions_changed + report.metrics[first_policy].decisions_same > 0:
            m = report.metrics[first_policy]
            change_rate = m.decisions_changed / (m.decisions_changed + m.decisions_same)

        # Build metrics summary
        metrics_summary = {}
        for name, m in report.metrics.items():
            metrics_summary[name] = {
                "total_receipts": m.total_receipts,
                "receipts_by_type": dict(m.receipts_by_type),
                "decision_change_rate": m.decision_change_rate,
                "repairs_triggered": sum(m.repairs_triggered.values()),
                "escalations_triggered": sum(m.escalations_triggered.values()),
                "invariant_violations": sum(m.invariant_violations.values()),
            }

        return cls(
            id=f"replay_{uuid.uuid4().hex[:12]}",
            timestamp=report.end_time,
            replay_window_start=report.replay_window_start,
            replay_window_end=report.replay_window_end,
            policies_compared=report.policies,
            total_receipts_replayed=total_receipts,
            total_diffs=len(report.diffs),
            total_violations=len(report.violations),
            decision_change_rate=change_rate,
            metrics_summary=metrics_summary,
            passed=passed,
            failure_reasons=failure_reasons or [],
        )


def emit_report_receipt(
    report: ReplayReport,
    output_path: str = ".loom/receipts/replay",
    passed: bool = True,
    failure_reasons: Optional[List[str]] = None,
) -> str:
    """
    Emit a ReplayReportReceipt to the receipts directory.

    Returns the path to the written receipt.
    """
    receipt = ReplayReportReceipt.from_report(report, passed, failure_reasons)

    # Ensure output directory exists
    out_dir = Path(output_path)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Write receipt
    out_file = out_dir / f"{receipt.id}.json"
    out_file.write_text(json.dumps(receipt.to_dict(), indent=2))

    return str(out_file)
