from __future__ import annotations

import gzip
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, Field

from nexus.core.events.bus import EventBus
from nexus.core.events.types import CheckpointWrittenEvent
from nexus.resilience.errors import CheckpointError


class CheckpointData(BaseModel):
    """Serialized checkpoint payload."""

    checkpoint_id: str = Field(default_factory=lambda: str(uuid4()))
    session_id: str
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    node_id: str
    agent_states: dict[str, Any]
    graph_state: dict[str, Any]
    token_ledgers: dict[str, Any]
    metadata: dict[str, Any] = Field(default_factory=dict)


class CheckpointManager:
    """Writes and restores gzip-compressed JSON checkpoints."""

    def __init__(self, bus: EventBus, checkpoint_dir: str = ".nexus/checkpoints") -> None:
        self.bus = bus
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)

    async def write(
        self,
        session_id: str,
        node_id: str,
        agent_states: dict,
        graph_state: dict,
        token_ledgers: dict,
        metadata: dict | None = None,
    ) -> str:
        """Write a compressed checkpoint and emit CheckpointWrittenEvent."""
        try:
            ckpt = CheckpointData(
                session_id=session_id,
                node_id=node_id,
                agent_states=agent_states,
                graph_state=graph_state,
                token_ledgers=token_ledgers,
                metadata=metadata or {},
            )
            path = self.checkpoint_dir / f"{ckpt.checkpoint_id}.ckpt.gz"
            payload = ckpt.model_dump_json().encode("utf-8")
            with gzip.open(path, "wb") as fh:
                fh.write(payload)

            await self.bus.publish_and_wait(
                CheckpointWrittenEvent(
                    source_id="checkpoint_manager",
                    checkpoint_id=ckpt.checkpoint_id,
                    path=str(path),
                    state_size_bytes=len(payload),
                )
            )
            return ckpt.checkpoint_id
        except OSError as exc:
            raise CheckpointError(f"Failed writing checkpoint: {exc}") from exc

    def load(self, checkpoint_id: str) -> CheckpointData:
        """Load checkpoint payload by id."""
        path = self.checkpoint_dir / f"{checkpoint_id}.ckpt.gz"
        if not path.exists():
            raise CheckpointError(f"Checkpoint not found: {checkpoint_id}")
        try:
            with gzip.open(path, "rb") as fh:
                data = fh.read().decode("utf-8")
            return CheckpointData.model_validate_json(data)
        except (OSError, ValueError) as exc:
            raise CheckpointError(f"Failed loading checkpoint {checkpoint_id}: {exc}") from exc

    def list_checkpoints(self, session_id: str | None = None) -> list[dict[str, Any]]:
        """List checkpoints with optional session filter."""
        out: list[dict[str, Any]] = []
        for path in sorted(self.checkpoint_dir.glob("*.ckpt.gz")):
            try:
                with gzip.open(path, "rb") as fh:
                    data = json.loads(fh.read().decode("utf-8"))
                if session_id and data.get("session_id") != session_id:
                    continue
                out.append(
                    {
                        "checkpoint_id": data.get("checkpoint_id"),
                        "session_id": data.get("session_id"),
                        "timestamp": data.get("timestamp"),
                        "node_id": data.get("node_id"),
                    }
                )
            except (OSError, ValueError):
                continue
        return out

    def delete(self, checkpoint_id: str) -> bool:
        """Delete checkpoint by id."""
        path = self.checkpoint_dir / f"{checkpoint_id}.ckpt.gz"
        if not path.exists():
            return False
        try:
            path.unlink()
            return True
        except OSError as exc:
            raise CheckpointError(f"Failed deleting checkpoint {checkpoint_id}: {exc}") from exc
