from __future__ import annotations

from dataiku import Folder

from backends.explorer.discovery.connection_info import (
    AbstractConnectionInfo,
    Neo4jConnectionInfo,
    get_database_type_from_snapshot_folder,
)
from backends.explorer.discovery.info_registry import get_connection_info_registry
from backends.explorer.discovery.scanner import ConnectionInfoFolderScanner
from backends.explorer.factories.db_handler_factory import DbHandlerFactory
from backends.explorer.metadata.metadata_manager import ExplorerMetadataManager
from graph.connections.db_handler import DbHandler
from graph.connections.neo4j.preset import Neo4jPresetConfig
from recipes.build_graph.assets import DATABASE_JSON_FILENAME, DB_KUZU_FILENAME
from solutions.graph.models import DbType, ExplorerMetadata


def get_connection_info_registry_from_inputs(
    db_folder: Folder | None = None,
    neo4j_configuration: Neo4jPresetConfig | None = None,
    database_name: str | None = None,
    database_type: DbType | None = None,
):
    if db_folder is not None:
        return get_connection_info_registry(db_folders=[db_folder], neo4j_preset=neo4j_configuration)

    if database_name:
        if database_type and database_type != "neo4j":
            raise ValueError("External databases without folder are only supported for Neo4j.")
        if not neo4j_configuration:
            raise ValueError("Neo4j configuration is required when selecting a database without folder.")
        return get_connection_info_registry(
            db_folders=None,
            neo4j_databases=[database_name],
            neo4j_preset=neo4j_configuration,
        )

    return get_connection_info_registry(db_folders=None, neo4j_databases=None, neo4j_preset=neo4j_configuration)


def _get_single_connection_info_from_registry(
    conn_info_registry,
    db_folder: Folder | None = None,
) -> AbstractConnectionInfo:
    conn_info_list: list[AbstractConnectionInfo] = conn_info_registry.list_connection_info()

    if not conn_info_list:
        if db_folder is not None:
            raise ValueError("No graph database found in the input folder.")
        raise ValueError("No graph database found for the selected configuration.")
    if len(conn_info_list) > 1:
        if db_folder is not None:
            raise ValueError("Multiple graph databases found in the input folder, please select a folder with a single graph.")
        raise ValueError("Multiple graph databases found for the selected configuration.")

    return conn_info_list[0]


def get_single_connection_info(
    db_folder: Folder | None = None,
    neo4j_configuration: Neo4jPresetConfig | None = None,
    database_name: str | None = None,
    database_type: DbType | None = None,
) -> AbstractConnectionInfo:
    conn_info_registry = get_connection_info_registry_from_inputs(
        db_folder=db_folder,
        neo4j_configuration=neo4j_configuration,
        database_name=database_name,
        database_type=database_type,
    )
    return _get_single_connection_info_from_registry(conn_info_registry, db_folder=db_folder)


def get_db_handler(
    db_folder: Folder | None = None,
    neo4j_configuration: Neo4jPresetConfig | None = None,
    database_name: str | None = None,
    database_type: DbType | None = None,
    snapshot_id: str | None = None,
) -> DbHandler:
    conn_info_registry = get_connection_info_registry_from_inputs(
        db_folder=db_folder,
        neo4j_configuration=neo4j_configuration,
        database_name=database_name,
        database_type=database_type,
    )
    conn_info: AbstractConnectionInfo | None
    if snapshot_id:
        conn_info = conn_info_registry.get_connection_info(snapshot_id)
        if not conn_info:
            raise ValueError(f"No connection info found for snapshot {snapshot_id}")
    else:
        conn_info = _get_single_connection_info_from_registry(conn_info_registry, db_folder=db_folder)

    if conn_info.type == "kuzu":
        return DbHandlerFactory.get_db_handler(conn_info)
    if conn_info.type == "neo4j":
        if not neo4j_configuration:
            raise ValueError("No Neo4j configuration provided")
        return DbHandlerFactory.get_db_handler(conn_info, neo4j_configuration)
    raise ValueError(f"Unknown database type '{conn_info.type}' for snapshot {conn_info.uid}")


