import dataiku

from graph.connections.neo4j.preset import Neo4jPresetConfig
from graph.exploration.queries.neo4j.list_databases import list_databases as list_neo4j_databases
from recipes.explore_graph.utils import (
    get_database_name_from_folder,
    get_database_type_from_folder,
    get_metadata,
    get_metadata_from_folder,
)


def _choices_from_metadata(snapshot_metadata, definition_key, group_key):
    groups = snapshot_metadata[definition_key]
    unique_groups = sorted({definition[group_key] for definition in groups})
    return {"choices": [{"label": group, "value": group} for group in unique_groups]}


def _get_snapshot_metadata_from_inputs(input_db_folder, neo4j_configuration, database_name):
    if input_db_folder is not None:
        return get_metadata_from_folder(input_db_folder, neo4j_configuration)
    if database_name and neo4j_configuration:
        return get_metadata(
            db_folder=None,
            neo4j_configuration=neo4j_configuration,
            database_name=database_name,
            database_type="neo4j",
        )
    return None


def _intersect_edge_properties(snapshot_metadata, selected_edge_groups):
    property_sets = [
        set(edge_definition.get("property_list", []))
        for edge_definition in snapshot_metadata["edge_definitions"]
        if edge_definition["edge_group"] in selected_edge_groups
    ]
    if not property_sets:
        return []
    return sorted(set.intersection(*property_sets))


def do(payload, config, plugin_config, inputs):
    parameter_name = payload["parameterName"]
    root_model = payload.get("rootModel", {})
    input_db_folder = dataiku.Folder(inputs[0]["fullName"]) if inputs else None

    neo4j_configuration = config.get("neo4j_preset")
    if neo4j_configuration:
        neo4j_configuration = Neo4jPresetConfig.from_dict(neo4j_configuration)

    result = {"choices": []}
    if parameter_name == "database_name":
        if input_db_folder is not None:
            database_name = get_database_name_from_folder(input_db_folder, neo4j_configuration)
            result = {"choices": [{"value": database_name, "label": database_name}]}
        elif neo4j_configuration:
            database_names = list_neo4j_databases(
                uri=neo4j_configuration.neo4j_uri,
                username=neo4j_configuration.neo4j_username,
                password=neo4j_configuration.get_neo4j_password(),
            )
            result = {"choices": [{"value": db_name, "label": db_name} for db_name in sorted(database_names)]}
    elif parameter_name == "database_type":
        if input_db_folder is not None:
            database_type = get_database_type_from_folder(input_db_folder)
            if not database_type:
                result = {"choices": []}
            else:
                label = "Built-in" if database_type == "kuzu" else "Neo4j"
                result = {"choices": [{"value": database_type, "label": label}]}
        else:
            # When no input folder is selected, only direct Neo4j access is available.
            result = {"choices": [{"value": "unmanaged_neo4j", "label": "Neo4j (unmanaged)"}]}
    elif parameter_name == "node_groups":
        if root_model:
            database_name = root_model.get("database_name")
            snapshot_metadata = _get_snapshot_metadata_from_inputs(
                input_db_folder,
                neo4j_configuration,
                database_name,
            )
            if snapshot_metadata is None:
                return result

            result = _choices_from_metadata(snapshot_metadata, "node_definitions", "node_group")
        else:
            return {"choices": []}
    elif parameter_name == "edge_groups":
        if root_model:
            database_name = root_model.get("database_name")
            snapshot_metadata = _get_snapshot_metadata_from_inputs(
                input_db_folder,
                neo4j_configuration,
                database_name,
            )
            if snapshot_metadata is None:
                return result

            result = _choices_from_metadata(snapshot_metadata, "edge_definitions", "edge_group")
        else:
            return {"choices": []}
    elif parameter_name == "weight_property":
        database_name = root_model.get("database_name")
        selected_edge_groups = set(root_model.get("edge_groups", []))
        snapshot_metadata = _get_snapshot_metadata_from_inputs(
            input_db_folder,
            neo4j_configuration,
            database_name,
        )
        if snapshot_metadata is None or not selected_edge_groups:
            return {"choices": [{"label": "No weight property", "value": ""}]}

        properties = _intersect_edge_properties(snapshot_metadata, selected_edge_groups)
        result = {
            "choices": [{"label": "No weight property", "value": ""}]
            + [{"label": prop, "value": prop} for prop in properties]
        }
    else:
        raise Exception(f"Unexpected parameter name '{parameter_name}'.")

    return result
