from __future__ import annotations

import logging
from threading import RLock
from typing import List, Tuple

from cachetools import LRUCache
from dataiku import Folder

from backends.explorer.connections.db_config_scanner import DbConfigScanner
from backends.explorer.connections.kuzu.db_handler_factory import KuzuDbHandlerFactory
from backends.explorer.connections.neo4j.db_handler_factory import Neo4jDbHandlerFactory
from backends.explorer.utils.presets import get_webapp_neo4j_preset
from graph.connections.db_handler import DbHandler
from solutions.graph.models import PresetNotSetException

logger = logging.getLogger(__name__)


class CleanupLRUCache(LRUCache):
    """LRU Cache that calls cleanup_handler() on evicted database handlers.

    This ensures proper resource cleanup (closing connections, releasing file handles)
    when handlers are evicted from the cache due to size limits.
    """

    def popitem(self) -> Tuple[str, DbHandler]:
        """Override to cleanup handler before removing from cache."""
        snapshot_id, handler = super().popitem()
        try:
            handler.cleanup_handler()
            logger.info(f"Cleaned up evicted handler for graph {snapshot_id}")
        except Exception as e:
            logger.exception(f"Error cleaning up handler for graph {snapshot_id}")

        return snapshot_id, handler

class ExplorerDbHandlerRegistry:
    """Centralized registry for managing database handlers with LRU caching.

    Uses a custom LRU cache that properly cleans up database handlers (closes connections,
    releases resources) when they are evicted from the cache.
    """
    def __init__(self, db_folders: List[Folder], cache_size: int = 128):
        self.db_config_scanner = DbConfigScanner(db_folders)

        self.cache_size = cache_size
        self._cache = CleanupLRUCache(maxsize=cache_size)
        self._cache_lock = RLock()

    def get_handler(self, snapshot_id: str) -> DbHandler:
        """Get or create a database handler for a given graph."""
        if not snapshot_id:
            raise ValueError("snapshot_id cannot be empty")

        with self._cache_lock:
            # Check if handler exists in cache
            if snapshot_id in self._cache:
                return self._cache[snapshot_id]  # type: ignore
            
            # Create handler only if not in cache
            handler = self._create_handler(snapshot_id)
            self._cache[snapshot_id] = handler
            return handler
    
    def _create_handler(self, snapshot_id: str) -> DbHandler:
        """Create a new database handler."""
        logger.debug(f"Creating handler for graph {snapshot_id}")
        db_type = self.db_config_scanner.get_db_type(snapshot_id)

        if db_type == "kuzu":
            return self._create_kuzu_handler(snapshot_id)
        elif db_type == "neo4j":
            return self._create_neo4j_handler(snapshot_id)
        else:
            raise ValueError(f"Unknown database type for snapshot {snapshot_id}")

    def _create_kuzu_handler(self, snapshot_id: str) -> DbHandler:
        """Create a Kuzu database handler for a snapshot."""
        db_folder = self.db_config_scanner.get_db_folder_by_snapshot_id(snapshot_id)
        return KuzuDbHandlerFactory(db_folder, snapshot_id).get_db_handler()

    def _create_neo4j_handler(self, snapshot_id: str) -> DbHandler:
        """Create a Neo4j database handler for a snapshot."""
        db_folder = self.db_config_scanner.get_db_folder_by_snapshot_id(snapshot_id)
        preset = get_webapp_neo4j_preset()
        if not preset:
            raise PresetNotSetException()

        return Neo4jDbHandlerFactory(db_folder, snapshot_id, preset).get_db_handler()

    def clear_cache(self) -> None:
        """Clear all cached handlers and cleanup their resources."""
        with self._cache_lock:
            for snapshot_id, handler in list(self._cache.items()):
                try:
                    handler.cleanup_handler()
                    logger.debug(f"Cleaned up handler for graph {snapshot_id}")
                except Exception as e:
                    logger.exception(f"Error cleaning up handler for graph {snapshot_id}: {e}")
            self._cache.clear()
            logger.info("DB handler cache cleared")

    def get_cache_info(self) -> dict:
        """Get information about the current cache state."""
        with self._cache_lock:
            return {
                "maxsize": self._cache.maxsize,
                "currsize": self._cache.currsize,
            }

_explorer_registry_instance: ExplorerDbHandlerRegistry | None = None
_registry_init_lock = RLock()

def get_db_handler_registry(db_folders: List[Folder]) -> ExplorerDbHandlerRegistry:
    global _explorer_registry_instance
    if _explorer_registry_instance:
        return _explorer_registry_instance

    with _registry_init_lock:
        if _explorer_registry_instance:
            return _explorer_registry_instance

        _explorer_registry_instance = ExplorerDbHandlerRegistry(db_folders)
        return _explorer_registry_instance


def get_handler_for_snapshot(snapshot_id: str, db_folders: List[Folder]) -> DbHandler:
    logger.debug("Cache info: %s", get_db_handler_registry(db_folders).get_cache_info())
    return get_db_handler_registry(db_folders).get_handler(snapshot_id)
