import copy
import json
import re
from enum import Enum

import yaml

from config.sku_config import SKU_CONFIG
from src.clients.tiworkload_client import Workload, get_workload_br_storage


class CloudProvider(Enum):
    AWS = "aws"
    GCP = "gcp"
    AZURE = "azure"
    ALIYUN = "aliyun"


RESOURCE_POOL_MAP = {
    ("aws", "us-west-2"): "eks-amd64-default",
    ("aws", "us-east-1"): "eks-useast1-default",
    ("gcp", "us-west1"): "gke-amd64-default",
    ("azure", "eastus2"): "azure-east-us-2",
    ("aliyun", "ap-southeast-1"): "alicloud-ap-southeast-1",
}

# Bench-toolset image mapping by cloud provider
BENCH_TOOLSET_IMAGE_MAP = {
    "aws": "946646677266.dkr.ecr.us-west-2.amazonaws.com/perf_testing/bench-toolset:latest",
    "gcp": "us-west1-docker.pkg.dev/qa-infra-dev/qa/perf_testing/bench-toolset:latest",
    "azure": None,  # Not supported yet
    "aliyun": None,  # Not supported yet
}

# Case ID and name mapping by cloud provider for bench-ch workload
BENCH_CH_CASE_MAP = {
    "aws": {"caseId": 3030027, "caseName": "perftest_bench_dbaas"},
    "gcp": {"caseId": 6720010, "caseName": "perftest_bench_dbaas_gcp"},
    "azure": None,  # Not supported yet
    "aliyun": None,  # Not supported yet
}


def _strip_yaml_comments(text):
    """Remove full-line YAML comments from a multi-line string."""
    if not isinstance(text, str):
        return text
    lines = []
    for line in text.splitlines():
        stripped = line.lstrip()
        if stripped.startswith("#"):
            continue
        lines.append(line)
    return "\n".join(lines)


def load_workflow(yaml_path):
    """Reads the react flow yaml output."""
    with open(yaml_path, "r") as f:
        return yaml.safe_load(f)


def load_sku_config_from_file(yaml_path):
    """Reads the SKU config yaml and returns a dict."""
    with open(yaml_path, "r") as f:
        return yaml.safe_load(f)


def sync_start_node_params(workflow):
    """
    Syncs startNodeParams from the __START__ node's inputs.
    startNodeParams and nodes[id=__START__].data.inputs are consistent
    except that inputs have a 'kind' field which startNodeParams do not.
    """
    nodes = workflow.get("nodes", [])
    start_node = next((n for n in nodes if n.get("id") == "__START__"), None)

    if not start_node:
        print("Warning: __START__ node not found in workflow.")
        return workflow

    # Get inputs from __START__ node
    inputs = start_node.get("data", {}).get("inputs", [])

    # Construct startNodeParams
    start_node_params = []
    for inp in inputs:
        # Deepcopy to avoid modifying the original input dict if we were to modify it in place
        param = copy.deepcopy(inp)
        # Remove 'kind' field as it's not present in startNodeParams
        if "kind" in param:
            del param["kind"]
        start_node_params.append(param)

    workflow["startNodeParams"] = start_node_params
    return workflow


def update_start_node_inputs(workflow, sku_config):
    """Updates __START__ node inputs based on sku_config."""
    nodes = workflow.get("nodes", [])
    start_node = next((n for n in nodes if n.get("id") == "__START__"), None)
    if not start_node:
        return

    inputs = start_node.get("data", {}).get("inputs", [])
    cluster_conf = sku_config.get("cluster", {})

    # Map sku_config cluster keys to __START__ input names
    mapping = {
        "cluster_version": "cluster_version",
        "org_id": "tidbcloud_org_id",
        "project_id": "tidbcloud_project_id",
    }

    for sku_key, input_name in mapping.items():
        if sku_key in cluster_conf:
            val = str(cluster_conf[sku_key])
            for inp in inputs:
                if inp.get("name") == input_name:
                    inp["value"] = val

    # Handle bench-ch specific parameters
    workload_config = sku_config.get("workload_config", {})
    if workload_config.get("template_type") == "bench-ch":
        ch_params = workload_config.get("params", {})
        user_email = sku_config.get("user_email")
        tiflash_count = cluster_conf.get("tiflash", {}).get("replicas", 1)

        # Update CH-specific parameters in __START__ inputs
        ch_mapping = {
            "threads": ch_params.get("threads"),
            "duration": ch_params.get("duration"),
            "warehouses": ch_params.get("warehouses"),
            "workload_node_memory": ch_params.get("workload_node_memory"),
            "workload_node_cpu": ch_params.get("workload_node_cpu"),
            "tiflash_replicas": str(tiflash_count),
            "email": user_email,
        }

        for input_name, value in ch_mapping.items():
            if value is not None:
                for inp in inputs:
                    if inp.get("name") == input_name:
                        inp["value"] = str(value)


