"""Cache manager for generation results."""

import json
import pickle
from dataclasses import asdict, is_dataclass
from enum import Enum
from pathlib import Path
from typing import Generic, TypeVar

from pydantic import BaseModel

from hezor_common.utilities.cache_manager.system_cache import get_system_cache_path

T = TypeVar("T")


class CacheType(Enum):
    """Cache type enumeration."""

    CREATION = "serialized_creations"
    CHAPTER = "serialized_chapters"
    SECTION = "serialized_sections"
    AGENT_GENERATION = "agent_generations"
    EXECUTION = "executions"
    HTTP_RESPONSES = "http_responses"

    # Generic cache type for other uses
    GENERIC = "generic_cache"


class CacheManager(Generic[T]):
    """Generic cache manager for generation results.

    This manager handles caching of generation results to avoid redundant
    computations. It uses hash-based file naming to identify cached results.

    Supports Pydantic models, dataclasses (JSON serialization), and regular Python
    objects (pickle serialization).

    Parameters
    ----------
    cache_type : CacheType
        The type of cache to manage
    base_dir : str | Path | None, optional
        Base cache directory. If None, uses get_system_cache_path() to find
        project cache dir. Defaults to None (recommended for project usage).
    model_class : Type[T] | None, optional
        Pydantic model or dataclass for JSON serialization. If provided, uses JSON;
        otherwise uses pickle.

    Examples
    --------
    >>> from hezor_common.data_model.creations.core import CreationGenerateResult
    >>> # Uses project cache directory automatically
    >>> manager = CacheManager[CreationGenerateResult](
    ...     CacheType.CREATION,
    ...     model_class=CreationGenerateResult
    ... )
    >>> # Check if cached result exists (by hash only)
    >>> input_hash = compute_hash(creation_model.model_dump_json())
    >>> cached = manager.get(input_hash)
    >>> if cached is None:
    ...     # Generate new result
    ...     result = generate_creation(query)
    ...     # Save to cache with identifier
    ...     manager.save("creation_id", input_hash, result)
    """

    def __init__(
        self,
        cache_type: CacheType,
        base_dir: str | Path | None = None,
        model_class: type[T] | None = None,
    ) -> None:
        """Initialize cache manager.

        Parameters
        ----------
        cache_type : CacheType
            The type of cache to manage
        base_dir : str | Path | None, optional
            Base cache directory. If None, uses get_system_cache_path() to find
            project cache dir. Defaults to None (recommended for project usage).
        model_class : Type[T] | None, optional
            Pydantic model or dataclass for JSON serialization
        """

        self.cache_type = cache_type

        # Use get_system_cache_path() if base_dir not specified
        if base_dir is None:
            self.base_dir = get_system_cache_path()
        else:
            self.base_dir = Path(base_dir)

        self.cache_dir = self.base_dir / cache_type.value
        self.model_class = model_class
        # Use JSON for Pydantic models or dataclasses
        self.use_json = model_class is not None and (
            is_dataclass(model_class)
            or (isinstance(model_class, type) and issubclass(model_class, BaseModel))
        )

        # Ensure cache directory exists
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def _get_extension(self) -> str:
        """Get file extension for cache type."""
        extensions = {
            CacheType.CREATION: ".creation",
            CacheType.CHAPTER: ".chapter",
            CacheType.SECTION: ".section",
            CacheType.EXECUTION: ".exec",
            CacheType.AGENT_GENERATION: ".gen",
            CacheType.HTTP_RESPONSES: ".http",
            CacheType.GENERIC: ".cache",
        }
        return extensions[self.cache_type]

    def _build_filename(self, identifier: str, content_hash: str) -> str:
        """Build cache filename.

        Parameters
        ----------
        identifier : str
            Unique identifier (e.g., creation_id, chapter_id, or dd_hhmm for agent)
        content_hash : str
            Hash of the input content

        Returns
        -------
        str
            Filename in format: {identifier}__{hash}.{extension}
        """
        ext = self._get_extension()
        return f"{identifier}__{content_hash}{ext}"

    def get_cache_path(self, identifier: str, content_hash: str) -> Path:
        """Get full cache file path.

        Parameters
        ----------
        identifier : str
            Unique identifier
        content_hash : str
            Hash of the input content

        Returns
        -------
        Path
            Full path to cache file
        """
        filename = self._build_filename(identifier, content_hash)
        return self.cache_dir / filename

    def _find_by_hash(self, content_hash: str) -> Path | None:
        """Find cache file by hash only.

        Parameters
        ----------
        content_hash : str
            Hash of the input content

        Returns
        -------
        Optional[Path]
            Path to cache file if found, None otherwise
        """
        ext = self._get_extension()
        pattern = f"*__{content_hash}{ext}"
        matches = list(self.cache_dir.glob(pattern))
        return matches[0] if matches else None

    def exists(self, content_hash: str) -> bool:
        """Check if cached result exists (by hash only).

        Parameters
        ----------
        content_hash : str
            Hash of the input content

        Returns
        -------
        bool
            True if cached result exists
        """
        return self._find_by_hash(content_hash) is not None

    def get(self, content_hash: str) -> T | None:
        """Get cached result if exists (by hash only).

        Parameters
        ----------
        content_hash : str
            Hash of the input content

        Returns
        -------
        Optional[T]
            Cached result if exists, None otherwise

        Examples
        --------
        >>> manager = CacheManager[dict](CacheType.SECTION)
        >>> result = manager.get("abc123hash")
        >>> if result is not None:
        ...     print("Cache hit!")
        """
        cache_path = self._find_by_hash(content_hash)
        if cache_path is None:
            return None

        try:
            if self.use_json and self.model_class is not None:
                # Try JSON first
                try:
                    with open(cache_path, encoding="utf-8") as f:
                        data = json.load(f)

                    # Handle Pydantic models
                    try:
                        if issubclass(self.model_class, BaseModel):
                            return self.model_class.model_validate(data)
                    except TypeError:
                        # issubclass raises TypeError if not a class
                        pass

                    # Handle dataclasses
                    if is_dataclass(self.model_class):
                        return self.model_class(**data)

                    # Fallback: return raw data
                    return data  # type: ignore[return-value]
                except (json.JSONDecodeError, UnicodeDecodeError):
                    # JSON failed, try pickle (fallback from save)
                    with open(cache_path, "rb") as f:
                        return pickle.load(f)
            else:
                # Use pickle deserialization
                with open(cache_path, "rb") as f:
                    return pickle.load(f)
        except Exception:
            # If cache is corrupted, return None
            return None

    def save(self, identifier: str, content_hash: str, result: T) -> Path:
        """Save result to cache.

        Parameters
        ----------
        identifier : str
            Unique identifier
        content_hash : str
            Hash of the input content
        result : T
            Result to cache

        Returns
        -------
        Path
            Path to saved cache file

        Examples
        --------
        >>> manager = CacheManager[dict](CacheType.CHAPTER)
        >>> cache_path = manager.save("chapter_456", "def456hash", chapter_result)
        >>> print(f"Saved to {cache_path}")
        """
        cache_path = self.get_cache_path(identifier, content_hash)

        if self.use_json:
            try:
                with open(cache_path, "w", encoding="utf-8") as f:
                    # Handle Pydantic models
                    if isinstance(result, BaseModel):
                        json.dump(result.model_dump(), f, ensure_ascii=False, indent=2)
                    # Handle dataclasses (check instance, not class)
                    elif is_dataclass(result) and not isinstance(result, type):
                        json.dump(asdict(result), f, ensure_ascii=False, indent=2)
                    else:
                        json.dump(result, f, ensure_ascii=False, indent=2)
            except (TypeError, ValueError):
                # JSON serialization failed, fallback to pickle
                with open(cache_path, "wb") as f:
                    pickle.dump(result, f)
        else:
            # Use pickle serialization
            with open(cache_path, "wb") as f:
                pickle.dump(result, f)

        return cache_path

    def delete(self, content_hash: str) -> bool:
        """Delete cached result (by hash only).

        Parameters
        ----------
        content_hash : str
            Hash of the input content

        Returns
        -------
        bool
            True if file was deleted, False if it didn't exist
        """
        cache_path = self._find_by_hash(content_hash)

        if cache_path is not None:
            cache_path.unlink()
            return True
        return False

    def list_all(self) -> list[Path]:
        """List all cached files for this cache type.

        Returns
        -------
        list[Path]
            List of all cache file paths
        """
        ext = self._get_extension()
        return list(self.cache_dir.glob(f"*{ext}"))

    def clear_all(self) -> int:
        """Clear all cached files for this cache type.

        Returns
        -------
        int
            Number of files deleted
        """
        files = self.list_all()
        count = 0
        for file_path in files:
            file_path.unlink()
            count += 1
        return count
