from __future__ import annotations

import logging

from dataiku import Folder

from backends.explorer.discovery.connection_info import (
    AbstractConnectionInfo,
    get_connection_info_from_folder,
)
from graph.connections.neo4j.preset import Neo4jPresetConfig
from recipes.build_graph.assets import BUILT_GRAPHS_PREFIX, DATABASE_JSON_FILENAME, DB_KUZU_FILENAME

logger = logging.getLogger(__name__)


class ConnectionInfoFolderScanner:
    """Scans DSS managed folders for graph database snapshots.

    Attributes:
        db_folders (list[Folder]): List of DSS folders to scan.
        neo4j_preset (Optional[Neo4jPresetConfig]): Credentials for Neo4j connections.
    """

    def __init__(self, db_folders: list[Folder], neo4j_preset: Neo4jPresetConfig | None = None):
        """Initializes the scanner.

        Args:
            db_folders (list[Folder]): List of DSS folders to scan.
            neo4j_preset (Optional[Neo4jPresetConfig]): Credentials for Neo4j connections.
        """
        self.db_folders = db_folders
        self.neo4j_preset = neo4j_preset

        self._snapshot_cache: dict[str, AbstractConnectionInfo] = {}
        self._cache_initialized = False

    def _initialize_cache(self):
        """Discovers all valid snapshots in the configured folders."""
        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:
                            self._snapshot_cache[snapshot_id] = get_connection_info_from_folder(
                                folder, snapshot_id, self.neo4j_preset
                            )

                    except Exception as e:
                        logger.debug(f"Could not extract snapshot ID from path {path}: {e}")

        self._cache_initialized = True

    def get_connection_infos(self) -> dict[str, AbstractConnectionInfo]:
        """Returns the mapping of snapshot IDs to connection information.

        Returns:
            dict[str, AbstractConnectionInfo]: The discovered connections.
        """
        self._initialize_cache()
        return self._snapshot_cache

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

        Args:
            path (str): The file path within the folder.

        Returns:
            str: The snapshot ID.

        Raises:
            ValueError: If the path does not match the expected 'built-graphs/ID/...' format.
        """
        # 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):
            pass

        raise ValueError(f"Invalid path format: {path}. Expected format: built-graphs/{{snapshot_id}}/{{db_file}}")