def update_testbed_config(workflow, sku_config):
    """Updates all testbed_create nodes spec based on sku_config."""
    nodes = workflow.get("nodes", [])
    testbed_nodes = [n for n in nodes if n.get("type") == "testbed_create"]

    if not testbed_nodes:
        return

    # Update resourcePool
    provider = sku_config.get("cloud_provider")
    region = sku_config.get("region")

    # Normalize provider to lowercase (e.g., "AWS" -> "aws", "GCP" -> "gcp")
    if provider:
        provider = provider.lower()

    # Strip cloud provider prefix from region if present (e.g., "gcp-us-west1" -> "us-west1")
    if region and "-" in region:
        parts = region.split("-", 1)
        if len(parts) == 2 and parts[0].lower() in ("aws", "gcp", "azure", "aliyun"):
            region = parts[1]

    resource_pool = RESOURCE_POOL_MAP.get((provider, region))

    cluster_conf = sku_config.get("cluster", {})

    # Process each testbed_create node
    for testbed_node in testbed_nodes:
        # Update resourcePool
        if resource_pool:
            testbed_node["data"]["spec"]["resourcePool"] = resource_pool

        # Update tc item spec
        items = testbed_node["data"]["spec"].get("items", [])
        tc_item = next((i for i in items if i.get("name") == "tc"), None)
        if not tc_item:
            continue

        # Parse the spec string (it is YAML content in the file)
        spec_str = tc_item.get("spec", "")
        if not spec_str:
            continue

        # Fix unquoted template variables for YAML parsing
        # Matches: key: {{...}} -> key: "{{...}}"
        # Regex: : (?!")(\{\{.*?\}\})
        # Replacement: : "$1"
        spec_str_fixed = re.sub(r': (?!")(\{\{.*?\}\})', r': "\1"', spec_str)

        try:
            spec_data = yaml.safe_load(spec_str_fixed)
        except yaml.YAMLError:
            # Fallback or error handling if needed
            continue

        # Update cloud provider (use normalized lowercase value)
        # Skip for NextGen clusters — provider is not a valid field
        if provider and tc_item.get("type") != "TIDBCLOUD_NEXTGEN_CLUSTER":
            spec_data["provider"] = provider

        # Update TiDB/TiKV/TiFlash/TiProxy config (Dedicated only)
        if tc_item.get("type") != "TIDBCLOUD_NEXTGEN_CLUSTER":
            if "tidb" in cluster_conf:
                spec_data["tidb"] = cluster_conf["tidb"]

            if "tikv" in cluster_conf:
                spec_data["tikv"] = cluster_conf["tikv"]

            tiflash_conf = cluster_conf.get("tiflash")
            if not tiflash_conf:
                if "tiflash" in spec_data:
                    del spec_data["tiflash"]
            elif "nodeProfileName" in tiflash_conf:
                spec_data["tiflash"] = tiflash_conf

            tiproxy_conf = cluster_conf.get("tiproxy")
            if not tiproxy_conf:
                if "tiproxy" in spec_data:
                    del spec_data["tiproxy"]
            elif "nodeProfileName" in tiproxy_conf:
                spec_data["tiproxy"] = tiproxy_conf

        # Dump back to string (using YAML to preserve format compatibility)
        tc_item["spec"] = yaml.safe_dump(spec_data)

        # Update benchtoolset workload node image (for bench-ch workloads)
        benchtoolset_item = next(
            (i for i in items if i.get("name") == "benchtoolset"), None
        )
        if benchtoolset_item and provider:
            bench_image = BENCH_TOOLSET_IMAGE_MAP.get(provider)
            if bench_image:  # Only update if we have a valid image for this provider
                bench_spec_str = benchtoolset_item.get("spec", "")
                if bench_spec_str:
                    try:
                        # Parse the benchtoolset spec (also YAML)
                        bench_spec_data = yaml.safe_load(bench_spec_str)
                        if isinstance(bench_spec_data, dict):
                            # Update container.image
                            if "container" in bench_spec_data:
                                bench_spec_data["container"]["image"] = bench_image
                                if provider == "gcp":
                                    bench_spec_data["container"]["tolerations"] = [
                                        {
                                            "key": "nodegroup",
                                            "operator": "Equal",
                                            "value": "default-n2-standard-8",
                                            "effect": "NoSchedule",
                                        }
                                    ]
                                    bench_spec_data["container"]["nodeSelector"] = {
                                        "nodegroup": "default-n2-standard-8"
                                    }
                                # Dump back to string
                                benchtoolset_item["spec"] = yaml.safe_dump(
                                    bench_spec_data
                                )
                    except yaml.YAMLError:
                        # Skip if parsing fails
                        pass


