import asyncio
import time
import logging
import random

logger = logging.getLogger(__name__)
from typing import Any, Dict, List, Optional, Tuple
from dataclasses import dataclass, field
from natsort import natsorted

# Supported sort modes: (sort_key, order)
# order: 'asc' for ascending, 'desc' for descending
SUPPORTED_SORT_MODES = [
    ('name', 'asc'),
    ('name', 'desc'),
    ('date', 'asc'),
    ('date', 'desc'),
    ('size', 'asc'),
    ('size', 'desc'),
    ('usage', 'asc'),
    ('usage', 'desc'),
    ('versions_count', 'asc'),
    ('versions_count', 'desc'),
]
# Is this in use?

DISPLAY_NAME_MODES = {"model_name", "file_name"}


@dataclass
class ModelCache:
    """Cache structure for model data with extensible sorting."""

    raw_data: List[Dict[str, Any]]
    folders: List[str]
    version_index: Dict[int, Dict[str, Any]] = field(default_factory=dict)
    model_id_index: Dict[int, List[Dict[str, Any]]] = field(default_factory=dict)
    # Multi-valued companion to version_index: every local file entry of a
    # CivitAI model version, so versions with several downloaded files stay
    # consistent (#1058).
    version_files_index: Dict[int, List[Dict[str, Any]]] = field(default_factory=dict)
    name_display_mode: str = "model_name"
    _lock: Any = field(init=False, repr=False, default=None)
    # Cache for last sort: (sort_key, order, seed) -> sorted list
    _last_sort: Tuple[Optional[str], str, Optional[str]] = field(
        init=False, repr=False, default=(None, "asc", None)
    )
    _last_sorted_data: List[Dict[str, Any]] = field(
        init=False, repr=False, default_factory=list
    )

    def __post_init__(self):
        self._lock = asyncio.Lock()
        self._normalize_raw_data()
        self.name_display_mode = self._normalize_display_mode(self.name_display_mode)
        # Default sort on init
        asyncio.create_task(self.resort())
        self.rebuild_version_index()

    @staticmethod
    def _normalize_display_mode(value: Optional[str]) -> str:
        if isinstance(value, str) and value in DISPLAY_NAME_MODES:
            return value
        return "model_name"

    @staticmethod
    def _ensure_string(value: Any) -> str:
        """Return a safe string representation for metadata fields."""

        if isinstance(value, str):
            return value
        if value is None:
            return ""
        return str(value)

    def _normalize_item(self, item: Dict[str, Any]) -> None:
        """Ensure core metadata fields are present and string typed."""

        if not isinstance(item, dict):
            return

        for field in ("model_name", "file_name", "folder"):
            if field in item:
                item[field] = self._ensure_string(item.get(field))

    def _normalize_raw_data(self) -> None:
        """Normalize every cached entry before it is consumed."""

        for item in self.raw_data:
            self._normalize_item(item)

    def _get_display_name(self, item: Dict[str, Any]) -> str:
        """Return the value used for name-based sorting based on display settings."""

        if self.name_display_mode == "file_name":
            primary = self._ensure_string(item.get("file_name"))
            fallback = self._ensure_string(item.get("model_name"))
        else:
            primary = self._ensure_string(item.get("model_name"))
            fallback = self._ensure_string(item.get("file_name"))

        candidate = primary or fallback
        return candidate or ""

    @staticmethod
    def _normalize_version_id(value: Any) -> Optional[int]:
        """Normalize a potential version identifier into an integer."""

        if isinstance(value, int):
            return value
        if isinstance(value, str):
            try:
                return int(value)
            except ValueError:
                return None
        return None

    def rebuild_version_index(self) -> None:
        """Rebuild the version and model indexes from the current raw data."""

        self.version_index = {}
        self.model_id_index = {}
        self.version_files_index = {}
        for item in self.raw_data:
            self.add_to_version_index(item)

    def add_to_version_index(self, item: Dict[str, Any]) -> None:
        """Register a cache item in the version/model indexes if possible."""

        civitai_data = item.get('civitai') if isinstance(item, dict) else None
        if not isinstance(civitai_data, dict):
            return

        version_id = self._normalize_version_id(civitai_data.get('id'))
        if version_id is None:
            return

        self.version_index[version_id] = item

        # Register in the multi-valued index, deduplicated by file_path (#1058)
        files = self.version_files_index.setdefault(version_id, [])
        for entry in files:
            if entry is item or (
                isinstance(entry, dict)
                and entry.get('file_path') == item.get('file_path')
            ):
                break
        else:
            files.append(item)

        model_id = self._normalize_version_id(civitai_data.get('modelId'))
        if model_id is None:
            return

        descriptor = self._build_version_descriptor(item, civitai_data, version_id)
        if descriptor is None:
            return

        versions = self.model_id_index.setdefault(model_id, [])
        for index, existing in enumerate(versions):
            if existing.get('versionId') == descriptor['versionId']:
                versions[index] = descriptor
                break
        else:
            versions.append(descriptor)

    def remove_from_version_index(self, item: Dict[str, Any]) -> None:
        """Remove a cache item from the version/model indexes if present."""

        civitai_data = item.get('civitai') if isinstance(item, dict) else None
        if not isinstance(civitai_data, dict):
            return

        version_id = self._normalize_version_id(civitai_data.get('id'))
        if version_id is None:
            return

        # Drop only this file's entry from the multi-valued index (#1058)
        files = self.version_files_index.get(version_id)
        if files:
            remaining = [
                entry
                for entry in files
                if not (
                    entry is item
                    or (
                        isinstance(entry, dict)
                        and entry.get('file_path') == item.get('file_path')
                    )
                )
            ]
            if remaining:
                self.version_files_index[version_id] = remaining
            else:
                self.version_files_index.pop(version_id, None)

        # A surviving sibling file keeps the version present in the indexes
        sibling = (self.version_files_index.get(version_id) or [None])[0]

        existing = self.version_index.get(version_id)
        if existing is item or (
            isinstance(existing, dict)
            and existing.get('file_path') == item.get('file_path')
        ):
            if sibling is not None:
                self.version_index[version_id] = sibling
            else:
                self.version_index.pop(version_id, None)

        model_id = self._normalize_version_id(civitai_data.get('modelId'))
        if model_id is None:
            return

        versions = self.model_id_index.get(model_id)
        if not versions:
            return

        if sibling is not None:
            # Update the descriptor to reflect the surviving sibling file
            descriptor = self._build_version_descriptor(
                sibling,
                sibling.get('civitai') if isinstance(sibling, dict) else {},
                version_id,
            )
            for index, existing_desc in enumerate(versions):
                if existing_desc.get('versionId') == version_id:
                    if descriptor is not None:
                        versions[index] = descriptor
                    break
            return

        filtered = [v for v in versions if v.get('versionId') != version_id]
        if filtered:
            self.model_id_index[model_id] = filtered
        else:
            self.model_id_index.pop(model_id, None)

    def _build_version_descriptor(
        self,
        item: Dict[str, Any],
        civitai_data: Dict[str, Any],
        version_id: int,
    ) -> Optional[Dict[str, Any]]:
        """Create a lightweight descriptor for a version entry."""

        model_name = self._ensure_string(civitai_data.get('name'))
        file_name = self._ensure_string(item.get('file_name'))
        return {
            'versionId': version_id,
            'name': model_name,
            'fileName': file_name,
        }

    def get_versions_by_model_id(self, model_id: Any) -> List[Dict[str, Any]]:
        """Return cached version descriptors for a given model ID."""

        normalized_id = self._normalize_version_id(model_id)
        if normalized_id is None:
            return []

        versions = self.model_id_index.get(normalized_id, [])
        return [dict(version) for version in versions]

    def get_files_by_version_id(self, version_id: Any) -> List[Dict[str, Any]]:
        """Return every local file entry for a CivitAI model version (#1058)."""

        normalized_id = self._normalize_version_id(version_id)
        if normalized_id is None:
            return []

        return list(self.version_files_index.get(normalized_id, []))

    async def resort(self):
        """Resort cached data according to last sort mode if set"""
        async with self._lock:
            sort_key, order, seed = self._last_sort
            if sort_key is not None:
                sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
                self._last_sorted_data = sorted_data
                # Update folder list
            # else: do nothing

            all_folders = {
                self._ensure_string(item.get('folder'))
                for item in self.raw_data
                if isinstance(item, dict)
            }
            self.folders = sorted(list(all_folders), key=lambda x: x.lower())
            self.rebuild_version_index()

    def _sort_data(self, data: List[Dict[str, Any]], sort_key: str, order: str, seed: Optional[str] = None) -> List[Dict[str, Any]]:
        """Sort data by sort_key and order"""
        start_time = time.perf_counter()
        reverse = (order == 'desc')
        if sort_key == 'name':
            # Natural sort by configured display name, case-insensitive, with file_path as tie-breaker
            result = natsorted(
                data,
                key=lambda x: (
                    self._get_display_name(x).lower(),
                    x.get('file_path', '').lower()
                ),
                reverse=reverse
            )
        elif sort_key == 'date':
            # Sort by modified timestamp, fallback to name and path for stability
            result = sorted(
                data,
                key=lambda x: (
                    x.get('modified', 0.0),
                    self._get_display_name(x).lower(),
                    x.get('file_path', '').lower()
                ),
                reverse=reverse
            )
        elif sort_key == 'size':
            # Sort by file size, fallback to name and path for stability
            result = sorted(
                data,
                key=lambda x: (
                    x.get('size', 0),
                    self._get_display_name(x).lower(),
                    x.get('file_path', '').lower()
                ),
                reverse=reverse
            )
        elif sort_key == 'usage':
            # Sort by usage count, fallback to 0, then name and path for stability
            return sorted(
                data,
                key=lambda x: (
                    x.get('usage_count', 0),
                    self._get_display_name(x).lower(),
                    x.get('file_path', '').lower()
                ),
                reverse=reverse
            )
        elif sort_key == 'random':
            # Random shuffle seeded for stable pagination: the same seed
            # always yields the same order, so successive page requests
            # stay consistent while browsing.
            rng = random.Random(seed or 'random')
            result = list(data)
            rng.shuffle(result)
        elif sort_key == 'versions_count':
            # Pre-dedup sort: fall back to name sort.
            # Actual re-sort by version_count happens in get_paginated_data after dedup.
            result = natsorted(
                data,
                key=lambda x: (
                    self._get_display_name(x).lower(),
                    x.get('file_path', '').lower()
                ),
                reverse=reverse
            )
        else:
            # Fallback: no sort
            result = list(data)
        
        duration = time.perf_counter() - start_time
        if duration > 0.05:
            logger.debug("ModelCache._sort_data(%s, %s) for %d items took %.3fs", sort_key, order, len(data), duration)
        return result

    async def get_sorted_data(self, sort_key: str = 'name', order: str = 'asc', seed: Optional[str] = None) -> List[Dict[str, Any]]:
        """Get sorted data by sort_key and order, using cache if possible"""
        async with self._lock:
            cache_key = (sort_key, order, seed)
            if cache_key == self._last_sort:
                return self._last_sorted_data
            
            start_time = time.perf_counter()
            sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
            self._last_sort = cache_key
            self._last_sorted_data = sorted_data
            
            duration = time.perf_counter() - start_time
            if duration > 0.1:
                logger.debug("ModelCache.get_sorted_data(%s, %s) took %.3fs", sort_key, order, duration)
            
            return sorted_data

    async def update_name_display_mode(self, display_mode: str) -> None:
        """Update the display mode used for name sorting and refresh cached results."""

        normalized = self._normalize_display_mode(display_mode)
        async with self._lock:
            if self.name_display_mode == normalized:
                return

            self.name_display_mode = normalized

            sort_key, order, seed = self._last_sort
            if sort_key == 'name':
                self._last_sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)

    async def update_preview_url(self, file_path: str, preview_url: str, preview_nsfw_level: int) -> bool:
        """Update preview_url for a specific model in all cached data
        
        Args:
            file_path: The file path of the model to update
            preview_url: The new preview URL
            preview_nsfw_level: The NSFW level of the preview
            
        Returns:
            bool: True if the update was successful, False if the model wasn't found
        """
        async with self._lock:
            # Update in raw_data
            for item in self.raw_data:
                if item['file_path'] == file_path:
                    item['preview_url'] = preview_url
                    item['preview_nsfw_level'] = preview_nsfw_level
                    break
            else:
                return False  # Model not found
                    
            return True

    async def clear_preview_by_path(self, preview_file_path: str) -> int:
        """Clear ``preview_url`` for every cached entry referencing a file path.

        When a preview file has been deleted from disk, this removes its
        reference from all matching cache entries so the next list-API
        response returns an empty ``preview_url`` instead of a stale URL
        that produces 404s.

        Returns the number of entries that were updated.
        """
        normalized = preview_file_path.replace("\\", "/")
        cleared = 0
        async with self._lock:
            for item in self.raw_data:
                cached_url = item.get("preview_url", "")
                if cached_url.replace("\\", "/") == normalized:
                    item["preview_url"] = ""
                    item["preview_nsfw_level"] = 0
                    cleared += 1
        return cleared