"""Persistence for checkpoint policies, history, and metadata."""

from __future__ import annotations

import os
import shutil
import stat
from pathlib import Path

import yaml

from gaol.checkpoint.models import (
    CheckpointHistory,
    CheckpointMetadata,
    CheckpointPolicy,
)
from gaol.config import get_data_dir


def _get_checkpoint_dir() -> Path:
    """Get the checkpoint data directory."""
    return get_data_dir() / "checkpoint"


class CheckpointStore:
    """Persistent storage for checkpoint policies, history, and metadata.

    Directory structure:
        ~/.local/share/gaol/checkpoint/
            policies/
                {container_name}.yaml
            history/
                {container_name}.yaml
            metadata/
                {container_name}/
                    {checkpoint_id}.yaml
            images/
                {container_name}/
                    {checkpoint_name}/
                        <CRIU image files>
    """

    def __init__(self, store_dir: Path | None = None) -> None:
        """Initialize the store."""
        self._store_dir = store_dir or _get_checkpoint_dir()
        self._policies_dir = self._store_dir / "policies"
        self._history_dir = self._store_dir / "history"
        self._metadata_dir = self._store_dir / "metadata"
        self._images_dir = self._store_dir / "images"

    @property
    def images_dir(self) -> Path:
        """Get the images directory."""
        return self._images_dir

    def _ensure_dirs(self) -> None:
        """Ensure store directories exist with proper permissions."""
        for dir_path in [
            self._store_dir,
            self._policies_dir,
            self._history_dir,
            self._metadata_dir,
            self._images_dir,
        ]:
            dir_path.mkdir(parents=True, exist_ok=True)
            os.chmod(dir_path, stat.S_IRWXU)  # 700

    def _write_yaml(self, path: Path, data: dict) -> None:
        """Write YAML with atomic write and proper permissions."""
        self._ensure_dirs()
        path.parent.mkdir(parents=True, exist_ok=True)
        temp_path = path.with_suffix(".tmp")
        with open(temp_path, "w") as f:
            yaml.dump(data, f, default_flow_style=False)
        os.chmod(temp_path, stat.S_IRUSR | stat.S_IWUSR)  # 600
        temp_path.rename(path)

    def _read_yaml(self, path: Path) -> dict | None:
        """Read YAML file if it exists."""
        if not path.exists():
            return None
        with open(path) as f:
            return yaml.safe_load(f) or {}

    # -------------------------------------------------------------------------
    # Policy operations
    # -------------------------------------------------------------------------

    def save_policy(self, policy: CheckpointPolicy) -> None:
        """Save or update a checkpoint policy."""
        path = self._policies_dir / f"{policy.container_name}.yaml"
        self._write_yaml(path, policy.model_dump(mode="json"))

    def get_policy(self, container_name: str) -> CheckpointPolicy | None:
        """Get checkpoint policy for a container."""
        path = self._policies_dir / f"{container_name}.yaml"
        data = self._read_yaml(path)
        if data is None:
            return None
        return CheckpointPolicy.model_validate(data)

    def delete_policy(self, container_name: str) -> bool:
        """Delete a policy. Returns True if deleted."""
        path = self._policies_dir / f"{container_name}.yaml"
        if path.exists():
            path.unlink()
            return True
        return False

    def list_policies(self) -> list[CheckpointPolicy]:
        """List all checkpoint policies."""
        if not self._policies_dir.exists():
            return []

        policies = []
        for path in self._policies_dir.glob("*.yaml"):
            try:
                data = self._read_yaml(path)
                if data:
                    policies.append(CheckpointPolicy.model_validate(data))
            except Exception:
                continue
        return policies

    def policy_exists(self, container_name: str) -> bool:
        """Check if a policy exists."""
        path = self._policies_dir / f"{container_name}.yaml"
        return path.exists()

    # -------------------------------------------------------------------------
    # History operations
    # -------------------------------------------------------------------------

    def save_history(self, history: CheckpointHistory) -> None:
        """Save checkpoint history for a container."""
        path = self._history_dir / f"{history.container_name}.yaml"
        self._write_yaml(path, history.model_dump(mode="json"))

    def get_history(self, container_name: str) -> CheckpointHistory:
        """Get checkpoint history for a container (creates if not exists)."""
        path = self._history_dir / f"{container_name}.yaml"
        data = self._read_yaml(path)
        if data is None:
            return CheckpointHistory(container_name=container_name)
        return CheckpointHistory.model_validate(data)

    def delete_history(self, container_name: str) -> bool:
        """Delete history for a container. Returns True if deleted."""
        path = self._history_dir / f"{container_name}.yaml"
        if path.exists():
            path.unlink()
            return True
        return False

    def list_histories(self) -> list[CheckpointHistory]:
        """List all checkpoint histories."""
        if not self._history_dir.exists():
            return []

        histories = []
        for path in self._history_dir.glob("*.yaml"):
            try:
                data = self._read_yaml(path)
                if data:
                    histories.append(CheckpointHistory.model_validate(data))
            except Exception:
                continue
        return histories

    # -------------------------------------------------------------------------
    # Metadata operations
    # -------------------------------------------------------------------------

    def save_checkpoint_metadata(self, metadata: CheckpointMetadata) -> None:
        """Save metadata for a checkpoint."""
        container_dir = self._metadata_dir / metadata.container_name
        path = container_dir / f"{metadata.checkpoint_id}.yaml"
        self._write_yaml(path, metadata.model_dump(mode="json"))

    def get_checkpoint_metadata(
        self, container_name: str, checkpoint_id: str
    ) -> CheckpointMetadata | None:
        """Get metadata for a specific checkpoint."""
        path = self._metadata_dir / container_name / f"{checkpoint_id}.yaml"
        data = self._read_yaml(path)
        if data is None:
            return None
        return CheckpointMetadata.model_validate(data)

    def get_checkpoint_by_name(
        self, container_name: str, checkpoint_name: str
    ) -> CheckpointMetadata | None:
        """Get metadata for a checkpoint by name."""
        container_dir = self._metadata_dir / container_name
        if not container_dir.exists():
            return None

        for path in container_dir.glob("*.yaml"):
            try:
                data = self._read_yaml(path)
                if data and data.get("checkpoint_name") == checkpoint_name:
                    return CheckpointMetadata.model_validate(data)
            except Exception:
                continue
        return None

    def list_checkpoints(self, container_name: str) -> list[CheckpointMetadata]:
        """List all checkpoints for a container."""
        container_dir = self._metadata_dir / container_name
        if not container_dir.exists():
            return []

        checkpoints = []
        for path in container_dir.glob("*.yaml"):
            try:
                data = self._read_yaml(path)
                if data:
                    checkpoints.append(CheckpointMetadata.model_validate(data))
            except Exception:
                continue

        # Sort by creation date (newest first)
        checkpoints.sort(key=lambda c: c.created_at, reverse=True)
        return checkpoints

    def delete_checkpoint_metadata(
        self, container_name: str, checkpoint_id: str
    ) -> bool:
        """Delete metadata for a checkpoint. Returns True if deleted."""
        path = self._metadata_dir / container_name / f"{checkpoint_id}.yaml"
        if path.exists():
            path.unlink()

            # Clean up empty container directory
            container_dir = path.parent
            if container_dir.exists() and not any(container_dir.iterdir()):
                container_dir.rmdir()

            return True
        return False

    def append_checkpoint(self, metadata: CheckpointMetadata) -> None:
        """Add a checkpoint to history and save metadata."""
        # Save metadata
        self.save_checkpoint_metadata(metadata)

        # Update history
        history = self.get_history(metadata.container_name)
        history.add_checkpoint(metadata)
        self.save_history(history)

    # -------------------------------------------------------------------------
    # Image directory operations
    # -------------------------------------------------------------------------

    def get_image_dir(self, container_name: str, checkpoint_name: str) -> Path:
        """Get the image directory path for a checkpoint.

        Returns the path (does not create the directory).
        """
        return self._images_dir / container_name / checkpoint_name

    def create_image_dir(self, container_name: str, checkpoint_name: str) -> Path:
        """Create and return the image directory for a checkpoint."""
        image_dir = self.get_image_dir(container_name, checkpoint_name)
        image_dir.mkdir(parents=True, exist_ok=True)
        os.chmod(image_dir, stat.S_IRWXU)  # 700
        return image_dir

    def delete_image_dir(self, container_name: str, checkpoint_name: str) -> bool:
        """Delete the image directory for a checkpoint. Returns True if deleted."""
        image_dir = self.get_image_dir(container_name, checkpoint_name)
        if image_dir.exists():
            shutil.rmtree(image_dir)

            # Clean up empty container directory
            container_dir = image_dir.parent
            if container_dir.exists() and not any(container_dir.iterdir()):
                container_dir.rmdir()

            return True
        return False

    def image_dir_exists(self, container_name: str, checkpoint_name: str) -> bool:
        """Check if the image directory exists."""
        return self.get_image_dir(container_name, checkpoint_name).exists()

    def get_image_dir_size(self, container_name: str, checkpoint_name: str) -> int:
        """Get the total size of a checkpoint image directory in bytes."""
        image_dir = self.get_image_dir(container_name, checkpoint_name)
        if not image_dir.exists():
            return 0

        total_size = 0
        for path in image_dir.rglob("*"):
            if path.is_file():
                total_size += path.stat().st_size
        return total_size

    def list_image_dirs(self, container_name: str) -> list[str]:
        """List all checkpoint image directories for a container."""
        container_dir = self._images_dir / container_name
        if not container_dir.exists():
            return []

        return [
            d.name for d in container_dir.iterdir()
            if d.is_dir()
        ]

    def delete_all_checkpoints(self, container_name: str) -> bool:
        """Delete all checkpoint data for a container.

        Removes metadata, history, images, and policy for a container.
        Returns True if any data was deleted.
        """
        deleted = False

        # Delete metadata
        container_metadata_dir = self._metadata_dir / container_name
        if container_metadata_dir.exists():
            shutil.rmtree(container_metadata_dir)
            deleted = True

        # Delete images
        container_images_dir = self._images_dir / container_name
        if container_images_dir.exists():
            shutil.rmtree(container_images_dir)
            deleted = True

        # Delete history
        if self.delete_history(container_name):
            deleted = True

        # Delete policy
        if self.delete_policy(container_name):
            deleted = True

        return deleted


# Global instance
_store: CheckpointStore | None = None


def get_checkpoint_store() -> CheckpointStore:
    """Get the global checkpoint store instance."""
    global _store
    if _store is None:
        _store = CheckpointStore()
    return _store
