import argparse
import json
import sys
from enum import Enum

import requests
from environs import Env
from requests.auth import HTTPDigestAuth

# Load environment variables
env = Env()
env.read_env()

TIDB_STAGING_TOKEN = env("TIDB_STAGING_TOKEN")

TCMS_AVAILABLE_REGIONS = {
    "aws": ["us-west-2", "us-east-1"],
    "gcp": ["us-west1"],
    "azure": ["eastus2"],
    "aliyun": ["ap-southeast-1"],
}

if not TIDB_STAGING_TOKEN:
    print("Error: TIDB_STAGING_TOKEN not found in .env")
    sys.exit(1)

try:
    PUBLIC_KEY, PRIVATE_KEY = TIDB_STAGING_TOKEN.split(":")
except ValueError:
    print("Error: TIDB_STAGING_TOKEN must be in format 'public_key:private_key'")
    sys.exit(1)

BASE_URL = "https://dedicated.staging.tidbapi.com/v1beta1"


class ComponentType(Enum):
    TIDB = "TIDB"
    TIKV = "TIKV"
    PD = "PD"
    TIPROXY = "TIPROXY"
    TIFLASH = "TIFLASH"


def get_regions(cloud_provider=None):
    """
    Get all regions, optionally filtered by cloud provider.
    Returns only regions that are available in TCMS_AVAILABLE_REGIONS.

    Args:
        cloud_provider: Optional filter like 'AWS', 'GCP', 'AZURE', 'ALIYUN'
    Returns:
        List of region objects matching the cloud provider and available in TCMS
    """
    url = f"{BASE_URL}/regions"
    headers = {"Accept": "application/json"}
    try:
        response = requests.get(
            url, headers=headers, auth=HTTPDigestAuth(PUBLIC_KEY, PRIVATE_KEY)
        )
        response.raise_for_status()
        all_regions = response.json().get("regions", [])

        if cloud_provider:
            # Filter regions by cloud provider prefix in regionId
            cloud_prefix = cloud_provider.lower()
            filtered_regions = [
                r
                for r in all_regions
                if r.get("regionId", "").startswith(f"{cloud_prefix}-")
            ]
        else:
            filtered_regions = all_regions

        # Filter by TCMS_AVAILABLE_REGIONS: only return regions that are in the allowlist
        tcms_filtered_regions = []
        for region in filtered_regions:
            region_id = region.get("regionId", "")
            # Extract cloud provider and region name from regionId (format: "aws-us-west-2")
            parts = region_id.split("-", 1)
            if len(parts) == 2:
                cloud = parts[0].lower()
                region_name = parts[1]
                # Check if this cloud provider has available regions and this region is in the list
                if (
                    cloud in TCMS_AVAILABLE_REGIONS
                    and region_name in TCMS_AVAILABLE_REGIONS[cloud]
                ):
                    tcms_filtered_regions.append(region)

        return tcms_filtered_regions
    except requests.exceptions.RequestException as e:
        print(f"Error fetching regions: {e}")
        if hasattr(e, "response") and e.response is not None:
            print(f"Response content: {e.response.text}")
        return []


def get_nodespecs(region_id, component_type=None):
    """
    Get node specs for a region, optionally filtered by component type.
    Args:
        region_id: Region ID like 'aws-us-west-2'
        component_type: Optional ComponentType enum to filter results
    Returns:
        List of nodeSpec objects
    """
    url = f"{BASE_URL}/regions/{region_id}/nodeSpecs?pageSize=100"
    headers = {"Accept": "application/json"}
    try:
        response = requests.get(
            url, headers=headers, auth=HTTPDigestAuth(PUBLIC_KEY, PRIVATE_KEY)
        )
        response.raise_for_status()
        all_specs = response.json().get("nodeSpecs", [])

        if component_type:
            filtered_specs = [
                s for s in all_specs if s.get("componentType") == component_type.value
            ]
            return filtered_specs

        return all_specs
    except requests.exceptions.RequestException as e:
        print(f"Error fetching node specs for {region_id}: {e}")
        if hasattr(e, "response") and e.response is not None:
            print(f"Response content: {e.response.text}")
        return []


def check_region(region_id):
    regions = get_regions()
    # Check if region_id matches 'name' in the region list
    # The API returns a list of region objects. We assume 'name' is the identifier like 'azure-eastus2'
    for region in regions:
        if region.get("regionId") == region_id:
            return True
    return False


