import json
import os
import pathlib
import shutil
import tarfile

import boto3
import zstandard

DIAG_DATA_PATH = "/tmp/diag_data"

S3_BUCKET_PACK = os.environ["S3_BUCKET_PACK"]


def get_component_config(component_type, diag_data_path):
    cluster_config_file = os.path.join(diag_data_path, "cluster.json")
    if not os.path.exists(cluster_config_file):
        print("file '%s' not exists" % cluster_config_file)
        raise Exception("file '%s' not exists" % cluster_config_file)
    try:
        with open(cluster_config_file, "r") as fr:
            cluster_info = json.loads(fr.read())
            cluster_topology = cluster_info["topology"]
            if component_type in cluster_topology and len(cluster_topology[component_type]) > 0:
                host = cluster_topology[component_type][0]["host"]
                if cluster_info["deploy_type"] == "tiup-cluster":
                    deploy_dir = cluster_topology[component_type][0]["attributes"]["deploy_dir"]
                elif cluster_info["deploy_type"] == "tidb-operator":
                    deploy_dir = cluster_topology[component_type][0]["attributes"]["pod"]
            else:
                host = None
                deploy_dir = None
        if host is not None and deploy_dir is not None:
            if cluster_info["deploy_type"] == "tiup-cluster":
                component_config_file = os.path.join(
                    diag_data_path, host + deploy_dir, "conf/config.json"
                )
            elif cluster_info["deploy_type"] == "tidb-operator":
                component_config_file = os.path.join(diag_data_path, deploy_dir, "conf/config.json")
            with open(component_config_file, "r") as fr:
                return fr.read()
        else:
            if component_type == "tiflash":
                return "{}"
            else:
                raise Exception("not find %s config" % component_type)
    except Exception as e:
        print(e)
        raise e


def get_cluster_info(diag_data_path):
    cluster_config_file = os.path.join(diag_data_path, "cluster.json")
    if not os.path.exists(cluster_config_file):
        print("file '%s' not exists" % cluster_config_file)
        raise Exception("file '%s' not exists" % cluster_config_file)
    try:
        cluster_info = {}
        with open(cluster_config_file, "r") as fr:
            cluster_topology = json.loads(fr.read())["topology"]
        pd_servers = []
        for pd_server in cluster_topology["pd"]:
            pd_servers.append(
                {
                    "host": pd_server["host"],
                    "port": pd_server["port"],
                    "status_port": pd_server["status_port"],
                }
            )
        tidb_servers = []
        for tidb_server in cluster_topology["tidb"]:
            tidb_servers.append(
                {
                    "host": tidb_server["host"],
                    "port": tidb_server["port"],
                    "status_port": tidb_server["status_port"],
                }
            )
        tikv_servers = []
        for tikv_server in cluster_topology["tikv"]:
            tikv_servers.append(
                {
                    "host": tikv_server["host"],
                    "port": tikv_server["port"],
                    "status_port": tikv_server["status_port"],
                }
            )
        tiflash_servers = []
        if "tiflash" in cluster_topology:
            for tiflash_server in cluster_topology["tiflash"]:
                tiflash_servers.append(
                    {
                        "host": tiflash_server["host"],
                        "port": tiflash_server["port"],
                        "status_port": tiflash_server["status_port"],
                    }
                )
        cluster_info["pd_servers"] = pd_servers
        cluster_info["tidb_servers"] = tidb_servers
        cluster_info["tikv_servers"] = tikv_servers
        cluster_info["tiflash_servers"] = tiflash_servers
        return json.dumps(cluster_info)
    except Exception as e:
        raise e


def decompress_zstandard_to_folder(input_file, destination_dir):
    input_file = pathlib.Path(input_file)
    with open(input_file, "rb") as compressed:
        decomp = zstandard.ZstdDecompressor()
        output_path = pathlib.Path(destination_dir) / input_file.stem
        with open(output_path, "wb") as destination:
            decomp.copy_stream(compressed, destination)


def tar_extract(tar_path, target_path):
    try:
        tar = tarfile.open(tar_path)
        file_names = tar.getnames()
        for file_name in file_names:
            tar.extract(file_name, target_path)
        tar.close()
    except Exception as e:
        raise e


def lambda_handler(event, context):
    try:
        # get event parameters
        cluster_id = event["kwargs"]["cluster"]
        package_id = event["kwargs"]["package_id"]
        # download diag data
        if not os.path.exists(DIAG_DATA_PATH):
            os.mkdir(DIAG_DATA_PATH)
        object_file = "%s/%s-conf" % (cluster_id, package_id)
        download_s3_file = "%s/%s-conf.tar.zst" % (DIAG_DATA_PATH, package_id)
        s3 = boto3.client("s3")
        s3.download_file(S3_BUCKET_PACK, object_file, download_s3_file)
        # decompress diag data
        decompress_zstandard_to_folder(download_s3_file, DIAG_DATA_PATH)
        tar_file = "%s/%s-conf.tar" % (DIAG_DATA_PATH, package_id)
        tar_extract(tar_file, DIAG_DATA_PATH)
        # get output
        output = {
            "pd_config": get_component_config("pd", DIAG_DATA_PATH),
            "tidb_config": get_component_config("tidb", DIAG_DATA_PATH),
            "tikv_config": get_component_config("tikv", DIAG_DATA_PATH),
            "tiflash_config": get_component_config("tiflash", DIAG_DATA_PATH),
            "cluster_info": get_cluster_info(DIAG_DATA_PATH),
        }
        return output
    except Exception as e:
        print(e)
        raise e
    finally:
        # cleanup diag data
        if os.path.exists(DIAG_DATA_PATH):
            shutil.rmtree(DIAG_DATA_PATH)


if __name__ == "__main__":
    response = lambda_handler(
        {
            "task_id": "1",
            "kwargs": {
                "cluster": "6862533764151471111",
                "package_id": "7bc876dd5df32193-7205709c4974f187-64d5ff67f5e8ffb7",
            },
        },
        None,
    )
    print(json.dumps(response))