import dataiku


def compare_versions(version1: str, version2: str) -> int:
    try:
        v1 = (version1 or "").strip()
        v2 = (version2 or "").strip()

        # Dev versions are always considered greater than any other version
        is_dev1 = v1.startswith("dev/")
        is_dev2 = v2.startswith("dev/")

        if is_dev1 and is_dev2:
            return 0
        if is_dev1:
            return 1
        if is_dev2:
            return -1

        # Strip pre-release identifiers (alpha, beta, rc, etc.)
        def strip_pre_release(v):
            return v.split("-")[0].split("+")[0]

        clean1 = strip_pre_release(v1)
        clean2 = strip_pre_release(v2)

        # Split and convert to numbers
        def to_int_list(s):
            return [int(x) if x.isdigit() else 0 for x in s.split(".")]

        parts1 = to_int_list(clean1)
        parts2 = to_int_list(clean2)

        max_length = max(len(parts1), len(parts2))
        for i in range(max_length):
            num1 = parts1[i] if i < len(parts1) else 0
            num2 = parts2[i] if i < len(parts2) else 0

            if num1 > num2:
                return 1
            if num1 < num2:
                return -1

        return 0
    except Exception as error:
        print("Error comparing versions:", error)
        return 0


def get_dataset_columns(project, dataset_name):
    dataset = project.get_dataset(dataset_name)
    schema = dataset.get_schema()
    return {"choices":[{"value": column['name'], "label": column['name']} for column in schema['columns']]}


def get_data_storage_options(project):
    choices = [
        {"value": "no_storage", "label": "No Storage"},
        {"value": "managed_folder", "label": "Managed Folder"},
    ]
    
    dss_version = dss_version = project.client.get_instance_info().raw.get("dssVersion", "0.0.0")
    if compare_versions(dss_version, "14.0.0") >= 0:
        choices.append({"value": "workload_folder", "label": "Built in (Workload Folder)"})
    return {"choices": choices}


def do(payload, config, plugin_config, inputs):
    parameter_name = payload["parameterName"]
    client = dataiku.api_client()
    project = client.get_default_project()

    if parameter_name == "llm_responses_column":
        dataset_name = config.get("llm_responses_dataset")
        return get_dataset_columns(project, dataset_name)
    elif parameter_name == "data_storage":
        return get_data_storage_options(project)
