"""
EntityReducer — Per-entity state tracking across multiple DAG runs.

While the Reducer tracks state within a single run, the EntityReducer
tracks state across runs for a specific entity (worker, portfolio,
venue, merchant, etc.).

Uses SnapChore-hashed checkpoints so state can be verified and
replayed if discrepancies are detected.

Examples:
  Dominion:   EntityReducer("employee", "worker-42") — tracks all payruns for a worker
  Stillpoint: EntityReducer("portfolio", "port-001") — tracks all executions for a portfolio
  Sonic:      EntityReducer("merchant", "merch-99")  — tracks all settlements for a merchant
  Grid:       EntityReducer("participant", "user-7")  — tracks all slot completions for a user
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from .artifact import ProofArtifact

# ---------------------------------------------------------------------------
# Entity state
# ---------------------------------------------------------------------------

@dataclass
class EntityCheckpoint:
    """A hashed checkpoint of entity state at a point in time.

    Checkpoints are taken every N runs (or at period close) and
    published as SmartBlocks.  Full replay only needed from last
    valid checkpoint.
    """
    entity_type: str
    entity_id: str
    run_id: str
    sequence: int
    state_hash: str
    total_runs: int
    total_proofs: int
    total_amount: float
    status_counts: Dict[str, int] = field(default_factory=dict)
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def to_dict(self) -> Dict[str, Any]:
        return {
            "entity_type": self.entity_type,
            "entity_id": self.entity_id,
            "run_id": self.run_id,
            "sequence": self.sequence,
            "state_hash": self.state_hash,
            "total_runs": self.total_runs,
            "total_proofs": self.total_proofs,
            "total_amount": self.total_amount,
            "status_counts": self.status_counts,
            "created_at": self.created_at.isoformat(),
        }


@dataclass
class EntityState:
    """Accumulated state for a single entity across runs."""
    entity_type: str
    entity_id: str
    total_runs: int = 0
    total_proofs: int = 0
    total_completed: int = 0
    total_failed: int = 0
    total_amount: float = 0.0
    run_ids: List[str] = field(default_factory=list)
    latest_run_id: Optional[str] = None
    latest_checkpoint: Optional[EntityCheckpoint] = None
    all_proofs: Dict[str, List[ProofArtifact]] = field(default_factory=dict)
    state_hash: str = ""

    def to_dict(self) -> Dict[str, Any]:
        return {
            "entity_type": self.entity_type,
            "entity_id": self.entity_id,
            "total_runs": self.total_runs,
            "total_proofs": self.total_proofs,
            "total_completed": self.total_completed,
            "total_failed": self.total_failed,
            "total_amount": self.total_amount,
            "run_ids": self.run_ids,
            "latest_run_id": self.latest_run_id,
            "state_hash": self.state_hash,
        }


# ---------------------------------------------------------------------------
# EntityReducer
# ---------------------------------------------------------------------------

class EntityReducer:
    """Tracks entity state across multiple DAG runs.

    Usage:
        reducer = EntityReducer("employee", "worker-42", snapchore=sbn_client.snapchore)

        # After each run completes:
        state = reducer.fold_run(
            run_id="run-123",
            proofs={"validate": proof1, "execute": proof2},
            amount=1500.00,
        )

        # Check state
        print(state.total_runs, state.total_amount)

        # Verify a checkpoint
        is_valid = reducer.verify_checkpoint(state.latest_checkpoint)
    """

    def __init__(
        self,
        entity_type: str,
        entity_id: str,
        snapchore: Any = None,
    ) -> None:
        self._entity_type = entity_type
        self._entity_id = entity_id
        self._snapchore = snapchore
        self._state = EntityState(entity_type=entity_type, entity_id=entity_id)
        self._checkpoints: List[EntityCheckpoint] = []

    @staticmethod
    def _compute_state_hash(state_data: Dict[str, Any]) -> str:
        """Compute a deterministic hash of the entity state."""
        raw = json.dumps(state_data, sort_keys=True, default=str, separators=(",", ":"))
        return hashlib.sha256(raw.encode()).hexdigest()

    def fold_run(
        self,
        run_id: str,
        proofs: Dict[str, ProofArtifact],
        amount: float = 0.0,
    ) -> EntityState:
        """Fold a completed run into the entity state."""
        self._state.total_runs += 1
        self._state.total_proofs += len(proofs)
        self._state.total_amount += amount
        self._state.run_ids.append(run_id)
        self._state.latest_run_id = run_id
        self._state.all_proofs[run_id] = list(proofs.values())

        # Count statuses
        for proof in proofs.values():
            if proof.succeeded:
                self._state.total_completed += 1
            else:
                self._state.total_failed += 1

        # Compute state hash
        hash_data = {
            "entity_type": self._entity_type,
            "entity_id": self._entity_id,
            "total_runs": self._state.total_runs,
            "total_amount": self._state.total_amount,
            "total_proofs": self._state.total_proofs,
        }
        self._state.state_hash = self._compute_state_hash(hash_data)

        # Create checkpoint
        status_counts = {
            "completed": self._state.total_completed,
            "failed": self._state.total_failed,
        }
        checkpoint = EntityCheckpoint(
            entity_type=self._entity_type,
            entity_id=self._entity_id,
            run_id=run_id,
            sequence=len(self._checkpoints),
            state_hash=self._state.state_hash,
            total_runs=self._state.total_runs,
            total_proofs=self._state.total_proofs,
            total_amount=self._state.total_amount,
            status_counts=status_counts,
        )
        self._checkpoints.append(checkpoint)
        self._state.latest_checkpoint = checkpoint

        return self._state

    def state(self) -> EntityState:
        return self._state

    def latest_checkpoint(self) -> Optional[EntityCheckpoint]:
        return self._checkpoints[-1] if self._checkpoints else None

    def checkpoints(self) -> List[EntityCheckpoint]:
        return list(self._checkpoints)

    def verify_checkpoint(self, checkpoint: EntityCheckpoint) -> bool:
        """Verify a checkpoint's state_hash matches recomputed state."""
        hash_data = {
            "entity_type": checkpoint.entity_type,
            "entity_id": checkpoint.entity_id,
            "total_runs": checkpoint.total_runs,
            "total_amount": checkpoint.total_amount,
            "total_proofs": checkpoint.total_proofs,
        }
        expected = self._compute_state_hash(hash_data)
        return checkpoint.state_hash == expected

    def reset(self) -> None:
        """Reset entity state (useful for testing or replay)."""
        self._state = EntityState(
            entity_type=self._entity_type,
            entity_id=self._entity_id,
        )
        self._checkpoints.clear()
