#!/usr/bin/env python3

import argparse
import difflib
import json
import yaml
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument("--mode", help="how to differ configuration")
parser.add_argument("--platform", help="run in github or shell")
args = parser.parse_args()

current_file_path = Path(__file__)
pools_dir = (
    f"{current_file_path.parent.parent.parent}/cd/manifests/manager/serverless-tidb/pools"
)
manager_pulumi_dir = f"{current_file_path.parent.parent.parent}/cd/projects/manager"
gateway_pulumi_dir = f"{current_file_path.parent.parent.parent}/cd/projects/tidb-gateway"
http_gateway_pulumi_dir = (
    f"{current_file_path.parent.parent.parent}/cd/projects/tidb-http-gateway"
)
pool_names = ["canary", "canary-vip", "default", "export", "vip"]


def extract_pool_configmap(path, key) -> list[str]:
    with open(path) as f:
        data = yaml.safe_load(f)["data"]
        return data.get(key, "").split("\n")


def highlight_diff_lines(text) -> bool:
    """
    Prints text lines with diff lines highlighted in yellow.
    Assumes diff lines start with '+' or '-'.
    """
    yellow = "\033[33m"
    red = "\033[31m"
    bold = "\033[1m"
    reset = "\033[0m"

    empty = True
    for line in text.splitlines():
        empty = False
        if args.platform == "github":
            print(line)
        elif line.startswith("+++ ") or line.startswith("--- "):
            print(f"{bold}{line}{reset}")
        elif line.startswith("+"):
            print(f"{yellow}{line}{reset}")
        elif line.startswith("-"):
            print(f"{red}{line}{reset}")
        else:
            print(line)

    return empty


def diff_tidb_pool_across_pools(key):
    """
    Use default pool as the standard one, then diff other pools with it in ap-northeast-1
    """
    filename = "pod-cm-prod-ap-northeast-1-f01.yaml"

    standard_file_name = f"{pools_dir}/default/{filename}"
    standard_config_file = extract_pool_configmap(standard_file_name, key)

    for pool_name in pool_names:
        pool_file_name = f"{pools_dir}/{pool_name}/{filename}"
        pool_config_file = extract_pool_configmap(pool_file_name, key)
        diff = difflib.unified_diff(
            standard_config_file,
            pool_config_file,
            fromfile=standard_file_name,
            tofile=pool_file_name,
            n=1,
            lineterm="",
        )
        if not highlight_diff_lines("\n".join(diff)):
            print()


def diff_tidb_pool_across_regions(key):
    """
    Use ap-northeast-1 region as the standard one, then diff other regions with it in every pool
    """
    eu = "pod-cm-prod-ap-northeast-1-f01.yaml"
    for pool_name in pool_names:
        standard_file_name = f"{pools_dir}/{pool_name}/{eu}"
        standard_config_file = extract_pool_configmap(standard_file_name, key)

        directory = Path(f"{pools_dir}/{pool_name}")
        for path in directory.iterdir():
            if path.name == eu or not path.name.startswith("pod-cm-prod-"):
                continue
            region_config_file = extract_pool_configmap(path, key)
            diff = difflib.unified_diff(
                standard_config_file,
                region_config_file,
                fromfile=standard_file_name,
                tofile=str(path),
                n=1,
                lineterm="",
            )
            if not highlight_diff_lines("\n".join(diff)):
                print()


tidb_pool_dev = "pod-cm-dev-us-east-1-f01.yaml"
tidb_pool_staging = "pod-cm-staging-us-east-1-f01.yaml"
tidb_pool_prod_ap_northeast_1 = "pod-cm-prod-ap-northeast-1-f01.yaml"
tidb_pool_prod_regions = [
    "pod-cm-prod-ap-northeast-1-f01.yaml",
    "pod-cm-prod-eu-central-1-f01.yaml",
    "pod-cm-prod-us-east-1-f01.yaml",
    "pod-cm-prod-us-east-1-f02.yaml",
    "pod-cm-prod-us-west-2-f01.yaml",
    "pod-cm-prod-us-west-2-f02.yaml",
    "pod-cm-prod-ap-southeast-1-f01.yaml",
]


def diff_tidb_pool(fromfile, tofile, key):
    for pool_name in pool_names:
        from_path = f"{pools_dir}/{pool_name}/{fromfile}"
        from_config = extract_pool_configmap(from_path, key)
        to_path = f"{pools_dir}/{pool_name}/{tofile}"
        to_config = extract_pool_configmap(to_path, key)

        diff = difflib.unified_diff(
            from_config,
            to_config,
            fromfile=from_path,
            tofile=to_path,
            n=1,
            lineterm="",
        )
        if not highlight_diff_lines("\n".join(diff)):
            print()


