from flask import Blueprint, request
from pydantic import BaseModel

from backends.editor.queries import get_source_datasets_config
from backends.editor.utils.webapp_config import webapp_config
from backends.utils import return_ok

datasets = Blueprint("datasets", __name__, url_prefix="/datasets")


@datasets.route("/get_datasets", methods=["GET"])
def get_datasets():
    config_datasets = get_source_datasets_config()

    return return_ok(
        {"nodes_datasets": config_datasets["nodes_datasets"], "edges_datasets": config_datasets["edges_datasets"]}
    )


class GetDatasetColumnsParams(BaseModel):
    dataset_name: str


@datasets.route("/get_dataset_columns", methods=["POST"])
def get_dataset_columns():
    params = GetDatasetColumnsParams(**request.get_json())

    config_datasets = get_source_datasets_config()
    all_datasets = config_datasets["nodes_datasets"] + config_datasets["edges_datasets"]

    dataset_name = params.dataset_name
    if dataset_name in all_datasets:
        return return_ok({"columns": webapp_config.get_columns(dataset_name)})
    else:
        raise ValueError("invalid dataset name {} ".format(str(dataset_name)))
