"""Local model cache/library for Kiln.

Agents save generated or downloaded models locally with tagged metadata
(source, prompt, dimensions, print history) so they can reuse them
across jobs without re-downloading or re-generating.

Files are stored in ``~/.kiln/model_cache/`` (override with
``KILN_MODEL_CACHE_DIR``).  Each file is stored under its SHA-256 hash
to enable automatic deduplication.
"""

from __future__ import annotations

import hashlib
import logging
import os
import secrets
import shutil
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

from kiln.persistence import KilnDB

logger = logging.getLogger(__name__)

_DEFAULT_CACHE_DIR = os.path.join(str(Path.home()), ".kiln", "model_cache")


@dataclass
class ModelCacheEntry:
    """A single cached 3D model file with metadata."""

    cache_id: str
    file_name: str
    file_path: str
    file_hash: str
    file_size_bytes: int
    source: str
    source_id: str | None = None
    prompt: str | None = None
    tags: list[str] = field(default_factory=list)
    dimensions: dict[str, float] | None = None
    print_count: int = 0
    last_printed_at: float | None = None
    created_at: float = 0.0
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Serialise to a plain dict suitable for JSON output."""
        d = asdict(self)
        return d


def _compute_sha256(file_path: str) -> str:
    """Compute the SHA-256 hex digest of a file."""
    h = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


class ModelCache:
    """Manages a local cache of 3D model files with metadata.

    Parameters:
        db: The :class:`KilnDB` instance for metadata storage.
        cache_dir: Directory to store cached model files.  Defaults to
            ``KILN_MODEL_CACHE_DIR`` or ``~/.kiln/model_cache/``.
    """

    def __init__(self, db: KilnDB, cache_dir: str | None = None) -> None:
        self._db = db
        self._cache_dir = cache_dir or os.environ.get("KILN_MODEL_CACHE_DIR", _DEFAULT_CACHE_DIR)
        os.makedirs(self._cache_dir, exist_ok=True)

    @property
    def cache_dir(self) -> str:
        """The filesystem path of the cache directory."""
        return self._cache_dir

    def add(
        self,
        file_path: str,
        *,
        source: str,
        source_id: str | None = None,
        prompt: str | None = None,
        tags: list[str] | None = None,
        dimensions: dict[str, float] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> ModelCacheEntry:
        """Add a model file to the cache.

        The file is copied into the cache directory under a subdirectory
        named after its SHA-256 hash (enabling automatic deduplication).
        If a file with the same hash already exists, the existing entry
        is returned without re-copying.

        Args:
            file_path: Path to the model file on disk.
            source: Origin of the model (e.g. ``"thingiverse"``,
                ``"meshy"``, ``"upload"``).
            source_id: Marketplace thing ID, generation job ID, etc.
            prompt: For AI-generated models, the prompt used.
            tags: Descriptive tags for search.
            dimensions: Bounding box in mm: ``{"x": ..., "y": ..., "z": ...}``.
            metadata: Arbitrary extra data.

        Returns:
            The :class:`ModelCacheEntry` for the cached file.

        Raises:
            FileNotFoundError: If *file_path* does not exist.
            ValueError: If *source* is empty.
        """
        if not os.path.isfile(file_path):
            raise FileNotFoundError(f"Model file not found: {file_path}")
        if not source:
            raise ValueError("source is required")

        file_hash = _compute_sha256(file_path)

        # Dedup check — return existing entry if the same file is cached.
        existing = self.get_by_hash(file_hash)
        if existing is not None:
            logger.info(
                "File already cached (hash=%s, cache_id=%s)",
                file_hash[:12],
                existing.cache_id,
            )
            return existing

        # Copy file into cache directory.
        file_name = os.path.basename(file_path)
        hash_dir = os.path.join(self._cache_dir, file_hash[:16])
        os.makedirs(hash_dir, exist_ok=True)
        dest_path = os.path.join(hash_dir, file_name)
        shutil.copy2(file_path, dest_path)
        file_size = os.path.getsize(dest_path)

        cache_id = secrets.token_hex(8)
        now = time.time()

        entry = ModelCacheEntry(
            cache_id=cache_id,
            file_name=file_name,
            file_path=dest_path,
            file_hash=file_hash,
            file_size_bytes=file_size,
            source=source,
            source_id=source_id,
            prompt=prompt,
            tags=tags or [],
            dimensions=dimensions,
            print_count=0,
            last_printed_at=None,
            created_at=now,
            metadata=metadata or {},
        )

        self._db.save_cache_entry(entry)
        logger.info("Cached model %s (%s, %d bytes)", cache_id, file_name, file_size)
        return entry

    def get(self, cache_id: str) -> ModelCacheEntry | None:
        """Return a cache entry by ID, or ``None`` if not found."""
        return self._db.get_cache_entry(cache_id)

    def get_by_hash(self, file_hash: str) -> ModelCacheEntry | None:
        """Return a cache entry by file hash (dedup check), or ``None``."""
        return self._db.get_cache_entry_by_hash(file_hash)

    def search(
        self,
        *,
        query: str | None = None,
        source: str | None = None,
        tags: list[str] | None = None,
        limit: int = 20,
    ) -> list[ModelCacheEntry]:
        """Search cached models by name, tags, source, or prompt text.

        Args:
            query: Free-text search against file name, prompt, and tags.
            source: Filter by source (e.g. ``"thingiverse"``).
            tags: Filter entries that contain ALL of these tags.
            limit: Maximum results to return.
        """
        return self._db.search_cache(
            query=query,
            source=source,
            tags=tags,
            limit=limit,
        )

    def record_print(self, cache_id: str) -> None:
        """Increment print_count and update last_printed_at for a cached model."""
        self._db.record_cache_print(cache_id)

    def delete(self, cache_id: str) -> bool:
        """Remove a model from the cache directory and DB.

        Returns ``True`` if the entry was found and deleted.
        """
        entry = self.get(cache_id)
        if entry is None:
            return False

        # Remove file from disk.
        try:
            if os.path.isfile(entry.file_path):
                os.remove(entry.file_path)
                # Clean up hash directory if empty.
                parent = os.path.dirname(entry.file_path)
                if os.path.isdir(parent) and not os.listdir(parent):
                    os.rmdir(parent)
        except OSError as exc:
            logger.warning("Could not remove cached file %s: %s", entry.file_path, exc)

        self._db.delete_cache_entry(cache_id)
        logger.info("Deleted cached model %s (%s)", cache_id, entry.file_name)
        return True

    def list_all(self, *, limit: int = 50, offset: int = 0) -> list[ModelCacheEntry]:
        """List all cached models, newest first."""
        return self._db.list_cache_entries(limit=limit, offset=offset)


# ---------------------------------------------------------------------------
# Module-level singleton
# ---------------------------------------------------------------------------

_model_cache: ModelCache | None = None


def get_model_cache() -> ModelCache:
    """Return the module-level :class:`ModelCache` singleton.

    Lazily created on first call, using the default :func:`get_db`
    instance and the default or env-configured cache directory.
    """
    global _model_cache
    if _model_cache is None:
        from kiln.persistence import get_db

        _model_cache = ModelCache(db=get_db())
    return _model_cache