def get_storages(node_specs, tidb_spec_key):
    """
    Get TIKV and TIFLASH node specs that match the vCpu count of the given tidb_spec_key.
    (Deprecated: Use get_available_components for more accurate filtering)
    """
    # Find the vCpu count for the given tidb_spec_key
    target_vcpu = 0
    for spec in node_specs:
        if spec.get("nodeSpecKey") == tidb_spec_key:
            target_vcpu = spec.get("vCpu")
            break

    if target_vcpu == 0:
        print(f"Warning: TiDB spec {tidb_spec_key} not found or has 0 vCpu")
        return []

    filtered_specs = []
    for spec in node_specs:
        if spec.get("componentType") in [
            ComponentType.TIKV.value,
            ComponentType.TIFLASH.value,
        ]:
            if spec.get("vCpu") == target_vcpu:
                filtered_specs.append(spec)

    return filtered_specs


def get_available_components(node_specs, tidb_spec_key):
    """
    Get available component node specs based on TiDB spec, following the cluster creation rules.

    Rules (from desired_comp_create_check.go:componentsSpecialCheck):

    2C/4C Component Rules:
    - 2 vCPU components can only work with other 2 vCPU components
    - 4 vCPU components can only work with other 4 vCPU components

    When TiDB has 2 or 4 vCPUs:
    - Only one TiDB node group is allowed
    - TiDB node quantity must be 1 or 2
    - TiFlash is not supported

    When TiDB has 2 vCPUs:
    - TiProxy is not supported

    When TiKV has 2 vCPUs:
    - Must have exactly 3 nodes

    When TiKV has 4 vCPUs:
    - Feature gate Multi4cTikv required to have > 3 nodes (simplified: allow 3+ nodes)

    For other vCPU counts (8C, 16C, etc.):
    - All components are allowed
    - TiKV/TiFlash can be any vCPU >= 8 (not restricted to match TiDB)
    - Matching vCPU specs are listed first (recommended default)

    Version Matching Rule (additional business rule, not from Go code):
    - All components must match TiDB's version suffix (e.g., V1, V2, V3)
    - Example: 8C16G-V2 TiDB can only use V2 components (8C16G-V2, 16C32G-V2, etc.)

    Args:
        node_specs: List of all node specs for the region
        tidb_spec_key: The selected TiDB nodeSpecKey (e.g., '2C4G-V2', '8C16G-V2')

    Returns:
        Dictionary with component types as keys and their available specs as values:
        {
            'tikv': [list of tikv specs],
            'tiflash': [list of tiflash specs],
            'pd': [list of pd specs],
            'tiproxy': [list of tiproxy specs],
            'constraints': {
                'tidb_count_range': (min, max),
                'tikv_count_range': (min, max),
                'tiflash_allowed': bool,
                'tiproxy_allowed': bool,
                'tikv_count_step': int  # TiKV must be multiple of this
            }
        }
    """
    # Validate input
    if not isinstance(node_specs, list):
        print(f"Error: node_specs must be a list, got {type(node_specs)}")
        return {"tikv": [], "tiflash": [], "pd": [], "tiproxy": [], "constraints": {}}

    # Find the TiDB spec
    tidb_spec = None
    for spec in node_specs:
        if spec.get("nodeSpecKey") == tidb_spec_key:
            tidb_spec = spec
            break

    if not tidb_spec:
        print(f"Warning: TiDB spec {tidb_spec_key} not found")
        return {"tikv": [], "tiflash": [], "pd": [], "tiproxy": [], "constraints": {}}

    tidb_vcpu = tidb_spec.get("vCpu", 0)

    # Get version from TiDB spec (API provides "version" field directly, e.g., "V1", "V2", "V3")
    # Version matching: all components must have the same version as TiDB
    tidb_version = tidb_spec.get("version")

    def version_matches(spec):
        """Check if spec version matches TiDB version

        Exception: PD and TiProxy are exempt from version matching.
        Only TiDB, TiKV, and TiFlash must have matching versions.
        """
        if not tidb_version:
            return True  # No version requirement
        # Handle edge case where spec might not be a dict
        if not isinstance(spec, dict):
            print(f"Warning: spec is not a dict, got {type(spec)}: {spec}")
            return False

        # PD and TiProxy are exceptions - they don't need to match version
        comp_type = spec.get("componentType")
        if comp_type in [ComponentType.PD.value, ComponentType.TIPROXY.value]:
            return True  # Skip version check for PD and TiProxy

        # For TiDB, TiKV, TiFlash - version must match
        spec_version = spec.get("version")
        if not spec_version:
            return False  # Spec has no version
        return spec_version == tidb_version

    # Initialize result structure
    result = {
        "tikv": [],
        "tiflash": [],
        "pd": [],
        "tiproxy": [],
        "constraints": {
            "tidb_count_range": (1, None),  # Default: min 1, no max
            "tikv_count_range": (3, None),  # Default: min 3, no max
            "tiflash_allowed": True,
            "tiproxy_allowed": True,
            "tikv_count_step": 3,  # TiKV must be multiple of 3
        },
    }

    # Apply 2C/4C special rules
    if tidb_vcpu in [2, 4]:
        # When TiDB is 2C or 4C
        result["constraints"]["tidb_count_range"] = (1, 2)  # Only 1 or 2 nodes
        result["constraints"]["tiflash_allowed"] = False  # TiFlash not supported

        if tidb_vcpu == 2:
            result["constraints"][
                "tiproxy_allowed"
            ] = False  # TiProxy not supported for 2C

        # Filter components: must match TiDB vCPU (2C or 4C) and version
        for spec in node_specs:
            spec_vcpu = spec.get("vCpu", 0)
            comp_type = spec.get("componentType")

            # Check version matching first
            if not version_matches(spec):
                continue  # Version must match

            if spec_vcpu != tidb_vcpu:
                continue  # Must match TiDB vCPU

            if comp_type == ComponentType.TIKV.value:
                result["tikv"].append(spec)
                if tidb_vcpu == 2:
                    # 2C TiKV must be exactly 3 nodes
                    result["constraints"]["tikv_count_range"] = (3, 3)
                elif tidb_vcpu == 4:
                    # 4C TiKV can be 3+ (simplified, no feature gate check)
                    result["constraints"]["tikv_count_range"] = (3, None)

            elif comp_type == ComponentType.PD.value:
                result["pd"].append(spec)

            elif (
                comp_type == ComponentType.TIPROXY.value
                and result["constraints"]["tiproxy_allowed"]
            ):
                result["tiproxy"].append(spec)

    else:
        # For 8C, 16C, and other vCPU counts: allow all components with vCPU >= 8 and matching version
        # Sort specs to prioritize matching vCPU (for default selection)
        matching_tikv = []
        other_tikv = []
        matching_tiflash = []
        other_tiflash = []

        for spec in node_specs:
            spec_vcpu = spec.get("vCpu", 0)
            comp_type = spec.get("componentType")

            # Check version matching first
            if not version_matches(spec):
                continue  # Version must match

            if comp_type == ComponentType.TIKV.value and spec_vcpu >= 8:
                if spec_vcpu == tidb_vcpu:
                    matching_tikv.append(spec)
                else:
                    other_tikv.append(spec)

            elif comp_type == ComponentType.TIFLASH.value and spec_vcpu >= 8:
                if spec_vcpu == tidb_vcpu:
                    matching_tiflash.append(spec)
                else:
                    other_tiflash.append(spec)

            elif comp_type == ComponentType.PD.value:
                # PD can have different vCPU (typically 4C or 8C)
                result["pd"].append(spec)

            elif comp_type == ComponentType.TIPROXY.value:
                # TiProxy can have different vCPU
                result["tiproxy"].append(spec)

        # Put matching vCPU specs first (for default selection), then others
        result["tikv"] = matching_tikv + other_tikv
        result["tiflash"] = matching_tiflash + other_tiflash

    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="TiDB Cloud Client")
    parser.add_argument("--cloud", required=True, help="Cloud provider (optional)")
    parser.add_argument(
        "--region", required=True, help="Region ID (e.g., azure-eastus2)"
    )
    parser.add_argument("--tidb", help="TiDB nodeSpecs like '8C16G'")

    args = parser.parse_args()

    region_id = f"{args.cloud}-{args.region}"

    if check_region(region_id):
        specs = get_nodespecs(region_id)
        # print(json.dumps(specs, indent=2))
    else:
        print(f"Region {args.region} in Cloud {args.cloud} does not exist")
    if args.tidb:
        storages = get_storages(specs, args.tidb)
        print(json.dumps(storages, indent=2))
    else:
        print(f"No available storages for TiDB spec '{args.tidb}'")