import dataiku

from graph.connections.neo4j.preset import Neo4jPresetConfig
from recipes.explore_graph.utils import (
    get_database_name_from_folder,
    get_database_type_from_folder,
    get_metadata_from_folder,
)


def do(payload, config, plugin_config, inputs):
    input_db_folder = dataiku.Folder(inputs[0]["fullName"])
    parameter_name = payload["parameterName"]
    
    if parameter_name == "database_name":
        database_name = get_database_name_from_folder(input_db_folder)
        return {"choices": [{"value": database_name, "label": database_name}]}
    elif parameter_name == "database_type":
        database_type = get_database_type_from_folder(input_db_folder)
        if not database_type:
            return {"choices": []}

        label = "Built-in" if database_type == "kuzu" else "Neo4j"

        return {"choices": [{"value": database_type, "label": label}]}
    elif parameter_name == "node_groups":
        if payload.get("rootModel", ""):
            neo4j_configuration = config.get("neo4j_preset")
            if neo4j_configuration:
                neo4j_configuration = Neo4jPresetConfig(**neo4j_configuration)

            snapshot_metadata = get_metadata_from_folder(input_db_folder, neo4j_configuration)
            if snapshot_metadata is None:
                raise Exception(f"Snapshot metadata for {input_db_folder} not found.")

            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]
            return {"choices": choices}

        return {"choices": []}
    elif parameter_name == "edge_groups":
        if payload.get("rootModel", ""):
            neo4j_configuration = config.get("neo4j_preset")
            if neo4j_configuration:
                neo4j_configuration = Neo4jPresetConfig(**neo4j_configuration)

            snapshot_metadata = get_metadata_from_folder(input_db_folder, neo4j_configuration)
            if snapshot_metadata is None:
                raise Exception(f"Snapshot metadata for {input_db_folder} not found.")

            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]
            return {"choices": choices}
    else:
        raise Exception(f"Unexpected parameter name '{parameter_name}'.")
