"""Checkpoint models and storage backends for workflow execution state."""

from __future__ import annotations

from abc import ABC, abstractmethod
from datetime import datetime
from pathlib import Path
from typing import Any, Generic, TypeVar
from uuid import uuid4

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.workflow.core.models import (
    StepStatus,
    WorkflowExecution,
)


class WorkflowCheckpoint(BaseModel):
    """Persisted snapshot of workflow execution state."""

    workflow_id: str
    workflow_version: str = "1.0.0"
    workflow_structure_hash: str
    checkpoint_id: str = Field(default_factory=lambda: str(uuid4()))
    created_at: datetime = Field(default_factory=datetime.now)
    checkpoint_type: str = "manual"
    execution: WorkflowExecution
    completed_step_ids: list[str] = Field(default_factory=list)
    pending_step_ids: list[str] = Field(default_factory=list)
    step_states: dict[str, dict[str, Any]] = Field(default_factory=dict)

    model_config = ConfigDict(extra="forbid")

    @classmethod
    def from_execution(
        cls,
        execution: WorkflowExecution,
        workflow_id: str,
        workflow_version: str,
        workflow_structure_hash: str,
        all_step_ids: list[str],
        checkpoint_type: str = "manual",
        step_states: dict[str, dict[str, Any]] | None = None,
    ) -> "WorkflowCheckpoint":
        completed_step_ids = [
            step_id
            for step_id, step_execution in execution.step_executions.items()
            if step_execution.status == StepStatus.COMPLETED
        ]
        pending_step_ids = [
            step_id
            for step_id in all_step_ids
            if step_id not in execution.step_executions
            or execution.step_executions[step_id].status
            in {StepStatus.PENDING, StepStatus.SUSPENDED}
        ]
        return cls(
            workflow_id=workflow_id,
            workflow_version=workflow_version,
            workflow_structure_hash=workflow_structure_hash,
            checkpoint_type=checkpoint_type,
            execution=execution.model_copy(deep=True),
            completed_step_ids=completed_step_ids,
            pending_step_ids=pending_step_ids,
            step_states=(step_states or {}).copy(),
        )


class CheckpointMetadata(BaseModel):
    """Lightweight checkpoint summary for listing operations."""

    checkpoint_id: str
    workflow_id: str
    workflow_version: str
    created_at: datetime
    checkpoint_type: str
    completed_steps: int
    pending_steps: int
    total_steps: int
    size_bytes: int | None = None

    model_config = ConfigDict(extra="forbid")

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint: WorkflowCheckpoint,
        size_bytes: int | None = None,
    ) -> "CheckpointMetadata":
        return cls(
            checkpoint_id=checkpoint.checkpoint_id,
            workflow_id=checkpoint.workflow_id,
            workflow_version=checkpoint.workflow_version,
            created_at=checkpoint.created_at,
            checkpoint_type=checkpoint.checkpoint_type,
            completed_steps=len(checkpoint.completed_step_ids),
            pending_steps=len(checkpoint.pending_step_ids),
            total_steps=len(checkpoint.completed_step_ids) + len(checkpoint.pending_step_ids),
            size_bytes=size_bytes,
        )


class CheckpointValidationResult(BaseModel):
    """Compatibility check result for workflow checkpoint resume."""

    is_valid: bool
    errors: list[str] = Field(default_factory=list)
    warnings: list[str] = Field(default_factory=list)
    can_resume: bool = False
    checkpoint_info: dict[str, Any] = Field(default_factory=dict)

    model_config = ConfigDict(extra="forbid")


CheckpointT = TypeVar("CheckpointT", bound=WorkflowCheckpoint)


class CheckpointStore(ABC, Generic[CheckpointT]):
    """Abstract storage backend for workflow checkpoints."""

    @abstractmethod
    async def save(self, checkpoint: CheckpointT) -> None:
        pass

    @abstractmethod
    async def load(self, checkpoint_id: str) -> CheckpointT | None:
        pass

    @abstractmethod
    async def load_latest(self, workflow_id: str) -> CheckpointT | None:
        pass

    @abstractmethod
    async def delete(self, checkpoint_id: str) -> bool:
        pass

    @abstractmethod
    async def list_metadata(
        self,
        workflow_id: str | None = None,
        limit: int = 100,
    ) -> list[CheckpointMetadata]:
        pass

    @abstractmethod
    async def cleanup_old(self, workflow_id: str, keep_last_n: int = 5) -> int:
        pass