def update_test_case_config(workflow, sku_config):
    """Updates test_case_call nodes based on sku_config (e.g., caseId and caseName for bench-ch)."""
    workload_config = sku_config.get("workload_config", {})

    # Only apply for bench-ch template
    if workload_config.get("template_type") != "bench-ch":
        return

    provider = sku_config.get("cloud_provider")
    if not provider:
        return

    # Normalize provider to lowercase
    provider = provider.lower()

    # Get the appropriate case config for this provider
    case_config = BENCH_CH_CASE_MAP.get(provider)
    if not case_config:
        # Provider not supported for bench-ch workload
        return

    # Find and update test_case_call nodes
    nodes = workflow.get("nodes", [])
    for node in nodes:
        if node.get("type") != "test_case_call":
            continue

        data = node.get("data", {})
        spec = data.get("spec", {})

        # Only update if it's a bench-ch related test case
        # We can identify it by checking if the current caseName starts with "perftest_bench_dbaas"
        current_case_name = spec.get("caseName", "")
        if current_case_name.startswith("perftest_bench_dbaas"):
            spec["caseId"] = case_config["caseId"]
            spec["caseName"] = case_config["caseName"]


def _deep_merge_dict(target, patch):
    """Recursively merge patch into target dict."""
    for key, value in patch.items():
        if isinstance(value, dict) and isinstance(target.get(key), dict):
            _deep_merge_dict(target[key], value)
        else:
            target[key] = value


def _apply_workload_override_to_node(node, node_override):
    """Apply workload/prepare overrides to a ti_workload_call node."""
    data = node.get("data") or {}
    spec = data.get("spec") or {}

    for field in ("workload", "prepare"):
        override_val = node_override.get(field)
        if not override_val:
            continue

        content = spec.get(field)
        if not content:
            continue

        # content is YAML string containing params, etc.
        try:
            content_no_comments = _strip_yaml_comments(content)
            parsed = yaml.safe_load(content_no_comments)
        except yaml.YAMLError:
            continue

        if not isinstance(parsed, dict):
            continue

        if not isinstance(override_val, dict):
            continue

        _deep_merge_dict(parsed, override_val)
        spec[field] = yaml.safe_dump(parsed)


def _apply_sleep_override_to_node(node, node_override):
    """Apply sleep.duration override to a sleep node."""
    data = node.get("data") or {}
    spec = data.get("spec") or {}
    sleep_cfg = node_override.get("sleep") or {}
    duration = sleep_cfg.get("duration")
    if duration is not None:
        spec["duration"] = duration


def update_workload_config(workflow, sku_config):
    """Apply per-node workload/sleep overrides based on sku_config['workload_config']."""
    overrides = sku_config.get("workload_config") or {}
    if not isinstance(overrides, dict):
        return

    # For bench-ch template, parameters are already handled in update_start_node_inputs
    if overrides.get("template_type") == "bench-ch":
        return

    node_overrides = overrides.get("nodes") or {}
    if not isinstance(node_overrides, dict):
        return

    nodes = workflow.get("nodes", [])
    nodes_by_name = {}
    for n in nodes:
        data = n.get("data") or {}
        name = data.get("name")
        if name:
            nodes_by_name[name] = n

    for name, node_override in node_overrides.items():
        if not isinstance(node_override, dict):
            continue
        node = nodes_by_name.get(name)
        if not node:
            continue

        node_type = node.get("type")
        if node_type == "ti_workload_call":
            _apply_workload_override_to_node(node, node_override)
        elif node_type == "sleep":
            _apply_sleep_override_to_node(node, node_override)


