from __future__ import annotations

from dataiku import Folder

from backends.explorer.discovery.connection_info import Neo4jConnectionInfo
from backends.explorer.discovery.scanner import ConnectionInfoFolderScanner
from graph.connections.connection_info import AbstractConnectionInfo
from graph.connections.neo4j.preset import Neo4jPresetConfig
from solutions.graph.models import DbType


class ConnectionInfoRegistry:
    """A central store for all available graph database connections.

    Supports connections discovered in DSS folders and those explicitly
    configured as external databases.
    """

    def __init__(self, neo4j_preset: Neo4jPresetConfig | None = None):
        """Initializes the registry.

        Args:
            neo4j_preset (Optional[Neo4jPresetConfig]): Default preset to use for Neo4j connections.
        """
        self.connection_infos: dict[str, AbstractConnectionInfo] = {}
        self.neo4j_preset: Neo4jPresetConfig | None = neo4j_preset

    def load_from_dataiku_folders(self, db_folders: list[Folder]):
        """Scans DSS folders and populates the registry.

        Args:
            db_folders (list[Folder]): The DSS folders containing graph data.
        """
        scanner = ConnectionInfoFolderScanner(db_folders, self.neo4j_preset)
        conn_infos = scanner.get_connection_infos()

        for uid, conn_info in conn_infos.items():
            self.connection_infos[uid] = conn_info

    def load_from_external_databases(self, db_names: list[str], db_type: DbType):
        """Adds explicitly configured external databases to the registry.

        Args:
            db_names (list[str]): Names of the databases.
            db_type (DbType): The type (currently only 'neo4j' supported for external).

        Raises:
            ValueError: If the database type is not supported for external connections.
        """
        for db_name in db_names:
            if db_type == "neo4j":
                if self.neo4j_preset:
                    self.connection_infos[db_name] = Neo4jConnectionInfo(db_name, db_name, self.neo4j_preset)
                else:
                    raise ValueError("Neo4j preset is required to load external databases.")
            else:
                raise ValueError(f"Unknown database type: {db_type}")

    def get_connection_info(self, uid: str) -> AbstractConnectionInfo | None:
        """Retrieves connection info for a given database name or snapshot ID.

        Args:
            uid (str): The identifier to look up.

        Returns:
            Optional[AbstractConnectionInfo]: The info object if found, else None.
        """
        return self.connection_infos.get(uid)
    
    def list_connection_info(self) -> list[AbstractConnectionInfo]:
        """Returns a list of all registered connection information objects.

        Returns:
            list[AbstractConnectionInfo]: All connection infos in the registry.
        """
        return [v for _, v in self.connection_infos.items()]


def get_connection_info_registry(
    db_folders: list[Folder] | None = None,
    neo4j_databases: list[str] | None = None,
    neo4j_preset: Neo4jPresetConfig | None = None
) -> ConnectionInfoRegistry:
    """Factory function to create and populate a ConnectionInfoRegistry.

    Args:
        db_folders (Optional[list[Folder]]): Folders to scan for connections.
        neo4j_databases (Optional[list[str]]): External Neo4j database names.
        neo4j_preset (Optional[Neo4jPresetConfig]): Preset for Neo4j connections.

    Returns:
        ConnectionInfoRegistry: The populated registry.
    """
    registry = ConnectionInfoRegistry(neo4j_preset)

    if db_folders:
        registry.load_from_dataiku_folders(db_folders)
    if neo4j_databases:
        registry.load_from_external_databases(neo4j_databases, "neo4j")

    return registry
