import certifi
import requests
from environs import Env
from requests.auth import HTTPDigestAuth
from tenacity import retry, stop_after_attempt, wait_exponential

env = Env()
env.read_env(".env")
PREPROCESS_HOST = env.str("PREPROCESS_HOST")
PREPROCESS_KEY = env.str("PREPROCESS_KEY")
COMPONENT_TYPES = ["TIDB", "TIKV", "TIFLASH"]


def set_to_list(obj):
    if isinstance(obj, dict):
        return {key: set_to_list(value) for key, value in obj.items()}
    elif isinstance(obj, set):
        return list(obj)
    else:
        return obj


@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def get_region_ids():
    url = f"https://{PREPROCESS_HOST}/v1beta1/regions"
    public_key, private_key = PREPROCESS_KEY.split(":")
    auth = HTTPDigestAuth(public_key, private_key)
    headers = {"Content-Type": "application/json"}
    response = requests.get(url, auth=auth, headers=headers, verify=certifi.where())
    response.raise_for_status()
    return (
        [region["regionId"] for region in response.json().get("regions")]
        if response.json().get("regions")
        else []
    )


@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def get_node_specs(region_id):
    public_key, private_key = PREPROCESS_KEY.split(":")
    auth = HTTPDigestAuth(public_key, private_key)
    headers = {"Content-Type": "application/json"}
    url = (
        f"https://{PREPROCESS_HOST}/v1beta1/regions/{region_id}/nodeSpecs?pageSize=100"
    )
    response = requests.get(url, auth=auth, headers=headers, verify=certifi.where())
    return response.json().get("nodeSpecs") if response.json().get("nodeSpecs") else []


def merge_node_specs():
    merged_node_specs = {
        "TIDB": {"nodeSpecKey": set()},
        "TIKV": {
            "nodeSpecKey": set(),
            "maxStorageSizeGi": 0,
            "minStorageSizeGi": 1e8,
            "storageTypes": set(),
        },
        "TIFLASH": {
            "nodeSpecKey": set(),
            "maxStorageSizeGi": 0,
            "minStorageSizeGi": 1e8,
            "storageTypes": set(),
        },
    }
    try:
        region_ids = get_region_ids()
    except requests.exceptions.RequestException as e:
        print(f"Error: {e}")

    for region_id in region_ids:
        try:
            node_specs = get_node_specs(region_id)
        except requests.exceptions.RequestException as e:
            print(f"Error: {e}")
        for node_spec in node_specs:
            if node_spec["componentType"] == "TIDB":
                merged_node_specs["TIDB"]["nodeSpecKey"].add(node_spec["nodeSpecKey"])
            elif node_spec["componentType"] == "TIKV":
                merged_node_specs["TIKV"]["nodeSpecKey"].add(node_spec["nodeSpecKey"])
                merged_node_specs["TIKV"]["maxStorageSizeGi"] = max(
                    merged_node_specs["TIKV"]["maxStorageSizeGi"],
                    node_spec["maxStorageSizeGi"],
                )
                merged_node_specs["TIKV"]["minStorageSizeGi"] = min(
                    merged_node_specs["TIKV"]["minStorageSizeGi"],
                    node_spec["minStorageSizeGi"],
                )
                merged_node_specs["TIKV"]["storageTypes"] |= set(
                    node_spec["storageTypes"]
                )
            elif node_spec["componentType"] == "TIFLASH":
                merged_node_specs["TIFLASH"]["nodeSpecKey"].add(
                    node_spec["nodeSpecKey"]
                )
                merged_node_specs["TIFLASH"]["maxStorageSizeGi"] = max(
                    merged_node_specs["TIFLASH"]["maxStorageSizeGi"],
                    node_spec["maxStorageSizeGi"],
                )
                merged_node_specs["TIFLASH"]["minStorageSizeGi"] = min(
                    merged_node_specs["TIFLASH"]["minStorageSizeGi"],
                    node_spec["minStorageSizeGi"],
                )
                merged_node_specs["TIFLASH"]["storageTypes"] |= set(
                    node_spec["storageTypes"]
                )
    storage_size_min = min(
        merged_node_specs["TIKV"]["minStorageSizeGi"],
        merged_node_specs["TIFLASH"]["minStorageSizeGi"],
    )
    storage_size_max = max(
        merged_node_specs["TIKV"]["maxStorageSizeGi"],
        merged_node_specs["TIFLASH"]["maxStorageSizeGi"],
    )
    tidb_speckey = list(merged_node_specs["TIDB"]["nodeSpecKey"])
    storage_sepckey = (
        merged_node_specs["TIKV"]["nodeSpecKey"]
        | merged_node_specs["TIFLASH"]["nodeSpecKey"]
    )
    # return set_to_list(merged_node_specs)
    return {
        "TIDB": {"nodeSpecKey": list(tidb_speckey)},
        "STORAGE": {
            "storageSizeGi": {"min": storage_size_min, "max": storage_size_max},
            "nodeSpecKey": list(storage_sepckey),
        },
    }


if __name__ == "__main__":
    specs = merge_node_specs()
    print(specs)