def _update_prepare_br_source(workflow, sku_config):
    """Update prepare.params.storage_url/s3_endpoint from TiWorkload datasets."""
    provider = sku_config.get("cloud_provider")
    region = sku_config.get("region")
    if not provider or not region:
        return

    # Normalize provider to lowercase (e.g., "AWS" -> "aws", "GCP" -> "gcp")
    provider = provider.lower()

    # Strip cloud provider prefix from region if present (e.g., "gcp-us-west1" -> "us-west1")
    if "-" in region:
        parts = region.split("-", 1)
        if len(parts) == 2 and parts[0].lower() in ("aws", "gcp", "azure", "aliyun"):
            region = parts[1]

    nodes = workflow.get("nodes", [])

    for node in nodes:
        if node.get("type") != "ti_workload_call":
            continue

        data = node.get("data") or {}
        spec = data.get("spec") or {}
        prepare_str = spec.get("prepare")
        if not prepare_str:
            continue

        workload_name = spec.get("workload_name")
        if not workload_name:
            continue

        workload_enum = None
        if workload_name == "tpcc":
            workload_enum = Workload.TPCC
        elif workload_name == "oltp_read_write":
            workload_enum = Workload.SYSBENCH
        elif workload_name == "tpch":
            workload_enum = Workload.TPCH

        if not workload_enum:
            continue

        # Fetch datasets from TiWorkload
        datasets = get_workload_br_storage(workload_enum, provider, region)
        if not datasets:
            continue

        dataset = datasets[0]
        storage_url = dataset.get("storage_url")
        s3_endpoint = dataset.get("s3_endpoint")
        if not storage_url and not s3_endpoint:
            continue

        try:
            prepare_no_comments = _strip_yaml_comments(prepare_str)
            parsed = yaml.safe_load(prepare_no_comments)
        except yaml.YAMLError:
            continue

        if not isinstance(parsed, dict):
            continue

        params = parsed.get("params") or {}
        if not isinstance(params, dict):
            continue

        if storage_url:
            params["storage_url"] = storage_url
        if s3_endpoint:
            params["s3_endpoint"] = s3_endpoint

        parsed["params"] = params
        spec["prepare"] = yaml.safe_dump(parsed)


def _normalize_ti_workload_specs(workflow):
    """Normalize ti_workload_call node specs by stripping YAML comments.

    For each ti_workload_call node, if spec.prepare/spec.workload exist and are
    YAML strings, we remove full-line comments and re-dump them via
    safe_load/safe_dump so that the final payload does not contain commented
    example options (e.g. --kespace=kespace-a) even when no overrides or
    TiWorkload BR replacements are applied.
    """
    nodes = workflow.get("nodes", [])

    for node in nodes:
        if node.get("type") != "ti_workload_call":
            continue

        data = node.get("data") or {}
        spec = data.get("spec") or {}

        for field in ("prepare", "workload"):
            content = spec.get(field)
            if not content or not isinstance(content, str):
                continue

            try:
                # Fix unquoted template variables for YAML parsing, e.g.:
                #   br_version: {{.inputs.cluster_version}}
                # ->
                #   br_version: "{{.inputs.cluster_version}}"
                content_fixed = re.sub(r': (?!")(\{\{.*?\}\})', r': "\1"', content)
                content_no_comments = _strip_yaml_comments(content_fixed)
                parsed = yaml.safe_load(content_no_comments)
            except yaml.YAMLError:
                # If parsing fails, keep original content to avoid breaking the workflow
                continue

            if not isinstance(parsed, dict):
                continue

            spec[field] = yaml.safe_dump(parsed)


def modify_workflow(workflow, sku_config):
    """
    Modifies the workflow based on sku_config.
    """
    update_start_node_inputs(workflow, sku_config)
    update_testbed_config(workflow, sku_config)
    update_test_case_config(workflow, sku_config)
    # First normalize original ti_workload_call specs so that any commented
    # example options are removed from prepare/workload blocks.
    _normalize_ti_workload_specs(workflow)
    update_workload_config(workflow, sku_config)
    _update_prepare_br_source(workflow, sku_config)


def generate_payload(workflow):
    """
    Generates the JSON payload for the API.
    """
    return json.dumps(workflow, indent=2)


def main():
    """Example usage: Generate workflow payload from config."""

    yaml_path = "config/test_workflow.yaml"
    try:
        workflow = load_workflow(yaml_path)

        # Remove startNodeParams if it exists, to ensure we regenerate it fresh
        if "startNodeParams" in workflow:
            del workflow["startNodeParams"]

        modify_workflow(workflow, SKU_CONFIG)
        workflow = sync_start_node_params(workflow)

        payload = generate_payload(workflow)
        print(payload)
    except FileNotFoundError as e:
        print(f"Error: File not found - {e}")
    except Exception as e:
        print(f"An error occurred: {e}")


if __name__ == "__main__":
    main()