pulumi_dev = "Pulumi.dev.yaml"
pulumi_staging = "Pulumi.staging.yaml"
pulumi_prod = "Pulumi.prod.yaml"
pulumi_prod_ap_northeast_1 = "Pulumi.prod-ap-northeast-1-f01.yaml"
pulumi_prod_regions = [
    "Pulumi.prod-ap-northeast-1-f01.yaml",
    "Pulumi.prod-us-west-2-f01.yaml",
    "Pulumi.prod-eu-central-1-f01.yaml",
    "Pulumi.prod-ap-southeast-1-f01.yaml",
    "Pulumi.prod-us-east-1-f01.yaml",
    "Pulumi.prod-us-east-1-f02.yaml",
]


def clear_unconcerned_keys(dic):
    unconcerned_keys = ["free-tier-stacks"]
    result = {}
    for k, v in dic["config"].items():
        skip = False
        for uk in unconcerned_keys:
            if uk in k:
                skip = True
        if skip:
            continue

        result[k] = v

    return result


def load_pulumi_config(directory, filename) -> dict:
    with open(f"{directory}/{filename}") as f:
        result = yaml.safe_load(f)
        if filename in pulumi_prod_regions:
            with open(f"{directory}/{pulumi_prod}") as prod:
                tmp = result
                result = yaml.safe_load(prod)
                for k, v in tmp["config"].items():
                    result["config"][k] = v

        return clear_unconcerned_keys(result)


def diff_pulumi(directory, fromfile, tofile):
    # a = (json.dumps(load_pulumi_config(directory, fromfile), indent=" ")).split("\n")
    # b = (json.dumps(load_pulumi_config(directory, tofile), indent=" ")).split("\n")

    # yaml will ignore the order of keys by serialize them in directory order
    a = (
        yaml.dump(
            load_pulumi_config(directory, fromfile),
        )
    ).split("\n")
    b = (
        yaml.dump(
            load_pulumi_config(directory, tofile),
        )
    ).split("\n")
    diff = difflib.unified_diff(
        a,
        b,
        fromfile=f"{directory}/{fromfile}",
        tofile=f"{directory}/{tofile}",
        n=1,
        lineterm="",
    )
    if not highlight_diff_lines("\n".join(diff)):
        print()


startup_key = "startup-script"
config_key = "config-file"


def main():
    # dev mode differ dev environment configuration with statging
    if args.mode == "dev":
        diff_pulumi(manager_pulumi_dir, pulumi_staging, pulumi_dev)
        diff_pulumi(gateway_pulumi_dir, pulumi_staging, pulumi_dev)
        diff_pulumi(http_gateway_pulumi_dir, pulumi_staging, pulumi_dev)
        diff_tidb_pool(tidb_pool_staging, tidb_pool_dev, startup_key)
        diff_tidb_pool(tidb_pool_staging, tidb_pool_dev, config_key)
    # staging mode differ staging environment configuration with prod ap-northeast-1
    elif args.mode == "staging":
        diff_pulumi(
            manager_pulumi_dir,
            pulumi_prod_ap_northeast_1,
            pulumi_staging,
        )
        diff_pulumi(
            gateway_pulumi_dir,
            pulumi_prod_ap_northeast_1,
            pulumi_staging,
        )
        diff_pulumi(
            http_gateway_pulumi_dir,
            pulumi_prod_ap_northeast_1,
            pulumi_staging,
        )
        diff_tidb_pool(tidb_pool_prod_ap_northeast_1, tidb_pool_staging, startup_key)
        diff_tidb_pool(tidb_pool_prod_ap_northeast_1, tidb_pool_staging, config_key)
    # prod mode differ production environment configuration in each region with ap-northeast-1
    elif args.mode == "prod":
        for region in pulumi_prod_regions:
            diff_pulumi(
                manager_pulumi_dir,
                pulumi_prod_ap_northeast_1,
                region,
            )
            diff_pulumi(
                gateway_pulumi_dir,
                pulumi_prod_ap_northeast_1,
                region,
            )
            diff_pulumi(
                http_gateway_pulumi_dir,
                pulumi_prod_ap_northeast_1,
                region,
            )
        for region in tidb_pool_prod_regions:
            diff_tidb_pool(tidb_pool_prod_ap_northeast_1, region, startup_key)
            diff_tidb_pool(tidb_pool_prod_ap_northeast_1, region, config_key)
    else:
        raise ValueError("unknown mode")


if __name__ == "__main__":
    main()