def get_db_handler_from_snapshot(
    input_db_folder: Folder | None,
    snapshot_id: str,
    neo4j_configuration: Neo4jPresetConfig | None = None,
    database_name: str | None = None,
    database_type: DbType | None = None,
) -> DbHandler:
    return get_db_handler(
        db_folder=input_db_folder,
        neo4j_configuration=neo4j_configuration,
        database_name=database_name,
        database_type=database_type,
        snapshot_id=snapshot_id,
    )


def get_metadata(
    db_folder: Folder | None = None,
    neo4j_configuration: Neo4jPresetConfig | None = None,
    database_name: str | None = None,
    database_type: DbType | None = None,
) -> ExplorerMetadata:
    conn_info_registry = get_connection_info_registry_from_inputs(
        db_folder=db_folder,
        neo4j_configuration=neo4j_configuration,
        database_name=database_name,
        database_type=database_type,
    )
    conn_info = _get_single_connection_info_from_registry(conn_info_registry, db_folder=db_folder)
    explorer_metadata = ExplorerMetadataManager(conn_info_registry)
    db_handler = get_db_handler(
        db_folder=db_folder,
        neo4j_configuration=neo4j_configuration,
        database_name=database_name,
        database_type=database_type,
        snapshot_id=conn_info.uid,
    )
    try:
        return explorer_metadata.get_snapshot_metadata(conn_info.uid, db_handler)
    finally:
        db_handler.cleanup_handler()


def get_metadata_from_folder(db_folder: Folder, neo4j_configuration: Neo4jPresetConfig | None = None) -> ExplorerMetadata:
    return get_metadata(db_folder=db_folder, neo4j_configuration=neo4j_configuration)


def get_metadata_from_database_name(
    database_name: str,
    neo4j_configuration: Neo4jPresetConfig,
) -> ExplorerMetadata:
    return get_metadata(
        db_folder=None,
        neo4j_configuration=neo4j_configuration,
        database_name=database_name,
        database_type="neo4j",
    )

def get_database_name_from_folder(
    db_folder: Folder,
    neo4j_configuration: Neo4jPresetConfig | None = None,
) -> str:
    conn_info = get_single_connection_info(db_folder=db_folder, neo4j_configuration=neo4j_configuration)

    if conn_info.type == "kuzu":
        return conn_info.uid
    if conn_info.type == "neo4j":
        db_info: Neo4jConnectionInfo = conn_info  # type: ignore[assignment]
        name = db_info.database_name
        if not name:
            raise ValueError(f"No database name found for Neo4j snapshot with name {name}.")
        return name
    raise ValueError(f"Unknown database type '{conn_info.type}' for snapshot {conn_info.uid}.")


def _get_single_snapshot_id_from_folder(db_folder: Folder) -> str:
    snapshot_ids: set[str] = set()
    for path in db_folder.list_paths_in_partition():
        normalized_path = path.lstrip("/")
        if normalized_path.endswith(f"/{DB_KUZU_FILENAME}") or normalized_path.endswith(f"/{DATABASE_JSON_FILENAME}"):
            snapshot_ids.add(ConnectionInfoFolderScanner._extract_snapshot_id_from_path(normalized_path))

    if not snapshot_ids:
        raise ValueError("No graph database found in the input folder.")
    if len(snapshot_ids) > 1:
        raise ValueError("Multiple graph databases found in the input folder, please select a folder with a single graph.")
    return next(iter(snapshot_ids))


def get_database_type_from_folder(
    db_folder: Folder,
) -> DbType:
    snapshot_id = _get_single_snapshot_id_from_folder(db_folder)
    return get_database_type_from_snapshot_folder(db_folder, snapshot_id)
