from pathlib import Path

import dataiku

from solutions.graph.explorer_builder import ExplorerMetadataManager


def do(payload, config, plugin_config, inputs):
    input_db_folder = dataiku.Folder(inputs[0]["fullName"])
    parameter_name = payload["parameterName"]

    if parameter_name == "path_to_db":
        db_paths_choices = []
        for path in input_db_folder.list_paths_in_partition():
            if path.lower().endswith(".kz"):
                rel_path = path.lstrip("/")
                db_paths_choices.append({"value": rel_path, "label": rel_path})

        return {"choices": db_paths_choices}
    elif parameter_name == "node_groups":
        if root_model := payload.get("rootModel", ""):
            if path_to_db := root_model.get("path_to_db", ""):
                # Hack: reconstruct configuration.json path and extract its content.
                # This is part of a time sensitive patch.
                snapshot_id = Path(path_to_db).parts[1]

                snapshot_metadata = ExplorerMetadataManager([input_db_folder])._get_snapshot_data(snapshot_id)
                if snapshot_metadata is None:
                    raise Exception(f"Snapshot metadata for id {snapshot_id} not found.")

                node_groups = snapshot_metadata["node_definitions"]

                choices = [{"label": k["node_group"], "value": k["node_group"]} for k in node_groups]
                return {"choices": choices}

        return {"choices": []}
    elif parameter_name == "edge_groups":
        if root_model := payload.get("rootModel", ""):
            if path_to_db := root_model.get("path_to_db", ""):
                # Hack: reconstruct configuration.json path and extract its content.
                # This is part of a time sensitive patch.
                snapshot_id = Path(path_to_db).parts[1]

                snapshot_metadata = ExplorerMetadataManager([input_db_folder])._get_snapshot_data(snapshot_id)
                if snapshot_metadata is None:
                    raise Exception(f"Snapshot metadata for id {snapshot_id} not found.")

                edge_groups = snapshot_metadata["edge_definitions"]

                choices = [{"label": k["edge_group"], "value": k["edge_group"]} for k in edge_groups]
                return {"choices": choices}

        return {"choices": []}
    else:
        raise Exception(f"Unexpected parameter name '{parameter_name}'.")
