from __future__ import annotations

import logging

from dataiku import Folder

from recipes.build_graph.assets import BUILT_GRAPHS_PREFIX, DATABASE_JSON_FILENAME, DB_KUZU_FILENAME
from solutions.graph.models import DbType

logger = logging.getLogger(__name__)


class DbConfigScanner:
    """Scanner for database configurations stored in managed folders."""
    def __init__(self, db_folders: list[Folder]):
        self.db_folders = db_folders
        self._snapshot_cache: dict[str, dict] = {}  # {snapshot_id: {folder, db_type}}
        self._cache_initialized = False
    
    def _initialize_cache(self):
        """Scans all folders once and builds a cache of snapshot information."""
        if self._cache_initialized:
            return
        
        for folder in self.db_folders:
            remote_paths = folder.list_paths_in_partition()
            
            for path in remote_paths:
                normalized_path = path.lstrip("/")
                
                # Check if path ends with database files (path-aware check)
                if (normalized_path.endswith(f"/{DB_KUZU_FILENAME}") or 
                    normalized_path.endswith(f"/{DATABASE_JSON_FILENAME}")):
                    try:
                        snapshot_id = self._extract_snapshot_id_from_path(normalized_path)
                        
                        if snapshot_id not in self._snapshot_cache:
                            # Read db_type from database.json if available
                            db_type = self._read_db_type_from_folder(folder, snapshot_id)
                            db_name = self._read_db_name_from_folder(folder, snapshot_id)
                            self._snapshot_cache[snapshot_id] = {
                                "folder": folder,
                                "db_type": db_type,
                                "db_name": db_name,
                            }
                        elif self._snapshot_cache[snapshot_id]["folder"] != folder:
                            logger.warning(
                                f"Multiple folders contain snapshot {snapshot_id}. Using the first one found."
                            )
                    except Exception as e:
                        logger.debug(f"Could not extract snapshot ID from path {path}: {e}")

        self._cache_initialized = True
    
    def _read_db_type_from_folder(self, folder: Folder, snapshot_id: str) -> DbType:
        """Reads the database type from database.json."""
        database_json_path = f"/{BUILT_GRAPHS_PREFIX}/{snapshot_id}/{DATABASE_JSON_FILENAME}"
        try:
            database_json = folder.read_json(database_json_path)
            database_type = database_json.get("type")

            if database_type == "kuzu":
                return "kuzu"
            elif database_type == "neo4j":
                return "neo4j"
            else:
                raise ValueError(
                    f"Unknown database type '{database_type}' for snapshot {snapshot_id}."
                )
        except Exception as e:
            # Exception raised if the database.json file does not exist
            logger.warning(f"Error reading database.json for snapshot {snapshot_id}: {e}. Defaulting to kuzu.")
            return "kuzu"
    
    def _read_db_name_from_folder(self, folder: Folder, snapshot_id: str) -> str | None:
        """Reads the database name from database.json."""
        database_json_path = f"/{BUILT_GRAPHS_PREFIX}/{snapshot_id}/{DATABASE_JSON_FILENAME}"
        try:
            database_json = folder.read_json(database_json_path)
            database_params = database_json.get("params")

            if not database_params:
                return None
            else:
                name: str = database_params["databaseName"]
                return name
        except Exception as e:
            # Exception raised if the database.json file does not exist
            logger.warning(f"Error reading database.json for snapshot {snapshot_id}: {e}.")
            return None

    def get_db_type(self, snapshot_id: str) -> DbType:
        """
        Determines the database type (Kuzu or Neo4j) for a given snapshot ID.

        Args:
            snapshot_id (str): The unique identifier of the snapshot.

        Returns:
            DbType: The type of the database, either "kuzu" or "neo4j".

        Raises:
            Exception: If no folder contains the snapshot.
        """
        self._initialize_cache()
        
        if snapshot_id not in self._snapshot_cache:
            raise Exception(f"No folder contains snapshot ID: {snapshot_id}")
        
        db_type: DbType = self._snapshot_cache[snapshot_id]["db_type"]
        return db_type

    def get_db_name(self, snapshot_id: str) -> str | None:
        self._initialize_cache()

        if snapshot_id not in self._snapshot_cache:
            raise Exception(f"No folder contains snapshot ID: {snapshot_id}")

        db_name: str | None = self._snapshot_cache[snapshot_id]["db_name"]
        return db_name

    def get_db_folder_by_snapshot_id(self, snapshot_id: str) -> Folder:
        """
        Finds the folder containing the database files for a given snapshot ID.

        Args:
            snapshot_id (str): The unique identifier of the snapshot.

        Returns:
            Folder: The managed folder containing the snapshot.

        Raises:
            Exception: If no folder contains the snapshot.
        """
        self._initialize_cache()
        
        if snapshot_id not in self._snapshot_cache:
            raise Exception(f"No folder contains snapshot ID: {snapshot_id}")
        
        return self._snapshot_cache[snapshot_id]["folder"]
    
    def get_snapshot_ids(self) -> list[str]:
        """
        Retrieves all snapshot IDs from all managed folders.

        Returns:
            list[str]: A list of unique snapshot IDs.
        """
        self._initialize_cache()
        return list(self._snapshot_cache.keys())

    @staticmethod
    def _extract_snapshot_id_from_path(path: str) -> str:
        """
        Extracts the snapshot ID from a path string.

        Expected path format: built-graphs/{snapshot_id}/db.kz or built-graphs/{snapshot_id}/database.json

        Args:
            path (str): The file path from which to extract the snapshot ID.

        Returns:
            str: The extracted snapshot ID.

        Raises:
            ValueError: If the path format is invalid.
        """
        # Normalize path by stripping leading slash
        normalized_path = path.lstrip("/")
        parts = normalized_path.split("/")
        
        # Find the snapshot ID between "built-graphs" and the database file
        try:
            built_graphs_index = parts.index(BUILT_GRAPHS_PREFIX)
            # The snapshot ID should be right after "built-graphs"
            if built_graphs_index + 1 < len(parts):
                return parts[built_graphs_index + 1]
        except (ValueError, IndexError):
            # If the expected structure is not present, fall through and raise a ValueError below.
            pass
        
        raise ValueError(f"Invalid path format: {path}. Expected format: built-graphs/{{snapshot_id}}/{{db_file}}")