class FileCheckpointStore(CheckpointStore[WorkflowCheckpoint]):
    """File-system checkpoint storage using JSON files per workflow."""

    def __init__(self, base_path: Path) -> None:
        self.base_path = Path(base_path)
        self.base_path.mkdir(parents=True, exist_ok=True)

    def _get_workflow_dir(self, workflow_id: str) -> Path:
        path = self.base_path / workflow_id
        path.mkdir(parents=True, exist_ok=True)
        return path

    def _get_checkpoint_path(self, workflow_id: str, checkpoint_id: str) -> Path:
        return self._get_workflow_dir(workflow_id) / f"{checkpoint_id}.json"

    async def save(self, checkpoint: WorkflowCheckpoint) -> None:
        path = self._get_checkpoint_path(checkpoint.workflow_id, checkpoint.checkpoint_id)
        path.write_text(checkpoint.model_dump_json(indent=2), encoding="utf-8")

    async def load(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
        for workflow_dir in self.base_path.iterdir():
            if not workflow_dir.is_dir():
                continue
            candidate = workflow_dir / f"{checkpoint_id}.json"
            if candidate.exists():
                return WorkflowCheckpoint.model_validate_json(
                    candidate.read_text(encoding="utf-8")
                )
        return None

    async def load_latest(self, workflow_id: str) -> WorkflowCheckpoint | None:
        workflow_dir = self._get_workflow_dir(workflow_id)
        files = sorted(
            workflow_dir.glob("*.json"),
            key=lambda path: path.stat().st_mtime,
            reverse=True,
        )
        if not files:
            return None
        return WorkflowCheckpoint.model_validate_json(files[0].read_text(encoding="utf-8"))

    async def delete(self, checkpoint_id: str) -> bool:
        for workflow_dir in self.base_path.iterdir():
            if not workflow_dir.is_dir():
                continue
            candidate = workflow_dir / f"{checkpoint_id}.json"
            if candidate.exists():
                candidate.unlink()
                return True
        return False

    async def list_metadata(
        self,
        workflow_id: str | None = None,
        limit: int = 100,
    ) -> list[CheckpointMetadata]:
        metadata: list[CheckpointMetadata] = []
        search_dirs = (
            [self._get_workflow_dir(workflow_id)]
            if workflow_id is not None
            else [path for path in self.base_path.iterdir() if path.is_dir()]
        )
        for workflow_dir in search_dirs:
            for checkpoint_path in workflow_dir.glob("*.json"):
                checkpoint = WorkflowCheckpoint.model_validate_json(
                    checkpoint_path.read_text(encoding="utf-8")
                )
                metadata.append(
                    CheckpointMetadata.from_checkpoint(
                        checkpoint,
                        size_bytes=checkpoint_path.stat().st_size,
                    )
                )
        metadata.sort(key=lambda item: item.created_at, reverse=True)
        return metadata[:limit]

    async def cleanup_old(self, workflow_id: str, keep_last_n: int = 5) -> int:
        workflow_dir = self._get_workflow_dir(workflow_id)
        files = sorted(
            workflow_dir.glob("*.json"),
            key=lambda path: path.stat().st_mtime,
            reverse=True,
        )
        stale_files = files[keep_last_n:]
        for path in stale_files:
            path.unlink()
        return len(stale_files)


class InMemoryCheckpointStore(CheckpointStore[WorkflowCheckpoint]):
    """Ephemeral in-memory checkpoint store, primarily for testing."""

    def __init__(self) -> None:
        self._checkpoints: dict[str, WorkflowCheckpoint] = {}
        self._by_workflow: dict[str, list[str]] = {}

    async def save(self, checkpoint: WorkflowCheckpoint) -> None:
        self._checkpoints[checkpoint.checkpoint_id] = checkpoint
        self._by_workflow.setdefault(checkpoint.workflow_id, []).append(
            checkpoint.checkpoint_id
        )

    async def load(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
        checkpoint = self._checkpoints.get(checkpoint_id)
        return checkpoint.model_copy(deep=True) if checkpoint else None

    async def load_latest(self, workflow_id: str) -> WorkflowCheckpoint | None:
        checkpoint_ids = self._by_workflow.get(workflow_id, [])
        checkpoints = [
            self._checkpoints[cid]
            for cid in checkpoint_ids
            if cid in self._checkpoints
        ]
        if not checkpoints:
            return None
        latest = max(checkpoints, key=lambda checkpoint: checkpoint.created_at)
        return latest.model_copy(deep=True)

    async def delete(self, checkpoint_id: str) -> bool:
        checkpoint = self._checkpoints.pop(checkpoint_id, None)
        if checkpoint is None:
            return False
        workflow_checkpoints = self._by_workflow.get(checkpoint.workflow_id, [])
        if checkpoint_id in workflow_checkpoints:
            workflow_checkpoints.remove(checkpoint_id)
        return True

    async def list_metadata(
        self,
        workflow_id: str | None = None,
        limit: int = 100,
    ) -> list[CheckpointMetadata]:
        checkpoints = (
            [
                self._checkpoints[cid]
                for cid in self._by_workflow.get(workflow_id, [])
                if cid in self._checkpoints
            ]
            if workflow_id is not None
            else list(self._checkpoints.values())
        )
        metadata = [CheckpointMetadata.from_checkpoint(checkpoint) for checkpoint in checkpoints]
        metadata.sort(key=lambda item: item.created_at, reverse=True)
        return metadata[:limit]

    async def cleanup_old(self, workflow_id: str, keep_last_n: int = 5) -> int:
        checkpoint_ids = self._by_workflow.get(workflow_id, [])
        checkpoints = [
            (cid, self._checkpoints[cid])
            for cid in checkpoint_ids
            if cid in self._checkpoints
        ]
        checkpoints.sort(key=lambda item: item[1].created_at, reverse=True)
        stale = checkpoints[keep_last_n:]
        for checkpoint_id, _ in stale:
            await self.delete(checkpoint_id)
        return len(stale)

    def clear(self) -> None:
        self._checkpoints.clear()
        self._by_workflow.clear()


class CheckpointConfig(BaseModel):
    """Configuration controlling workflow checkpoint auto-save behavior."""

    store: CheckpointStore[WorkflowCheckpoint] = Field(default_factory=InMemoryCheckpointStore)
    auto_save: bool = True
    save_interval_steps: int = 1
    auto_cleanup: bool = False
    keep_last_n: int = 5

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")

__all__ = [
    "CheckpointConfig",
    "CheckpointMetadata",
    "CheckpointStore",
    "CheckpointValidationResult",
    "FileCheckpointStore",
    "InMemoryCheckpointStore",
    "WorkflowCheckpoint",
]
