"""
Merkle Epoch Manager

Batches proof chain entries into epochs for efficient anchoring to Rekor.

Default: 60-second epochs with automatic rotation.
"""

from __future__ import annotations

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

from receipts.proof_chain import ProofChainEntry

logger = logging.getLogger(__name__)


@dataclass
class MerkleEpoch:
    """A batch of entries anchored together."""

    epoch_id: str               # "2025-12-04T07:35:00"
    start_time: datetime
    end_time: Optional[datetime] = None
    entries: List[str] = field(default_factory=list)  # Entry IDs (hashes)
    merkle_root: Optional[str] = None

    # External anchoring (Rekor)
    rekor_uuid: Optional[str] = None
    rekor_log_index: Optional[int] = None
    rekor_timestamp: Optional[datetime] = None
    inclusion_proof: Optional[Dict[str, Any]] = None

    # TSA timestamp
    tsa_token: Optional[str] = None
    tsa_timestamp: Optional[datetime] = None

    def add_entry(self, entry_id: str) -> None:
        """Add an entry to this epoch."""
        self.entries.append(entry_id)

    def compute_root(self) -> str:
        """Compute Merkle root of all entries."""
        if not self.entries:
            self.merkle_root = "sha256:" + hashlib.sha256(b"empty_epoch").hexdigest()
            return self.merkle_root

        # Convert to leaf hashes (strip sha256: prefix)
        leaves = []
        for entry_id in self.entries:
            hash_hex = entry_id.replace("sha256:", "")
            leaves.append(bytes.fromhex(hash_hex))

        # Build Merkle tree
        while len(leaves) > 1:
            if len(leaves) % 2 == 1:
                leaves.append(leaves[-1])  # Duplicate last for odd count

            new_leaves = []
            for i in range(0, len(leaves), 2):
                combined = leaves[i] + leaves[i + 1]
                new_leaves.append(hashlib.sha256(combined).digest())
            leaves = new_leaves

        self.merkle_root = "sha256:" + leaves[0].hex()
        logger.debug(
            "Computed Merkle root: epoch=%s, entries=%d, root=%s",
            self.epoch_id,
            len(self.entries),
            self.merkle_root[:24] + "...",
        )
        return self.merkle_root

    def close(self) -> None:
        """Close the epoch and compute root."""
        self.end_time = datetime.now(timezone.utc)
        self.compute_root()
        logger.info(
            "Closed epoch: id=%s, entries=%d, duration=%s",
            self.epoch_id,
            len(self.entries),
            (self.end_time - self.start_time).total_seconds(),
        )

    def is_anchored(self) -> bool:
        """Check if epoch has been anchored to Rekor."""
        return self.rekor_uuid is not None

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to dict."""
        result = {
            "epoch_id": self.epoch_id,
            "start_time": self.start_time.isoformat(),
            "end_time": self.end_time.isoformat() if self.end_time else None,
            "entry_count": len(self.entries),
            "entries": self.entries,
            "merkle_root": self.merkle_root,
        }

        if self.rekor_uuid:
            result["rekor"] = {
                "uuid": self.rekor_uuid,
                "log_index": self.rekor_log_index,
                "timestamp": self.rekor_timestamp.isoformat() if self.rekor_timestamp else None,
            }

        if self.tsa_token:
            result["tsa"] = {
                "token": self.tsa_token,
                "timestamp": self.tsa_timestamp.isoformat() if self.tsa_timestamp else None,
            }

        return result

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "MerkleEpoch":
        """Deserialize from dict."""
        epoch = cls(
            epoch_id=data["epoch_id"],
            start_time=datetime.fromisoformat(data["start_time"]),
            end_time=datetime.fromisoformat(data["end_time"]) if data.get("end_time") else None,
            entries=data.get("entries", []),
            merkle_root=data.get("merkle_root"),
        )

        if "rekor" in data and data["rekor"]:
            epoch.rekor_uuid = data["rekor"].get("uuid")
            epoch.rekor_log_index = data["rekor"].get("log_index")
            if data["rekor"].get("timestamp"):
                epoch.rekor_timestamp = datetime.fromisoformat(data["rekor"]["timestamp"])

        if "tsa" in data and data["tsa"]:
            epoch.tsa_token = data["tsa"].get("token")
            if data["tsa"].get("timestamp"):
                epoch.tsa_timestamp = datetime.fromisoformat(data["tsa"]["timestamp"])

        return epoch


class EpochManager:
    """
    Manages Merkle epochs with automatic rotation.

    Default: 60-second epochs.
    """

    EPOCH_DURATION_SECONDS = 60

    def __init__(self, storage_dir: Optional[Path] = None):
        self.storage_dir = storage_dir
        self._current_epoch: Optional[MerkleEpoch] = None
        self._closed_epochs: List[MerkleEpoch] = []

        # Load existing epochs from storage
        if storage_dir:
            self._load_epochs()

    def _epoch_id(self, timestamp: datetime) -> str:
        """Generate epoch ID from timestamp (rounded to epoch boundary)."""
        epoch_start = timestamp.replace(
            second=(timestamp.second // self.EPOCH_DURATION_SECONDS) * self.EPOCH_DURATION_SECONDS,
            microsecond=0,
        )
        return epoch_start.strftime("%Y-%m-%dT%H:%M:%S")

    def _load_epochs(self) -> None:
        """Load epochs from storage directory."""
        if not self.storage_dir or not self.storage_dir.exists():
            return

        for epoch_file in sorted(self.storage_dir.glob("*.json")):
            try:
                with open(epoch_file) as f:
                    data = json.load(f)
                epoch = MerkleEpoch.from_dict(data)
                if epoch.end_time:  # Only load closed epochs
                    self._closed_epochs.append(epoch)
            except Exception as e:
                logger.warning("Failed to load epoch file %s: %s", epoch_file, e)

        logger.info("Loaded %d epochs from storage", len(self._closed_epochs))

    def add_entry(self, entry: ProofChainEntry) -> MerkleEpoch:
        """Add an entry to the current epoch, rotating if needed."""
        now = datetime.now(timezone.utc)
        epoch_id = self._epoch_id(now)

        # Check if we need a new epoch
        if self._current_epoch is None or self._current_epoch.epoch_id != epoch_id:
            # Close current epoch
            if self._current_epoch:
                self._close_current_epoch()

            # Start new epoch
            self._current_epoch = MerkleEpoch(
                epoch_id=epoch_id,
                start_time=now,
            )
            logger.debug("Started new epoch: %s", epoch_id)

        # Add entry
        self._current_epoch.add_entry(entry.id)

        return self._current_epoch

    def _close_current_epoch(self) -> None:
        """Close and persist the current epoch."""
        if not self._current_epoch:
            return

        if not self._current_epoch.entries:
            logger.debug("Skipping empty epoch: %s", self._current_epoch.epoch_id)
            return

        self._current_epoch.close()
        self._closed_epochs.append(self._current_epoch)

        # Persist if storage configured
        if self.storage_dir:
            self._persist_epoch(self._current_epoch)

    def _persist_epoch(self, epoch: MerkleEpoch) -> None:
        """Persist epoch to storage."""
        if not self.storage_dir:
            return

        self.storage_dir.mkdir(parents=True, exist_ok=True)

        # Use epoch_id as filename (replace : with -)
        safe_id = epoch.epoch_id.replace(":", "-")
        epoch_file = self.storage_dir / f"{safe_id}.json"

        with open(epoch_file, "w") as f:
            json.dump(epoch.to_dict(), f, indent=2)

        logger.debug("Persisted epoch: %s -> %s", epoch.epoch_id, epoch_file)

    def force_close(self) -> Optional[MerkleEpoch]:
        """Force close current epoch (e.g., for critical events)."""
        if self._current_epoch and self._current_epoch.entries:
            self._close_current_epoch()
            closed = self._closed_epochs[-1] if self._closed_epochs else None
            self._current_epoch = None
            return closed
        return None

    def get_current_epoch(self) -> Optional[MerkleEpoch]:
        """Get the current (open) epoch."""
        return self._current_epoch

    def get_closed_epochs(self) -> List[MerkleEpoch]:
        """Get all closed epochs."""
        return self._closed_epochs.copy()

    def get_epochs_for_anchoring(self, max_count: int = 10) -> List[MerkleEpoch]:
        """Get closed epochs that haven't been anchored to Rekor."""
        return [
            e for e in self._closed_epochs
            if not e.is_anchored()
        ][:max_count]

    def mark_anchored(
        self,
        epoch_id: str,
        rekor_uuid: str,
        rekor_log_index: int,
        rekor_timestamp: Optional[datetime] = None,
    ) -> bool:
        """Mark an epoch as anchored to Rekor."""
        for epoch in self._closed_epochs:
            if epoch.epoch_id == epoch_id:
                epoch.rekor_uuid = rekor_uuid
                epoch.rekor_log_index = rekor_log_index
                epoch.rekor_timestamp = rekor_timestamp or datetime.now(timezone.utc)

                if self.storage_dir:
                    self._persist_epoch(epoch)

                logger.info(
                    "Marked epoch as anchored: id=%s, rekor_uuid=%s, log_index=%d",
                    epoch_id,
                    rekor_uuid,
                    rekor_log_index,
                )
                return True

        logger.warning("Epoch not found for anchoring: %s", epoch_id)
        return False

    def get_epoch(self, epoch_id: str) -> Optional[MerkleEpoch]:
        """Get a specific epoch by ID."""
        if self._current_epoch and self._current_epoch.epoch_id == epoch_id:
            return self._current_epoch

        for epoch in self._closed_epochs:
            if epoch.epoch_id == epoch_id:
                return epoch

        return None

    def stats(self) -> Dict[str, Any]:
        """Get epoch manager statistics."""
        return {
            "current_epoch": self._current_epoch.epoch_id if self._current_epoch else None,
            "current_entries": len(self._current_epoch.entries) if self._current_epoch else 0,
            "closed_epochs": len(self._closed_epochs),
            "unanchored_epochs": len(self.get_epochs_for_anchoring()),
            "anchored_epochs": len([e for e in self._closed_epochs if e.is_anchored()]),
            "total_entries": sum(len(e.entries) for e in self._closed_epochs),
        }


# Global epoch manager
_epoch_manager: Optional[EpochManager] = None


def get_epoch_manager(storage_dir: Optional[Path] = None) -> EpochManager:
    """Get or create the global epoch manager."""
    global _epoch_manager
    if _epoch_manager is None:
        _epoch_manager = EpochManager(storage_dir=storage_dir)
    return _epoch_manager


def reset_epoch_manager() -> None:
    """Reset the global epoch manager (for testing)."""
    global _epoch_manager
    _epoch_manager = None


__all__ = [
    "MerkleEpoch",
    "EpochManager",
    "get_epoch_manager",
    "reset_epoch_manager",
]
