import copy
import json
from typing import Any, Dict, List

from utils.preprocessing_biz import merge_node_specs


def resolve_ref(spec: Dict[str, Any], ref: str) -> Dict[str, Any]:
    parts = ref.split("/")
    current = spec
    for part in parts[1:]:
        current = current[part]
    return copy.deepcopy(current)


def resolve_refs(obj: Any, spec: Dict[str, Any]) -> Any:
    if isinstance(obj, dict):
        new_obj: Dict[str, Any] = {}
        ref_resolved: Dict[str, Any] = {}  # keep value inside ref
        for key, value in obj.items():
            if key == "$ref":
                ref_content = resolve_ref(spec, value)
                ref_resolved = resolve_refs(ref_content, spec)
                # new_obj.update(ref_resolved)  # keep value outside ref
            else:
                new_obj[key] = resolve_refs(value, spec)
        new_obj = {**new_obj, **ref_resolved}  # keep value inside ref
        return new_obj
    elif isinstance(obj, list):
        return [resolve_refs(item, spec) for item in obj]
    else:
        return obj


def resolve_allOf(obj: Any) -> Any:
    if isinstance(obj, dict):
        if "allOf" in obj:
            merged: Dict[str, Any] = {}
            for item in obj["allOf"]:
                item_resolved = resolve_allOf(item)
                if isinstance(item_resolved, dict):
                    deep_merge(merged, item_resolved)

            for key, value in obj.items():
                if key != "allOf":
                    deep_merge(merged, {key: resolve_allOf(value)})
            return merged
        else:
            return {k: resolve_allOf(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [resolve_allOf(item) for item in obj]
    else:
        return obj


def deep_merge(d1: Dict[str, Any], d2: Dict[str, Any]) -> Dict[str, Any]:
    for k, v in d2.items():
        if isinstance(v, dict):
            d1[k] = deep_merge(d1.get(k, {}), v)
        elif isinstance(v, list):
            if k in d1:
                d1[k].extend(v)
            else:
                d1[k] = v
        else:
            d1[k] = v
    return d1


def enhance_spec_definition(
    origin_spec: Dict[str, Any],
    enhanced_def: str,
    enhanced_keys: List[str],
    ref_spec: Dict[str, Any],
    ref_key: str,
) -> None:
    """
    Enhance a specific definition in the API specification.

    :param origin_spec: The original API specification to be enhanced.
    :param enhanced_def: The key of the definition to be enhanced in the spec.
    :param enhanced_keys: List of keys within the definition to be enhanced.
    :param ref_spec: Reference specification containing enhancement data.
    :param ref_key: Key in the reference spec to use for enhancements.
    """
    definition = origin_spec["definitions"].get(enhanced_def)
    if definition and "properties" in definition:
        for key in enhanced_keys:
            if key in definition["properties"]:
                if key == "nodeSpecKey":
                    definition["properties"][key]["enum"] = ref_spec[ref_key][
                        "nodeSpecKey"
                    ]
                elif key == "storageSizeGi":
                    definition["properties"][key]["maximum"] = ref_spec[ref_key][
                        "storageSizeGi"
                    ]["max"]
                    definition["properties"][key]["minimum"] = ref_spec[ref_key][
                        "storageSizeGi"
                    ]["min"]


def add_biz_restrictions(spec: Dict[str, Any]) -> Dict[str, Any]:
    """
    Add business restrictions to the API specification.
    """
    ref_spec = merge_node_specs()

    # Define the settings to update
    enhancements = [
        # Storage related enhancements
        (
            "v1beta1ClusterStorageNodeSetting",
            ["storageSizeGi", "nodeSpecKey"],
            "STORAGE",
        ),
        (
            "v1beta1UpdateClusterRequestStorageNodeSetting",
            ["storageSizeGi", "nodeSpecKey"],
            "STORAGE",
        ),
        # TiDB related enhancements
        ("v1beta1ClusterTidbNodeSetting", ["nodeSpecKey"], "TIDB"),
        ("v1beta1UpdateClusterRequestTidbNodeSetting", ["nodeSpecKey"], "TIDB"),
        (
            "v1beta1AdvancedModifyDedicatedClusterRequestTidbNodeSetting",
            ["nodeSpecKey"],
            "TIDB",
        ),
    ]

    # Update each setting
    for enhanced_def, enhanced_keys, ref_key in enhancements:
        enhance_spec_definition(spec, enhanced_def, enhanced_keys, ref_spec, ref_key)

    return spec


def unwrap_refs(spec: Dict[str, Any]) -> Dict[str, Any]:
    spec = add_biz_restrictions(spec)

    ref_resolved_spec = resolve_refs(copy.deepcopy(spec), spec)
    allOf_resolved_spec = resolve_allOf(ref_resolved_spec)
    return allOf_resolved_spec