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 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")
            if input_db_folder is not None:
                snapshot_metadata = get_metadata_from_folder(input_db_folder, neo4j_configuration)
            elif database_name and neo4j_configuration:
                snapshot_metadata = get_metadata(
                    db_folder=None,
                    neo4j_configuration=neo4j_configuration,
                    database_name=database_name,
                    database_type="neo4j",
                )
            else:
                return result

            node_groups = snapshot_metadata["node_definitions"]
            unique_groups = sorted(list(set(k["node_group"] for k in node_groups)))
            choices = [{"label": group, "value": group} for group in unique_groups]
            result = {"choices": choices}
    elif parameter_name == "edge_groups":
        if root_model:
            database_name = root_model.get("database_name")
            if input_db_folder is not None:
                snapshot_metadata = get_metadata_from_folder(input_db_folder, neo4j_configuration)
            elif database_name and neo4j_configuration:
                snapshot_metadata = get_metadata(
                    db_folder=None,
                    neo4j_configuration=neo4j_configuration,
                    database_name=database_name,
                    database_type="neo4j",
                )
            else:
                return result

            edge_groups = snapshot_metadata["edge_definitions"]
            unique_groups = sorted(list(set(k["edge_group"] for k in edge_groups)))
            choices = [{"label": group, "value": group} for group in unique_groups]
            result = {"choices": choices}
    else:
        raise Exception(f"Unexpected parameter name '{parameter_name}'.")

    return result
