entity_database_service.py

"""
Service layer for entity database file-based operations

This module provides high-level services for loading, managing, and accessing
manually curated entity reference images from the filesystem.

Implements async-first operations following the MemecoinService pattern.
"""

import asyncio
import base64
import logging
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from src.constants import ENTITY_DATABASE_ROOT
from src.models.entity.entity_database_models import DatabaseEntity, EntityDatabaseSnapshot
from src.models.entity.entity_models import EntityType

logger = logging.getLogger(__name__)


class EntityDatabaseService:
    """
    Service for filesystem-based entity database operations

    This service handles loading entity data from the filesystem, managing
    the entity_database directory structure, and providing caching for
    performance optimization.

    Implements async-first operations with true async methods.
    All blocking I/O wrapped with asyncio.to_thread().

    Directory structure expected:
        res/entity_database/
            {category_name}/
                {entity_name}/
                    example_001.jpg
                    example_002.jpg
                    ...

    Example:
        res/entity_database/
            logo/
                bnb_logo/
                    example_001.jpg
                    example_002.jpg
            person/
                vitalik_buterin/
                    example_001.jpg
    """

    def __init__(
        self,
        database_root: Path = ENTITY_DATABASE_ROOT,
        cache_timeout: int = 60,
    ):
        """
        Initialize the entity database service

        Args:
            database_root: Root directory for entity database (default from constants)
            cache_timeout: Cache timeout in seconds (default 60)
        """
        self.database_root = database_root
        self.cache_timeout = cache_timeout
        self._lock = asyncio.Lock()
        self._snapshot_cache: Optional[Tuple[EntityDatabaseSnapshot, float]] = None

    async def _ensure_database_directory(self) -> Path:
        """
        Ensure entity database directory exists

        Returns:
            Path to entity database root

        Raises:
            PermissionError: If directory cannot be created or accessed
        """
        try:
            # Check if directory exists (blocking I/O)
            dir_exists = await asyncio.to_thread(self.database_root.exists)

            if not dir_exists:
                logger.warning(
                    f"Entity database directory does not exist: {self.database_root}"
                )
                logger.info("Creating entity database directory...")
                await asyncio.to_thread(self.database_root.mkdir, parents=True, exist_ok=True)

            # Verify it's a directory
            is_dir = await asyncio.to_thread(self.database_root.is_dir)
            if not is_dir:
                raise PermissionError(
                    f"Entity database path exists but is not a directory: {self.database_root}"
                )

            return self.database_root

        except Exception as e:
            logger.error(f"Failed to ensure entity database directory: {e}")
            raise

    def _is_cache_valid(
        self, cache_entry: Optional[Tuple[EntityDatabaseSnapshot, float]]
    ) -> bool:
        """
        Check if cache entry is still valid

        Args:
            cache_entry: Tuple of (snapshot, timestamp) or None

        Returns:
            True if cache is valid, False otherwise
        """
        if cache_entry is None:
            return False

        _, timestamp = cache_entry
        return time.time() - timestamp < self.cache_timeout

    async def _scan_category_directory(
        self, category_path: Path, category: EntityType
    ) -> List[DatabaseEntity]:
        """
        Scan a category directory and load all entities (supports recursive subdirectories)

        Handles both flat and hierarchical directory structures:
        - Flat: logo/bnb_logo/example_001.jpg
        - Hierarchical: logo/crypto/bnb_chain/bnb_chain_001.jpg

        Args:
            category_path: Path to category directory
            category: EntityType for this category

        Returns:
            List of DatabaseEntity objects for this category
        """

        def _scan_entity_directory(
            entity_dir: Path, subcategory: str | None = None
        ) -> DatabaseEntity | None:
            """
            Scan a single entity directory for images

            Args:
                entity_dir: Path to entity directory
                subcategory: Optional subcategory name (e.g., 'crypto', 'brand')

            Returns:
                DatabaseEntity if images found, None otherwise
            """
            entity_name = entity_dir.name
            display_name = entity_name.replace("_", " ").title()

            # Find all image files (ANY .jpg or .png, not just "example_*")
            example_paths = []

            # Glob for .jpg files
            for image_file in entity_dir.glob("*.jpg"):
                if image_file.is_file():
                    example_paths.append(image_file)

            # Glob for .png files
            for image_file in entity_dir.glob("*.png"):
                if image_file.is_file():
                    example_paths.append(image_file)

            # Sort for consistent ordering
            example_paths.sort()

            # Skip entities with no images
            if not example_paths:
                logger.debug(
                    f"No images found for entity: {entity_name} in {category.value}"
                    + (f"/{subcategory}" if subcategory else "")
                )
                return None

            # Create database entity with subcategory
            entity = DatabaseEntity(
                name=entity_name,
                display_name=display_name,
                category=category,
                subcategory=subcategory,
                example_image_paths=example_paths,
                aliases=[],  # Could be loaded from metadata file in future
            )

            logger.debug(
                f"Loaded entity: {display_name} ({len(example_paths)} images)"
                + (f" [subcategory: {subcategory}]" if subcategory else "")
            )
            return entity

        def _scan_sync() -> List[DatabaseEntity]:
            """Synchronous filesystem scanning (wrapped in to_thread)"""
            entities = []

            if not category_path.exists() or not category_path.is_dir():
                logger.debug(
                    f"Category directory not found or not a directory: {category_path}"
                )
                return entities

            # Iterate through items in category directory
            for item in category_path.iterdir():
                if not item.is_dir():
                    continue

                # Check if this is a subcategory directory or entity directory
                # Heuristic: If directory contains subdirectories, it's a subcategory
                # Otherwise, it's an entity directory
                subdirs = [d for d in item.iterdir() if d.is_dir()]

                if subdirs:
                    # This is a subcategory directory (e.g., logo/crypto/)
                    subcategory_name = item.name
                    logger.debug(
                        f"Found subcategory: {category.value}/{subcategory_name}"
                    )

                    # Scan all entity directories in this subcategory
                    for entity_dir in subdirs:
                        entity = _scan_entity_directory(entity_dir, subcategory_name)
                        if entity:
                            entities.append(entity)
                else:
                    # This is a direct entity directory (e.g., person/vitalik_buterin/)
                    entity = _scan_entity_directory(item, subcategory=None)
                    if entity:
                        entities.append(entity)

            return entities

        try:
            entities = await asyncio.to_thread(_scan_sync)
            logger.info(
                f"Loaded {len(entities)} entities from category: {category.value}"
            )
            return entities

        except Exception as e:
            logger.error(f"Failed to scan category {category.value}: {e}")
            return []

    async def load_database_snapshot(
        self, use_cache: bool = True
    ) -> EntityDatabaseSnapshot:
        """
        Load complete entity database snapshot from filesystem

        Args:
            use_cache: Whether to use cached snapshot if available

        Returns:
            EntityDatabaseSnapshot with all entities organized by category

        Raises:
            PermissionError: If database directory cannot be accessed
        """
        async with self._lock:
            # Use cache if valid
            if use_cache and self._is_cache_valid(self._snapshot_cache):
                logger.debug("Returning cached entity database snapshot")
                return self._snapshot_cache[0]

            logger.info("Loading entity database from filesystem...")

            # Ensure database directory exists
            await self._ensure_database_directory()

            # Scan all category directories
            entities_by_category: Dict[EntityType, List[DatabaseEntity]] = {}
            total_entity_count = 0
            total_image_count = 0
            categories_loaded = []

            for entity_type in EntityType:
                category_path = self.database_root / entity_type.value

                # Scan this category
                entities = await self._scan_category_directory(category_path, entity_type)

                if entities:
                    entities_by_category[entity_type] = entities
                    categories_loaded.append(entity_type)
                    total_entity_count += len(entities)

                    # Count images
                    for entity in entities:
                        total_image_count += len(entity.example_image_paths)

            # Create snapshot
            snapshot = EntityDatabaseSnapshot(
                entities_by_category=entities_by_category,
                total_entity_count=total_entity_count,
                total_image_count=total_image_count,
                categories_loaded=categories_loaded,
            )

            # Update cache
            self._snapshot_cache = (snapshot, time.time())

            logger.info(
                f"✅ Loaded entity database: {total_entity_count} entities, "
                f"{total_image_count} images across {len(categories_loaded)} categories"
            )

            return snapshot

    async def get_all_entities_by_category(
        self, use_cache: bool = True
    ) -> Dict[EntityType, List[DatabaseEntity]]:
        """
        Get all entities organized by category

        Args:
            use_cache: Whether to use cached data if available

        Returns:
            Dictionary mapping EntityType to list of DatabaseEntity objects
        """
        snapshot = await self.load_database_snapshot(use_cache=use_cache)
        return snapshot.entities_by_category

    async def get_entity_by_name(
        self,
        name: str,
        category: EntityType = None,
        use_cache: bool = True,
    ) -> Optional[DatabaseEntity]:
        """
        Find entity by name or alias (case-insensitive)

        Args:
            name: Entity name to search for
            category: Optional category to narrow search
            use_cache: Whether to use cached data

        Returns:
            DatabaseEntity if found, None otherwise
        """
        snapshot = await self.load_database_snapshot(use_cache=use_cache)
        return snapshot.find_entity_by_name(name, category)

    async def find_entity_by_display_name(
        self,
        display_name: str,
        category: EntityType = None,
        use_cache: bool = True,
    ) -> Optional[DatabaseEntity]:
        """
        Find entity by display name (case-insensitive)

        Used for resolving LLM-returned display names back to database keys.
        Searches only the display_name field, not aliases or directory names.

        Args:
            display_name: Display name to search for (e.g., "Bnb Chain", "Vitalik Buterin")
            category: Optional category to narrow search
            use_cache: Whether to use cached data

        Returns:
            DatabaseEntity if found by display_name, None otherwise

        Example:
            entity = await service.find_entity_by_display_name("Bnb Chain", EntityType.BRAND_LOGO)
            if entity:
                print(entity.name)  # "bnb_chain"
        """
        snapshot = await self.load_database_snapshot(use_cache=use_cache)
        search_lower = display_name.lower().strip()

        if category:
            # Search only in specified category
            entities = snapshot.get_entities_by_category(category)
            for entity in entities:
                if entity.display_name.lower().strip() == search_lower:
                    return entity
        else:
            # Search all categories
            for entities in snapshot.entities_by_category.values():
                for entity in entities:
                    if entity.display_name.lower().strip() == search_lower:
                        return entity

        return None

    async def get_entities_by_category(
        self,
        category: EntityType,
        use_cache: bool = True,
    ) -> List[DatabaseEntity]:
        """
        Get all entities for a specific category

        Args:
            category: EntityType to retrieve
            use_cache: Whether to use cached data

        Returns:
            List of DatabaseEntity objects for this category
        """
        snapshot = await self.load_database_snapshot(use_cache=use_cache)
        return snapshot.get_entities_by_category(category)

    async def get_entity_image_paths(
        self,
        entity_name: str,
        category: EntityType = None,
        use_cache: bool = True,
    ) -> List[Path]:
        """
        Get file paths for entity's example images

        Args:
            entity_name: Entity name to search for
            category: Optional category to narrow search
            use_cache: Whether to use cached data

        Returns:
            List of Path objects to example images, empty list if not found
        """
        entity = await self.get_entity_by_name(entity_name, category, use_cache)

        if entity:
            return entity.example_image_paths
        else:
            return []

    async def get_entity_images_as_base64(
        self,
        entity_name: str,
        category: EntityType = None,
        use_cache: bool = True,
    ) -> List[str]:
        """
        Get entity's example images as base64-encoded strings

        Args:
            entity_name: Entity name to search for
            category: Optional category to narrow search
            use_cache: Whether to use cached data

        Returns:
            List of base64-encoded image strings, empty list if not found
        """
        image_paths = await self.get_entity_image_paths(entity_name, category, use_cache)

        if not image_paths:
            return []

        def _load_and_encode_images() -> List[str]:
            """Load images and encode to base64 (blocking I/O)"""
            encoded_images = []

            for image_path in image_paths:
                try:
                    if not image_path.exists() or not image_path.is_file():
                        logger.warning(f"Image file not found: {image_path}")
                        continue

                    with open(image_path, "rb") as f:
                        image_data = f.read()
                        encoded = base64.b64encode(image_data).decode("utf-8")
                        encoded_images.append(encoded)

                except Exception as e:
                    logger.error(f"Failed to load image {image_path}: {e}")
                    continue

            return encoded_images

        try:
            encoded_images = await asyncio.to_thread(_load_and_encode_images)
            logger.debug(
                f"Loaded and encoded {len(encoded_images)} images for entity: {entity_name}"
            )
            return encoded_images

        except Exception as e:
            logger.error(
                f"Failed to load images for entity {entity_name}: {e}"
            )
            return []

    async def refresh_cache(self) -> None:
        """
        Force full refresh of cached data

        Reloads entire database from filesystem, bypassing cache.
        Use sparingly - only when filesystem changes are detected.
        """
        logger.info("🔄 Refreshing entity database cache from filesystem")
        async with self._lock:
            self._snapshot_cache = None

        # Pre-load cache
        await self.load_database_snapshot(use_cache=False)

    async def get_database_stats(self) -> Dict[str, any]:
        """
        Get statistics about the entity database

        Returns:
            Dictionary with statistics:
            - total_entities: Total number of entities
            - total_images: Total number of example images
            - categories: Number of categories with entities
            - entities_by_category: Breakdown per category
        """
        snapshot = await self.load_database_snapshot(use_cache=True)

        return {
            "total_entities": snapshot.total_entity_count,
            "total_images": snapshot.total_image_count,
            "categories": len(snapshot.categories_loaded),
            "entities_by_category": snapshot.get_category_stats(),
            "database_root": str(self.database_root),
        }

    async def has_entity(
        self,
        entity_name: str,
        category: EntityType = None,
        use_cache: bool = True,
    ) -> bool:
        """
        Check if entity exists in database

        Args:
            entity_name: Entity name to check
            category: Optional category to narrow search
            use_cache: Whether to use cached data

        Returns:
            True if entity exists, False otherwise
        """
        entity = await self.get_entity_by_name(entity_name, category, use_cache)
        return entity is not None


# Global service instance (singleton pattern)
_entity_database_service_instance: Optional[EntityDatabaseService] = None


def get_entity_database_service() -> EntityDatabaseService:
    """
    Get the global entity database service instance

    Returns:
        EntityDatabaseService instance (singleton)
    """
    global _entity_database_service_instance

    if _entity_database_service_instance is None:
        _entity_database_service_instance = EntityDatabaseService()

    return _entity_database_service_instance