# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "pymysql",
#   "rich",
# ]
# ///
"""
TiDB one-box skill — deploy, validate, and destroy a TiDB-X on EKS dev environment.

Commands:
    bootstrap    terraform init + validate + plan (no apply; review before deploy)
    deploy       terraform apply for the full stack (EKS + add-ons + TiDB-X)
    destroy      destroy all Terraform-managed resources with automatic remediations
    scale        scale TiDB-X replicas and/or EKS node group size
    validate     check TiDB-X cluster health via kubectl + status endpoint
    verify-sql   verify SQL pushdown to TiKV via pymysql
    stop         pause the environment by scaling the EKS node group to zero
    start        resume the environment by restoring the EKS node group size
    refresh-auth refresh AWS SSO + upsert EKS access entry + update kubeconfig

Usage:
    uv run tidb-onebox.py --repo-root /path/to/one-box <command> [options]
"""

import argparse
import json
import re
import shlex
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

from rich.console import Console

console = Console()

TERRAFORM_CMD = ["rtk", "terraform"] if shutil.which("rtk") else ["terraform"]
TERRAFORM_CMD_DISPLAY = " ".join(TERRAFORM_CMD)

GROUP_RESOURCES = [
    "pdgroups.core.pingcap.com",
    "tikvgroups.core.pingcap.com",
    "tidbgroups.core.pingcap.com",
    "tsogroups.core.pingcap.com",
    "schedulinggroups.core.pingcap.com",
    "tiproxygroups.core.pingcap.com",
    "tiflashgroups.core.pingcap.com",
    "tikvworkergroups.core.pingcap.com",
    "resourcemanagergroups.core.pingcap.com",
]

EKS_ADMIN_POLICY = "arn:aws:eks::aws:cluster-access-policy/AmazonEKSClusterAdminPolicy"

SCALE_COMPONENTS = [
    {
        "arg": "pd",
        "display": "pd",
        "tf_var": "scale_pd_replicas",
        "resource": "pdgroups.core.pingcap.com",
        "name": "pd",
    },
    {
        "arg": "tikv",
        "display": "tikv",
        "tf_var": "scale_tikv_replicas",
        "resource": "tikvgroups.core.pingcap.com",
        "name": "tikv",
    },
    {
        "arg": "tidb",
        "display": "tidb",
        "tf_var": "scale_tidb_replicas",
        "resource": "tidbgroups.core.pingcap.com",
        "name": "tidb",
    },
    {
        "arg": "tso",
        "display": "tso",
        "tf_var": "scale_pd_tso_replicas",
        "resource": "tsogroups.core.pingcap.com",
        "name": "tso",
    },
    {
        "arg": "scheduling",
        "display": "scheduling",
        "tf_var": "scale_pd_scheduling_replicas",
        "resource": "schedulinggroups.core.pingcap.com",
        "name": "scheduling",
    },
    {
        "arg": "tiproxy",
        "display": "tiproxy",
        "tf_var": "scale_tiproxy_replicas",
        "resource": "tiproxygroups.core.pingcap.com",
        "name": "tiproxy",
    },
    {
        "arg": "tiflash_write",
        "display": "tiflash-write",
        "tf_var": "scale_tiflash_write_replicas",
        "resource": "tiflashgroups.core.pingcap.com",
        "name": "tiflash-write",
    },
    {
        "arg": "tiflash_compute",
        "display": "tiflash-compute",
        "tf_var": "scale_tiflash_compute_replicas",
        "resource": "tiflashgroups.core.pingcap.com",
        "name": "tiflash-compute",
    },
]


# ---------------------------------------------------------------------------
# Subprocess helpers
# ---------------------------------------------------------------------------


_SHELL_META = re.compile(r"[|<>&;`$()*?\[\]]|\s2>|\s>>")


def _split(cmd: str | list[str]) -> list[str]:
    """Convert a command to an argv list.

    For plain commands we tokenize with shlex (executed without a shell). For
    commands that rely on shell features (pipes, redirections, globs, etc.) we
    fall back to invoking ``bash -c`` as the argv head so we still avoid
    ``shell=True`` (which Semgrep flags as a security risk).
    """
    if isinstance(cmd, list):
        return cmd
    if _SHELL_META.search(cmd):
        return ["bash", "-c", cmd]
    return shlex.split(cmd)


def run(cmd: str | list[str], *, cwd: str | None = None, check: bool = True) -> int:
    """Stream a command to the terminal; return exit code."""
    argv = _split(cmd)
    result = subprocess.run(argv, cwd=cwd)
    if check and result.returncode != 0:
        raise subprocess.CalledProcessError(result.returncode, argv)
    return result.returncode


def run_tee(
    cmd: str | list[str],
    *,
    cwd: str | None = None,
    check: bool = True,
    log_prefix: str = "cmd",
    announce: str | None = None,
) -> tuple[int, str]:
    """Stream a command to stdout while teeing output to a temp log file."""
    argv = _split(cmd)
    with tempfile.NamedTemporaryFile(
        mode="wb",
        prefix=f"{log_prefix}-",
        suffix=".log",
        dir="/tmp",
        delete=False,
    ) as logf:
        log_path = logf.name
        if announce:
            console.print(f"[dim]{announce}: log saved to {log_path}[/dim]")
        proc = subprocess.Popen(
            argv,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
        assert proc.stdout is not None
        for chunk in iter(lambda: proc.stdout.read(4096), b""):
            sys.stdout.buffer.write(chunk)
            sys.stdout.flush()
            logf.write(chunk)
        rc = proc.wait()
    if check and rc != 0:
        raise RuntimeError(f"Command failed. Log saved to {log_path}: {' '.join(argv)}")
    return rc, log_path


def capture(cmd: str | list[str], *, cwd: str | None = None, check: bool = True) -> str:
    """Run a command silently; return stdout stripped."""
    argv = _split(cmd)
    result = subprocess.run(argv, cwd=cwd, capture_output=True, text=True)
    if check and result.returncode != 0:
        raise RuntimeError(f"Command failed: {argv}\n{result.stderr.strip()}")
    return result.stdout.strip()


def capture3(cmd: str | list[str], *, cwd: str | None = None) -> tuple[int, str, str]:
    """Return (returncode, stdout, stderr)."""
    argv = _split(cmd)
    r = subprocess.run(argv, cwd=cwd, capture_output=True, text=True)
    return r.returncode, r.stdout.strip(), r.stderr.strip()


def tf_cmd(*args: str) -> list[str]:
    return [*TERRAFORM_CMD, *args]


def tf_shell(args: str) -> str:
    return f"{TERRAFORM_CMD_DISPLAY} {args}"


# ---------------------------------------------------------------------------
# tfvars / state helpers
# ---------------------------------------------------------------------------


def get_tfvar(env_dir: str, key: str) -> str:
    """Read a simple scalar from dev.tfvars."""
    tfvars = Path(env_dir) / "dev.tfvars"
    if not tfvars.exists():
        return ""
    for line in tfvars.read_text().splitlines():
        m = re.match(rf'^\s*{re.escape(key)}\s*=\s*"?([^"#\n]+)"?\s*(?:#.*)?$', line)
        if m:
            return m.group(1).strip()
    return ""


def tf_state_has(address: str, *, cwd: str) -> bool:
    rc, out, _ = capture3(tf_shell("state list 2>/dev/null"), cwd=cwd)
    return address in out.splitlines()


def get_vpc_id(env_dir: str) -> str:
    addr = "module.network.module.vpc.aws_vpc.this[0]"
    if not tf_state_has(addr, cwd=env_dir):
        return ""
    return capture(
        tf_shell(
            f"state show '{addr}' | awk -F'\"' '/^[[:space:]]*id[[:space:]]*=/ {{print $2; exit}}'"
        ),
        cwd=env_dir,
    )


def get_eks_cluster_name(env_dir: str) -> str:
    """EKS cluster name = ${project_name}-${environment} (matches main.tf locals)."""
    project = get_tfvar(env_dir, "project_name") or "tidb-onebox"
    env = get_tfvar(env_dir, "environment") or "dev"
    return f"{project}-{env}"


def get_tf_output(env_dir: str, name: str, *, default: str = "") -> str:
    rc, out, _ = capture3(tf_cmd("output", "-raw", name), cwd=env_dir)
    if rc == 0 and out:
        return out
    return default


def get_nodegroup_defaults(env_dir: str) -> tuple[int, int, int]:
    """Return (desired, min, max) from dev.tfvars."""
    desired = int(get_tfvar(env_dir, "node_desired_size") or "0")
    min_size = int(get_tfvar(env_dir, "node_min_size") or "0")
    max_size = int(get_tfvar(env_dir, "node_max_size") or "0")
    return desired, min_size, max_size


def list_nodegroups(cluster: str, region: str) -> list[str]:
    out = capture(
        f"aws eks list-nodegroups --cluster-name {cluster} --region {region} "
        "--query 'nodegroups' --output text",
        check=False,
    )
    return out.split() if out else []


def resolve_nodegroup(cluster: str, region: str, requested: str | None = None) -> str:
    if requested:
        return requested
    nodegroups = list_nodegroups(cluster, region)
    if not nodegroups:
        raise RuntimeError(f"No managed node groups found for cluster '{cluster}'.")
    if len(nodegroups) > 1:
        raise RuntimeError(
            f"Multiple node groups found for cluster '{cluster}': {', '.join(nodegroups)}. "
            "Re-run with --nodegroup-name."
        )
    return nodegroups[0]


def get_nodegroup_scaling(
    cluster: str, nodegroup: str, region: str
) -> dict[str, int | str]:
    out = capture(
        f"aws eks describe-nodegroup --cluster-name {cluster} --nodegroup-name {nodegroup} "
        f"--region {region} --output json"
    )
    nodegroup_data = json.loads(out)["nodegroup"]
    scaling = nodegroup_data["scalingConfig"]
    return {
        "desired": int(scaling["desiredSize"]),
        "min": int(scaling["minSize"]),
        "max": int(scaling["maxSize"]),
        "status": str(nodegroup_data["status"]),
    }


def update_nodegroup_scaling(
    cluster: str,
    nodegroup: str,
    region: str,
    *,
    desired: int,
    min_size: int,
    max_size: int,
) -> None:
    run(
        f"aws eks update-nodegroup-config --cluster-name {cluster} "
        f"--nodegroup-name {nodegroup} --region {region} "
        f"--scaling-config minSize={min_size},maxSize={max_size},desiredSize={desired}"
    )


def wait_for_nodegroup_scaling(
    cluster: str,
    nodegroup: str,
    region: str,
    *,
    desired: int,
    timeout_seconds: int = 1200,
) -> dict[str, int | str]:
    deadline = time.time() + timeout_seconds
    last = get_nodegroup_scaling(cluster, nodegroup, region)
    while time.time() < deadline:
        current = get_nodegroup_scaling(cluster, nodegroup, region)
        last = current
        console.print(
            f"  Nodegroup {nodegroup}: status={current['status']} desired={current['desired']} "
            f"min={current['min']} max={current['max']}"
        )
        if current["status"] == "ACTIVE" and current["desired"] == desired:
            return current
        time.sleep(15)
    raise RuntimeError(
        f"Timed out waiting for node group '{nodegroup}' to reach desiredSize={desired}. "
        f"Last status={last['status']} desired={last['desired']} min={last['min']} max={last['max']}"
    )


def get_k8s_nodegroup_nodes(nodegroup: str) -> list[dict]:
    selector = f"eks.amazonaws.com/nodegroup={nodegroup}"
    rc, out, err = capture3(f"kubectl get nodes -l {selector} -o json")
    if rc != 0:
        raise RuntimeError(
            f"Failed to query Kubernetes nodes for node group '{nodegroup}': {err}"
        )
    return json.loads(out).get("items", [])


def count_ready_schedulable_nodes(nodes: list[dict]) -> int:
    ready = 0
    for node in nodes:
        if node.get("spec", {}).get("unschedulable"):
            continue
        conditions = node.get("status", {}).get("conditions", [])
        if any(
            c.get("type") == "Ready" and c.get("status") == "True" for c in conditions
        ):
            ready += 1
    return ready


def wait_for_k8s_nodegroup_state(
    nodegroup: str,
    *,
    expected_registered: int | None = None,
    expected_ready_schedulable: int | None = None,
    exact_ready_schedulable: bool = False,
    timeout_seconds: int = 1200,
) -> None:
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        nodes = get_k8s_nodegroup_nodes(nodegroup)
        registered = len(nodes)
        ready_schedulable = count_ready_schedulable_nodes(nodes)
        console.print(
            f"  Kubernetes nodes for {nodegroup}: registered={registered} ready_schedulable={ready_schedulable}"
        )
        if expected_registered is not None and registered != expected_registered:
            time.sleep(15)
            continue
        if expected_ready_schedulable is not None and (
            ready_schedulable < expected_ready_schedulable
            or (
                exact_ready_schedulable
                and ready_schedulable != expected_ready_schedulable
            )
        ):
            time.sleep(15)
            continue
        return
    raise RuntimeError(
        f"Timed out waiting for Kubernetes nodes in node group '{nodegroup}'. "
        f"Expected registered={expected_registered} ready_schedulable={expected_ready_schedulable}."
    )


def wait_for_ready_pods(namespace: str, *, timeout_seconds: int = 1200) -> None:
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        rc, out, err = capture3(f"kubectl get pods -n {namespace} -o json")
        if rc != 0:
            raise RuntimeError(
                f"Failed to query pods in namespace '{namespace}': {err}"
            )

        items = json.loads(out).get("items", [])
        total = len(items)
        ready = 0
        all_ready = total > 0

        for item in items:
            statuses = item.get("status", {}).get("containerStatuses") or []
            if statuses and all(status.get("ready") for status in statuses):
                ready += 1
                continue
            all_ready = False

        console.print(f"  Pods in {namespace}: ready={ready}/{total}")
        if all_ready:
            return

        time.sleep(10)

    raise RuntimeError(
        f"Timed out waiting for pods in namespace '{namespace}' to become Ready."
    )


def get_namespace_pod_status(namespace: str) -> tuple[int, int]:
    rc, out, err = capture3(f"kubectl get pods -n {namespace} -o json")
    if rc != 0:
        raise RuntimeError(f"Failed to query pods in namespace '{namespace}': {err}")

    total = 0
    ready = 0
    for item in json.loads(out).get("items", []):
        metadata = item.get("metadata", {})
        if metadata.get("deletionTimestamp"):
            continue
        total += 1
        statuses = item.get("status", {}).get("containerStatuses") or []
        if statuses and all(status.get("ready") for status in statuses):
            ready += 1
    return total, ready


def get_namespaced_resource_json(
    resource: str, name: str, namespace: str
) -> dict | None:
    rc, out, _ = capture3(
        f"kubectl get {resource} {name} -n {namespace} -o json --ignore-not-found"
    )
    if rc != 0 or not out.strip():
        return None
    data = json.loads(out)
    if not data.get("kind"):
        return None
    return data


def wait_for_scaled_workload(
    namespace: str,
    requested_targets: dict[str, int],
    *,
    timeout_seconds: int = 900,
    stall_timeout_seconds: int = 180,
) -> None:
    deadline = time.time() + timeout_seconds
    last_summary = ""
    last_summary_snapshot = ""
    last_changed_at = time.time()
    while time.time() < deadline:
        expected_pods = 0
        targets_applied = True
        target_specs_applied = True
        summary_parts = []

        for component in SCALE_COMPONENTS:
            obj = get_namespaced_resource_json(
                component["resource"], component["name"], namespace
            )
            target = requested_targets.get(component["arg"])
            if obj is None:
                if target == 0:
                    summary_parts.append(f"{component['display']}=absent")
                    continue
                if target is not None:
                    targets_applied = False
                    target_specs_applied = False
                    summary_parts.append(f"{component['display']}=missing")
                continue

            spec = obj.get("spec", {})
            status = obj.get("status", {})
            replicas = int(spec.get("replicas", 0))
            current_replicas = int(status.get("currentReplicas", replicas))
            ready_replicas = int(status.get("readyReplicas", current_replicas))
            expected_pods += current_replicas
            summary_parts.append(
                f"{component['display']}={replicas}/{current_replicas}/{ready_replicas}"
            )
            if target is not None and replicas != target:
                targets_applied = False
                target_specs_applied = False
            if current_replicas != replicas or ready_replicas != current_replicas:
                targets_applied = False

        total_pods, ready_pods = get_namespace_pod_status(namespace)
        last_summary = (
            f"{' '.join(summary_parts)} pods={ready_pods}/{total_pods} "
            f"expected_pods={expected_pods}"
        )
        console.print(f"  Scale status {namespace}: {last_summary}")

        if last_summary != last_summary_snapshot:
            last_changed_at = time.time()

        if targets_applied and total_pods == expected_pods and ready_pods == total_pods:
            return

        # TiKV scale-in can legitimately sit in a steady-looking state while PD
        # migrates leaders/regions. Once the spec is applied and pods are Ready,
        # fall back to the overall timeout instead of failing on the shorter
        # stall timeout.
        status_only_catchup = (
            target_specs_applied
            and total_pods == expected_pods
            and ready_pods == total_pods
        )
        if (
            not status_only_catchup
            and time.time() - last_changed_at >= stall_timeout_seconds
        ):
            raise RuntimeError(
                f"Scale stalled for {stall_timeout_seconds}s in namespace '{namespace}'. "
                f"Last status: {last_summary}"
            )

        last_summary_snapshot = last_summary
        time.sleep(10)

    raise RuntimeError(
        f"Timed out waiting for scaled workload to settle in namespace '{namespace}'. "
        f"Last status: {last_summary}"
    )


def get_first_pod_name(namespace: str, selector: str) -> str:
    rc, out, err = capture3(
        f"kubectl get pods -n {namespace} -l {shlex.quote(selector)} -o json"
    )
    if rc != 0:
        raise RuntimeError(
            f"Failed to query pods in namespace '{namespace}' for selector '{selector}': {err}"
        )
    items = json.loads(out).get("items", [])
    if not items:
        raise RuntimeError(
            f"No pods found in namespace '{namespace}' for selector '{selector}'."
        )
    return items[0]["metadata"]["name"]


def get_pd_replication_max_replicas(namespace: str) -> int:
    pd_pod = get_first_pod_name(namespace, "pingcap.com/component=pd")
    rc, out, err = capture3(
        "kubectl exec "
        f"-n {shlex.quote(namespace)} {shlex.quote(pd_pod)} -- "
        "curl -fsS http://127.0.0.1:2379/pd/api/v1/config/replicate"
    )
    if rc != 0:
        raise RuntimeError(f"Failed to query PD replication config: {err}")
    return int(json.loads(out)["max-replicas"])


def set_pd_replication_max_replicas(namespace: str, replicas: int) -> None:
    pd_pod = get_first_pod_name(namespace, "pingcap.com/component=pd")
    payload = json.dumps({"max-replicas": replicas})
    rc, _, err = capture3(
        "kubectl exec "
        f"-n {shlex.quote(namespace)} {shlex.quote(pd_pod)} -- "
        "curl -fsS -X POST http://127.0.0.1:2379/pd/api/v1/config/replicate "
        "-H 'Content-Type: application/json' "
        f"-d {shlex.quote(payload)}"
    )
    if rc != 0:
        raise RuntimeError(f"Failed to update PD replication config: {err}")

    deadline = time.time() + 120
    while time.time() < deadline:
        current = get_pd_replication_max_replicas(namespace)
        console.print(f"  PD replication max-replicas={current}")
        if current == replicas:
            return
        time.sleep(5)

    raise RuntimeError(
        f"Timed out waiting for PD replication max-replicas to become {replicas}."
    )


def get_tikv_instances(namespace: str) -> list[dict]:
    rc, out, err = capture3(
        f"kubectl get tikvs.core.pingcap.com -n {namespace} -o json"
    )
    if rc != 0:
        raise RuntimeError(f"Failed to query TiKV instances in '{namespace}': {err}")
    return json.loads(out).get("items", [])


def get_pd_region_check_count(namespace: str, name: str) -> int:
    pd_pod = get_first_pod_name(namespace, "pingcap.com/component=pd")
    rc, out, err = capture3(
        "kubectl exec "
        f"-n {shlex.quote(namespace)} {shlex.quote(pd_pod)} -- "
        f"curl -fsS http://127.0.0.1:2379/pd/api/v1/regions/check/{shlex.quote(name)}"
    )
    if rc != 0:
        raise RuntimeError(f"Failed to query PD region check '{name}': {err}")
    return int(json.loads(out).get("count", 0))


def ensure_tikv_scale_out_safe(namespace: str) -> None:
    blocking = []
    for item in get_tikv_instances(namespace):
        name = item.get("metadata", {}).get("name", "<unknown>")
        spec = item.get("spec", {})
        status = item.get("status", {})
        state = status.get("state", "")
        if spec.get("offline"):
            blocking.append(f"{name}: spec.offline=true")
        if state in {"Removing", "Removed"}:
            blocking.append(f"{name}: state={state}")

    offline_peer_count = get_pd_region_check_count(namespace, "offline-peer")
    if offline_peer_count > 0:
        blocking.append(f"PD offline-peer regions={offline_peer_count}")

    if blocking:
        raise RuntimeError(
            "Refusing TiKV scale-out while a previous TiKV scale-in is still in progress or incomplete: "
            + "; ".join(blocking)
        )


def list_namespaced_resources(resource: str, namespace: str) -> list[str]:
    rc, out, _ = capture3(
        f"kubectl get {resource} -n {namespace} -o name --ignore-not-found"
    )
    if rc != 0 or not out.strip():
        return []
    return [line for line in out.splitlines() if line]


def list_namespace_pvs(namespace: str) -> list[str]:
    rc, out, err = capture3("kubectl get pv -o json")
    if rc != 0:
        raise RuntimeError(f"Failed to query PVs: {err}")
    pvs = []
    for item in json.loads(out).get("items", []):
        claim_ref = item.get("spec", {}).get("claimRef") or {}
        if claim_ref.get("namespace") == namespace:
            name = item.get("metadata", {}).get("name")
            if name:
                pvs.append(name)
    return pvs


def wait_for_namespace_storage_cleanup(
    namespace: str = "tidb-x-cluster", *, timeout_seconds: int = 1800
) -> None:
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        rc, _, _ = capture3(f"kubectl get ns {namespace}")
        if rc != 0:
            return

        pods = list_namespaced_resources("pods", namespace)
        pvcs = list_namespaced_resources("pvc", namespace)
        pvs = list_namespace_pvs(namespace)
        console.print(
            f"  Namespace cleanup {namespace}: pods={len(pods)} pvcs={len(pvcs)} pvs={len(pvs)}"
        )
        if not pods and not pvcs and not pvs:
            return
        time.sleep(10)

    raise RuntimeError(
        f"Timed out waiting for namespace '{namespace}' pods/PVCs/PVs to clean up."
    )


def eks_cluster_exists(cluster: str, region: str) -> bool:
    rc, _, _ = capture3(f"aws eks describe-cluster --name {cluster} --region {region}")
    return rc == 0


def ensure_destroy_kube_access(
    cluster: str, region: str, sso_region: str, namespace: str = "tidb-x-cluster"
) -> bool:
    rc, _, _ = capture3(f"kubectl auth can-i get pods -n {shlex.quote(namespace)}")
    if rc == 0:
        return True

    run(f"aws eks update-kubeconfig --name {cluster} --region {region}", check=False)
    rc, _, _ = capture3(f"kubectl auth can-i get pods -n {shlex.quote(namespace)}")
    if rc == 0:
        return True

    principal_arn = resolve_principal_arn(sso_region)
    restore_eks_access(cluster, principal_arn, region)
    run(f"aws eks update-kubeconfig --name {cluster} --region {region}", check=False)
    rc, _, _ = capture3(f"kubectl auth can-i get pods -n {shlex.quote(namespace)}")
    return rc == 0


def destroy_tidb_x_workload(
    env_dir: str,
    cluster: str,
    region: str,
    sso_region: str,
    namespace: str = "tidb-x-cluster",
) -> None:
    address = "module.tidb_x.helm_release.tidb_cluster"
    if not tf_state_has(address, cwd=env_dir):
        return
    if not eks_cluster_exists(cluster, region):
        return
    if not ensure_destroy_kube_access(cluster, region, sso_region, namespace):
        raise RuntimeError(
            "Cannot access the cluster to gracefully destroy TiDB-X before infrastructure teardown."
        )

    console.rule("[bold cyan]Destroy TiDB-X workload[/bold cyan]")
    rc_apply, apply_log = run_tee(
        tf_cmd(
            "destroy",
            "-no-color",
            f"-target={address}",
            "-var-file=dev.tfvars",
            "-auto-approve",
        ),
        cwd=env_dir,
        check=False,
        log_prefix="terraform-destroy-tidb-cluster",
        announce="destroy tidb cluster",
    )
    if rc_apply != 0:
        raise RuntimeError(
            f"Failed to destroy TiDB-X workload cleanly before infrastructure teardown. Log saved to {apply_log}."
        )

    console.print(
        "[dim]Waiting for TiDB-X pods, PVCs, and PVs to disappear before tearing down EKS...[/dim]"
    )
    wait_for_namespace_storage_cleanup(namespace)
    console.print("[green]TiDB-X workload cleanup complete.[/green]")


# ---------------------------------------------------------------------------
# S3 pre-emptying
# ---------------------------------------------------------------------------


def preflight_s3_empty(env_dir: str, region: str) -> None:
    """Empty the S3 bucket only if Terraform owns it (i.e. it appears in state).

    Uses `terraform state list` so we don't depend on the now-misleading
    `s3_create_bucket` tfvar: if `aws_s3_bucket.tidb_x` is present in state,
    Terraform created the bucket and will try to delete it on destroy, so we
    must empty it first. Otherwise (bucket pre-existed and was only referenced)
    we leave the data alone.
    """
    if not tf_state_has("module.tidb_x.aws_s3_bucket.tidb_x[0]", cwd=env_dir):
        return
    bucket = get_tfvar(env_dir, "s3_bucket_name")
    if not bucket:
        return
    rc, _, _ = capture3(f"aws s3api head-bucket --bucket {bucket} --region {region}")
    if rc != 0:
        return
    console.print(
        f"[bold]S3 bucket '[cyan]{bucket}[/cyan]' is Terraform-managed — emptying before destroy...[/bold]"
    )
    run(f"aws s3 rm s3://{bucket} --recursive --region {region}")
    console.print(f"[green]Bucket '{bucket}' emptied.[/green]")


def _delete_external_bucket(env_dir: str, region: str) -> None:
    """Empty + delete the configured S3 bucket if it is NOT in Terraform state.

    Used by `destroy --also-delete-external-bucket` to fully clean up buckets
    that were auto-created by `preflight_s3_ensure` (and therefore left behind
    by `terraform destroy`).
    """
    bucket = get_tfvar(env_dir, "s3_bucket_name")
    if not bucket:
        return
    if tf_state_has("module.tidb_x.aws_s3_bucket.tidb_x[0]", cwd=env_dir):
        return  # Terraform already deleted it
    rc, _, _ = capture3(f"aws s3api head-bucket --bucket {bucket} --region {region}")
    if rc != 0:
        return
    console.print(
        f"[bold]Deleting external S3 bucket '[cyan]{bucket}[/cyan]'...[/bold]"
    )
    run(f"aws s3 rm s3://{bucket} --recursive --region {region}", check=False)
    run(f"aws s3api delete-bucket --bucket {bucket} --region {region}", check=False)


def preflight_s3_ensure(env_dir: str, region: str) -> None:
    """Ensure the configured S3 bucket exists before deploy.

    Looks up s3_bucket_name from dev.tfvars. If the bucket already exists, no-op.
    Otherwise creates it; failure (e.g. wrong AWS profile, name conflict) aborts.
    """
    bucket = get_tfvar(env_dir, "s3_bucket_name")
    if not bucket:
        return
    rc, _, err = capture3(f"aws s3api head-bucket --bucket {bucket} --region {region}")
    if rc == 0:
        console.print(
            f"[green]S3 bucket '[cyan]{bucket}[/cyan]' already exists.[/green]"
        )
        return
    console.print(
        f"[bold]S3 bucket '[cyan]{bucket}[/cyan]' not found — creating in region '{region}'...[/bold]"
    )
    if region == "us-east-1":
        cmd = f"aws s3api create-bucket --bucket {bucket} --region {region}"
    else:
        cmd = (
            f"aws s3api create-bucket --bucket {bucket} --region {region} "
            f"--create-bucket-configuration LocationConstraint={region}"
        )
    rc2, _, err2 = capture3(cmd)
    if rc2 != 0:
        console.print(
            f"[bold red]Failed to create S3 bucket '{bucket}': {err2 or err}[/bold red]\n"
            "Check that the AWS profile/account is correct and the bucket name is globally unique."
        )
        sys.exit(1)
    console.print(f"[green]Bucket '{bucket}' created.[/green]")


# ---------------------------------------------------------------------------
# Classic ELB cleanup
# ---------------------------------------------------------------------------


def list_classic_elbs(vpc_id: str, region: str) -> list[str]:
    out = capture(
        f"aws elb describe-load-balancers --region {region} "
        f'--query "LoadBalancerDescriptions[?VPCId==`{vpc_id}`].LoadBalancerName" '
        f"--output text",
        check=False,
    )
    return out.split() if out else []


def preflight_elb(vpc_id: str, region: str, force: bool) -> None:
    if not vpc_id:
        return
    names = list_classic_elbs(vpc_id, region)
    if not names:
        return
    console.print(f"[yellow]Classic ELBs detected in VPC {vpc_id}: {names}[/yellow]")
    if not force:
        console.print(
            "Tip: rerun with --force-cleanup-elb to delete them automatically."
        )
        return
    for name in names:
        console.print(f"  Deleting ELB: {name}")
        run(
            f"aws elb delete-load-balancer --region {region} --load-balancer-name {name}"
        )
    for attempt in range(1, 25):
        remaining = int(
            capture(
                f"aws ec2 describe-network-interfaces --region {region} "
                f"--filters 'Name=vpc-id,Values={vpc_id}' "
                f"--query 'length(NetworkInterfaces[?RequesterManaged==`true`])' --output text",
                check=False,
            )
            or "0"
        )
        console.print(
            f"  Waiting for ELB ENIs to disappear: remaining={remaining} (attempt {attempt})"
        )
        if remaining == 0:
            return
        time.sleep(10)
    console.print("[yellow]Warning: timed out waiting for ELB ENIs.[/yellow]")


# ---------------------------------------------------------------------------
# Destroy remediations
# ---------------------------------------------------------------------------


def strip_tidb_x_finalizers(namespace: str = "tidb-x-cluster") -> bool:
    """Patch finalizers off all TiDB-X CRs in the namespace so a stuck
    Terminating namespace can finish deleting. Returns True if anything was
    patched.
    """
    rc, _, _ = capture3(f"kubectl get ns {namespace}")
    if rc != 0:
        return False
    patched = False
    rc, resources, _ = capture3(
        "kubectl api-resources --api-group=core.pingcap.com --verbs=list --namespaced -o name"
    )
    if rc != 0 or not resources.strip():
        return False
    for resource in resources.splitlines():
        if not resource:
            continue
        rc, out, _ = capture3(
            f"kubectl get {resource} -n {namespace} -o name --ignore-not-found"
        )
        if rc != 0 or not out.strip():
            continue
        for name in out.splitlines():
            if not name:
                continue
            rc_patch, _, _ = capture3(
                f"kubectl patch -n {namespace} {name} --type=merge "
                f'-p \'{{"metadata":{{"finalizers":[]}}}}\''
            )
            if rc_patch == 0:
                patched = True
    return patched


def force_delete_namespace_pods(namespace: str = "tidb-x-cluster") -> bool:
    rc, out, _ = capture3(f"kubectl get pods -n {namespace} -o name --ignore-not-found")
    if rc != 0 or not out.strip():
        return False
    deleted = False
    for name in out.splitlines():
        if not name:
            continue
        rc_del, _, _ = capture3(
            f"kubectl delete -n {namespace} {name} --force --grace-period=0 --wait=false"
        )
        if rc_del == 0:
            deleted = True
    return deleted


def remediate_known_issues(env_dir: str) -> bool:
    """Remove state entries that are known to cause timeouts. Returns True if anything was removed."""
    repaired = False

    # Helm release: uninstall may timeout after 5 min even though it succeeded
    if tf_state_has("module.tidb_x.helm_release.tidb_cluster", cwd=env_dir):
        console.print(
            "[yellow]Removing tidb_cluster helm release from state (known uninstall timeout)...[/yellow]"
        )
        run(
            tf_cmd("state", "rm", "module.tidb_x.helm_release.tidb_cluster"),
            cwd=env_dir,
        )
        run(
            "kubectl delete pods --all -n tidb-x-cluster --force --grace-period=0",
            check=False,
        )
        repaired = True

    # Namespace stuck Terminating because TiDB-X CRs still hold finalizers
    # (operator already gone, no controller to reconcile them).
    if strip_tidb_x_finalizers():
        console.print(
            "[yellow]Stripped finalizers from leftover TiDB-X CRs to unblock namespace deletion.[/yellow]"
        )
        repaired = True

    if force_delete_namespace_pods():
        console.print(
            "[yellow]Force-deleted leftover pods in tidb-x-cluster to unblock namespace deletion.[/yellow]"
        )
        repaired = True

    return repaired


def is_unauthorized(text: str) -> bool:
    return bool(
        re.search(
            r"Unauthorized|forbidden|401|403|credentials|authentication",
            text,
            re.IGNORECASE,
        )
    )


# ---------------------------------------------------------------------------
# EKS access entry
# ---------------------------------------------------------------------------


def resolve_principal_arn(
    sso_role_region: str | None = None, *, aws_profile: str | None = None
) -> str:
    p = f"--profile {aws_profile}" if aws_profile else ""
    current_arn = capture(f"aws sts get-caller-identity {p} --query Arn --output text")
    account_id = capture(
        f"aws sts get-caller-identity {p} --query Account --output text"
    )
    console.print(f"  Caller ARN : {current_arn}")
    console.print(f"  Account ID : {account_id}")
    if re.match(r"arn:aws:sts::\d+:assumed-role/", current_arn):
        role = re.sub(
            r"arn:aws:sts::[0-9]+:assumed-role/([^/]+)/.*", r"\1", current_arn
        )
        rc, role_arn, _ = capture3(
            f"aws iam get-role --role-name {shlex.quote(role)} {p} "
            "--query 'Role.Arn' --output text"
        )
        if rc == 0 and role_arn.startswith(f"arn:aws:iam::{account_id}:role/"):
            return role_arn
        if sso_role_region:
            return (
                f"arn:aws:iam::{account_id}:role/"
                f"aws-reserved/sso.amazonaws.com/{sso_role_region}/{role}"
            )
        raise RuntimeError(
            "Could not resolve the stable IAM role ARN from the current assumed-role "
            "session. Re-run with --sso-role-region if your SSO role lives in a "
            "different region than the cluster."
        )
    return current_arn


def restore_eks_access(cluster: str, principal_arn: str, region: str) -> None:
    console.print(f"Restoring EKS admin access for: {principal_arn}")
    rc, _, _ = capture3(
        f"aws eks describe-access-entry --cluster-name {cluster} "
        f"--principal-arn '{principal_arn}' --region {region}"
    )
    if rc != 0:
        run(
            f"aws eks create-access-entry --cluster-name {cluster} "
            f"--principal-arn '{principal_arn}' --region {region}"
        )
    run(
        f"aws eks associate-access-policy --cluster-name {cluster} "
        f"--principal-arn '{principal_arn}' --policy-arn {EKS_ADMIN_POLICY} "
        f"--access-scope type=cluster --region {region}",
        check=False,
    )
    console.print("[green]EKS access restored.[/green]")


# ---------------------------------------------------------------------------
# Auth preflight (used by validate / verify-sql / stop)
# ---------------------------------------------------------------------------


def preflight_auth() -> None:
    """Fail fast with a friendly message if AWS or kubectl auth is broken."""
    rc, _, err = capture3("aws sts get-caller-identity")
    if rc != 0:
        console.print(
            "[bold red]AWS credentials not valid:[/bold red] "
            f"{err.splitlines()[-1] if err else 'unknown error'}\n"
            "Refresh your credentials, then re-run the command. "
            "If your profile uses SSO, run [bold]refresh-auth[/bold]."
        )
        sys.exit(1)
    rc, _, err = capture3("kubectl auth can-i get pods --all-namespaces")
    if rc != 0:
        console.print(
            "[bold red]kubectl cannot reach the cluster:[/bold red] "
            f"{err.splitlines()[-1] if err else 'unknown error'}\n"
            "Run [bold]refresh-auth[/bold] to update kubeconfig and EKS access entry."
        )
        sys.exit(1)


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------


def cmd_bootstrap(args) -> None:
    env_dir = str(Path(args.repo_root) / "terraform/environments/dev")

    if not (Path(env_dir) / "dev.tfvars").exists():
        import shutil

        shutil.copy(Path(env_dir) / "dev.tfvars.example", Path(env_dir) / "dev.tfvars")
        console.print(
            "[yellow]Created dev.tfvars from template — review and edit image URIs and s3_bucket_name before deploying.[/yellow]"
        )

    run(tf_cmd("init", "-no-color"), cwd=env_dir)
    run(tf_cmd("fmt", "-recursive"), cwd=env_dir)
    run(tf_cmd("validate", "-no-color"), cwd=env_dir)
    _, plan_log = run_tee(
        tf_cmd("plan", "-no-color", "-var-file=dev.tfvars", "-out=tfplan"),
        cwd=env_dir,
        log_prefix="terraform-bootstrap-plan",
        announce="bootstrap plan",
    )
    console.print(f"\n[bold green]Plan written to {env_dir}/tfplan.[/bold green]")
    console.print(f"[dim]Bootstrap plan: log saved to {plan_log}[/dim]")
    console.print("Review the plan above, then run [bold]deploy[/bold] to apply.")


def cmd_deploy(args) -> None:
    env_dir = str(Path(args.repo_root) / "terraform/environments/dev")

    if not (Path(env_dir) / "dev.tfvars").exists():
        console.print(
            f"[red]Missing dev.tfvars — run [bold]bootstrap[/bold] first.[/red]"
        )
        sys.exit(1)

    region = get_tfvar(env_dir, "aws_region") or "us-west-2"
    cluster_name = get_eks_cluster_name(env_dir)

    run(tf_cmd("init", "-no-color"), cwd=env_dir)
    run(tf_cmd("validate", "-no-color"), cwd=env_dir)

    preflight_s3_ensure(env_dir, region)

    console.rule("[bold cyan]terraform apply[/bold cyan]")
    rc_apply, apply_log = run_tee(
        tf_cmd("apply", "-no-color", "-var-file=dev.tfvars", "-auto-approve"),
        cwd=env_dir,
        check=False,
        log_prefix="terraform-deploy-apply",
        announce="deploy apply",
    )
    if rc_apply != 0:
        console.print(
            "[bold red]Apply failed.[/bold red] If the failure was transient "
            "(e.g. EKS API timeout, ECR throttling), it is safe to re-run "
            "[bold]deploy[/bold] — Terraform is idempotent and will resume from "
            f"the current state. Log saved to [bold]{apply_log}[/bold]"
        )
        sys.exit(rc_apply)

    out_name = capture(tf_cmd("output", "-raw", "eks_cluster_name"), cwd=env_dir)
    run(f"aws eks update-kubeconfig --name {out_name} --region {region}")
    console.print(
        f"[bold green]Deploy complete. Kubeconfig updated for cluster '{out_name}'.[/bold green]"
    )


def cmd_destroy(args) -> None:
    env_dir = str(Path(args.repo_root) / "terraform/environments/dev")

    if not (Path(env_dir) / "dev.tfvars").exists():
        console.print(f"[red]Missing dev.tfvars.[/red]")
        sys.exit(1)

    region = get_tfvar(env_dir, "aws_region") or "us-west-2"
    cluster_name = get_eks_cluster_name(env_dir)
    sso_region = args.sso_role_region or region
    namespace = get_tf_output(
        env_dir, "tidb_x_cluster_namespace", default="tidb-x-cluster"
    )

    run(tf_cmd("init", "-no-color"), cwd=env_dir)
    run(tf_cmd("validate", "-no-color"), cwd=env_dir)

    preflight_s3_empty(env_dir, region)

    vpc_id = get_vpc_id(env_dir)
    preflight_elb(vpc_id, region, args.force_cleanup_elb)

    if not args.auto_approve:
        confirm = input("Type 'destroy' to continue: ").strip()
        if confirm != "destroy":
            console.print("Aborted.")
            sys.exit(0)

    destroy_tidb_x_workload(env_dir, cluster_name, region, sso_region, namespace)

    for attempt in range(1, 4):
        console.rule(f"[bold yellow]Destroy attempt {attempt}/3[/bold yellow]")

        # Proactively remove known timeout-prone state entries before planning
        remediate_known_issues(env_dir)

        # Plan (streaming)
        rc_plan, plan_log = run_tee(
            tf_cmd(
                "plan",
                "-no-color",
                "-destroy",
                "-var-file=dev.tfvars",
                "-out=destroy.tfplan",
            ),
            cwd=env_dir,
            check=False,
            log_prefix=f"terraform-destroy-plan-{attempt}",
            announce=f"destroy plan attempt {attempt}",
        )
        if rc_plan != 0:
            console.print(f"[red]Destroy plan failed. Log saved to {plan_log}.[/red]")
            sys.exit(rc_plan)

        rc_apply, apply_log = run_tee(
            tf_cmd("apply", "-no-color", "destroy.tfplan"),
            cwd=env_dir,
            check=False,
            log_prefix=f"terraform-destroy-apply-{attempt}",
            announce=f"destroy apply attempt {attempt}",
        )

        if rc_apply == 0:
            for f in ["destroy.tfplan", "tfplan"]:
                (Path(env_dir) / f).unlink(missing_ok=True)
            if args.also_delete_external_bucket:
                _delete_external_bucket(env_dir, region)
            console.print(
                "[bold green]All Terraform-managed resources destroyed.[/bold green]"
            )
            return

        apply_output = (
            Path(apply_log).read_text(errors="replace")
            if Path(apply_log).exists()
            else ""
        )
        console.print(
            f"[yellow]Attempt {attempt} failed. Log saved to [bold]{apply_log}[/bold]. Checking for known remediations...[/yellow]"
        )

        # After apply failure, check if Helm/gp3 entries are now removable
        if remediate_known_issues(env_dir):
            continue

        # Kubernetes provider Unauthorized (EKS access entry removed early)
        if is_unauthorized(apply_output):
            console.print(
                "[yellow]Kubernetes provider auth error — restoring EKS access entry...[/yellow]"
            )
            principal_arn = resolve_principal_arn(sso_region)
            restore_eks_access(cluster_name, principal_arn, region)
            continue

        console.print("[red]No known remediation applies.[/red]")
        break

    console.print(
        "[bold red]All destroy attempts exhausted. Manual intervention required.[/bold red]"
    )
    if vpc_id:
        run(
            f"aws ec2 describe-network-interfaces --region {region} "
            f"--filters 'Name=vpc-id,Values={vpc_id}' --output table",
            check=False,
        )
    sys.exit(1)


def cmd_validate(args) -> None:
    preflight_auth()
    ns, cluster, svc = args.namespace, args.cluster, args.status_service

    console.rule("[bold]Cluster CR[/bold]")
    run(f"kubectl get clusters.core.pingcap.com {cluster} -n {ns}", check=False)

    console.rule("[bold]Group resources[/bold]")
    for r in GROUP_RESOURCES:
        run(f"kubectl get {r} -n {ns}", check=False)

    console.rule("[bold]Waiting for all pods Ready (20 min timeout)[/bold]")
    wait_for_ready_pods(ns, timeout_seconds=1200)

    console.rule("[bold]Pods[/bold]")
    run(f"kubectl get pods -n {ns} -o wide")

    console.rule("[bold]Services[/bold]")
    run(f"kubectl get svc -n {ns}")

    console.rule("[bold]TiDB status endpoint (port 10080)[/bold]")
    pf = subprocess.Popen(
        ["kubectl", "port-forward", "-n", ns, f"svc/{svc}", "10080:10080"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    try:
        for _ in range(10):
            rc, out, _ = capture3("curl -sf http://127.0.0.1:10080/status")
            if rc == 0:
                console.print(out)
                console.print(
                    "[bold green]Cluster basic validation complete.[/bold green]"
                )
                return
            time.sleep(2)
        console.print("[bold red]Failed to reach TiDB status endpoint.[/bold red]")
        sys.exit(1)
    finally:
        pf.terminate()
        pf.wait()


def cmd_verify_sql(args) -> None:
    import pymysql

    preflight_auth()
    ns, svc, port = args.namespace, args.service, args.port
    console.print(
        f"Starting port-forward to svc/{svc} in namespace {ns} (local port {port})..."
    )

    pf = subprocess.Popen(
        ["kubectl", "port-forward", "-n", ns, f"svc/{svc}", f"{port}:4000"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    try:
        conn = None
        for _ in range(20):
            try:
                conn = pymysql.connect(
                    host="127.0.0.1",
                    port=port,
                    user=args.user,
                    password=args.password or "",
                    connect_timeout=2,
                )
                break
            except Exception:
                time.sleep(2)

        if conn is None:
            console.print("[bold red]Could not connect to TiDB.[/bold red]")
            sys.exit(1)

        ts = int(time.time())
        db, tbl = f"tikv_verify_{ts}", f"t_{ts}"

        with conn:
            with conn.cursor() as cur:
                cur.execute("SET SESSION tidb_isolation_read_engines = 'tikv'")
                cur.execute(f"CREATE DATABASE IF NOT EXISTS `{db}`")
                cur.execute(f"USE `{db}`")
                cur.execute(
                    f"CREATE TABLE `{tbl}` (id BIGINT PRIMARY KEY, v INT NOT NULL, pad VARCHAR(32) NOT NULL)"
                )
                cur.executemany(
                    f"INSERT INTO `{tbl}` (id, v, pad) VALUES (%s, %s, %s)",
                    [
                        (i, i, w)
                        for i, w in enumerate(
                            [
                                "alpha",
                                "bravo",
                                "charlie",
                                "delta",
                                "echo",
                                "foxtrot",
                                "golf",
                                "hotel",
                            ],
                            1,
                        )
                    ],
                )
                cur.execute(
                    f"EXPLAIN ANALYZE SELECT /*+ READ_FROM_STORAGE(TIKV[t]) */ SUM(v) "
                    f"FROM `{tbl}` AS t WHERE v >= 3"
                )
                plan = "\n".join(str(r) for r in cur.fetchall())
                console.print("[bold]EXPLAIN ANALYZE:[/bold]")
                console.print(plan)

                if "cop[tikv]" in plan:
                    console.print(
                        "[bold green]Success: cop[tikv] present — query pushed down to TiKV.[/bold green]"
                    )
                else:
                    console.print(
                        "[bold red]Failed: cop[tikv] not found in plan.[/bold red]"
                    )
                    sys.exit(1)

                cur.execute(f"DROP DATABASE IF EXISTS `{db}`")
    finally:
        pf.terminate()
        pf.wait()


def cmd_stop(args) -> None:
    rc, _, err = capture3("aws sts get-caller-identity")
    if rc != 0:
        console.print(
            "[bold red]AWS credentials not valid:[/bold red] "
            f"{err.splitlines()[-1] if err else 'unknown error'}"
        )
        sys.exit(1)
    env_dir = str(Path(args.repo_root) / "terraform/environments/dev")
    region = get_tfvar(env_dir, "aws_region") or "us-west-2"
    cluster = args.cluster or get_eks_cluster_name(env_dir)
    nodegroup = resolve_nodegroup(cluster, region, args.nodegroup_name)
    current = get_nodegroup_scaling(cluster, nodegroup, region)
    console.print(
        f"Pausing cluster '[bold]{cluster}[/bold]' by scaling node group "
        f"'[bold]{nodegroup}[/bold]' to zero..."
    )
    console.print(
        f"  Current scaling: desired={current['desired']} min={current['min']} max={current['max']}"
    )
    if current["desired"] == 0 and current["min"] == 0:
        console.print("[green]Node group already paused.[/green]")
        return
    update_nodegroup_scaling(
        cluster,
        nodegroup,
        region,
        desired=0,
        min_size=0,
        max_size=max(int(current["max"]), 1),
    )
    wait_for_nodegroup_scaling(cluster, nodegroup, region, desired=0)
    wait_for_k8s_nodegroup_state(nodegroup, expected_registered=0)
    console.print("[green]Environment paused.[/green]")


def cmd_start(args) -> None:
    rc, _, err = capture3("aws sts get-caller-identity")
    if rc != 0:
        console.print(
            "[bold red]AWS credentials not valid:[/bold red] "
            f"{err.splitlines()[-1] if err else 'unknown error'}"
        )
        sys.exit(1)
    env_dir = str(Path(args.repo_root) / "terraform/environments/dev")
    region = get_tfvar(env_dir, "aws_region") or "us-west-2"
    cluster = args.cluster or get_eks_cluster_name(env_dir)
    nodegroup = resolve_nodegroup(cluster, region, args.nodegroup_name)
    desired_default, min_default, max_default = get_nodegroup_defaults(env_dir)
    desired = args.desired_size if args.desired_size is not None else desired_default
    min_size = args.min_size if args.min_size is not None else min_default
    max_size = args.max_size if args.max_size is not None else max_default

    if desired < 0 or min_size < 0 or max_size < 0:
        raise RuntimeError("desired/min/max sizes must be non-negative.")
    if min_size > desired:
        raise RuntimeError("min_size cannot be greater than desired_size.")
    if desired > max_size:
        raise RuntimeError("desired_size cannot be greater than max_size.")

    current = get_nodegroup_scaling(cluster, nodegroup, region)
    console.print(
        f"Resuming cluster '[bold]{cluster}[/bold]' by scaling node group "
        f"'[bold]{nodegroup}[/bold]' to desired={desired}, min={min_size}, max={max_size}..."
    )
    console.print(
        f"  Current scaling: desired={current['desired']} min={current['min']} max={current['max']}"
    )
    if (
        current["desired"] == desired
        and current["min"] == min_size
        and current["max"] == max_size
        and current["status"] == "ACTIVE"
    ):
        console.print("[green]Node group already at the requested size.[/green]")
        return
    update_nodegroup_scaling(
        cluster,
        nodegroup,
        region,
        desired=desired,
        min_size=min_size,
        max_size=max_size,
    )
    wait_for_nodegroup_scaling(cluster, nodegroup, region, desired=desired)
    wait_for_k8s_nodegroup_state(nodegroup, expected_ready_schedulable=desired)
    console.print("[green]Environment resumed.[/green]")


def cmd_scale(args) -> None:
    rc, _, err = capture3("aws sts get-caller-identity")
    if rc != 0:
        console.print(
            "[bold red]AWS credentials not valid:[/bold red] "
            f"{err.splitlines()[-1] if err else 'unknown error'}"
        )
        sys.exit(1)

    env_dir = str(Path(args.repo_root) / "terraform/environments/dev")
    region = get_tfvar(env_dir, "aws_region") or "us-west-2"
    cluster = args.cluster or get_eks_cluster_name(env_dir)
    nodegroup = resolve_nodegroup(cluster, region, args.nodegroup_name)
    namespace = get_tf_output(
        env_dir, "tidb_x_cluster_namespace", default="tidb-x-cluster"
    )

    component_targets: dict[str, int] = {}
    terraform_vars: list[str] = []
    for component in SCALE_COMPONENTS:
        value = getattr(args, component["arg"])
        if value is None:
            continue
        if value < 0:
            raise RuntimeError(f"{component['display']} replicas must be non-negative.")
        component_targets[component["arg"]] = value
        terraform_vars.extend(["-var", f"{component['tf_var']}={value}"])

    desired_default, min_default, max_default = get_nodegroup_defaults(env_dir)
    desired = args.desired_size if args.desired_size is not None else desired_default
    min_size = args.min_size if args.min_size is not None else min_default
    max_size = args.max_size if args.max_size is not None else max_default
    node_target_requested = any(
        value is not None for value in (args.desired_size, args.min_size, args.max_size)
    )

    if desired < 0 or min_size < 0 or max_size < 0:
        raise RuntimeError("desired/min/max sizes must be non-negative.")
    if min_size > desired:
        raise RuntimeError("min_size cannot be greater than desired_size.")
    if desired > max_size:
        raise RuntimeError("desired_size cannot be greater than max_size.")
    if not component_targets and not node_target_requested:
        raise RuntimeError(
            "No scale changes requested. Pass component replica flags and/or node size flags."
        )

    current = get_nodegroup_scaling(cluster, nodegroup, region)
    node_expand = node_target_requested and desired > int(current["desired"])
    node_shrink = node_target_requested and desired < int(current["desired"])

    tikv_target = component_targets.get("tikv")
    current_tikv = get_namespaced_resource_json(
        "tikvgroups.core.pingcap.com", "tikv", namespace
    )
    current_tikv_replicas = None
    if current_tikv is not None:
        current_tikv_replicas = int(current_tikv.get("spec", {}).get("replicas", 0))
    if (
        tikv_target is not None
        and current_tikv_replicas is not None
        and tikv_target > current_tikv_replicas
    ):
        ensure_tikv_scale_out_safe(namespace)

    if component_targets:
        summary = ", ".join(
            f"{component['display']}={component_targets[component['arg']]}"
            for component in SCALE_COMPONENTS
            if component["arg"] in component_targets
        )
        console.print(f"[bold]Requested TiDB-X scale:[/bold] {summary}")
    if node_target_requested:
        console.print(
            "[bold]Requested node group scale:[/bold] "
            f"desired={desired} min={min_size} max={max_size}"
        )
        console.print(
            f"  Current scaling: desired={current['desired']} min={current['min']} max={current['max']}"
        )

    def wait_for_requested_node_state() -> None:
        if desired == 0:
            wait_for_k8s_nodegroup_state(nodegroup, expected_registered=0)
        else:
            wait_for_k8s_nodegroup_state(
                nodegroup,
                expected_registered=desired,
                expected_ready_schedulable=desired,
                exact_ready_schedulable=True,
            )

    if node_expand:
        console.print(
            f"[bold]Pre-scaling node group '{nodegroup}'[/bold] to desired={desired}, min={min_size}, max={max_size} before workload expansion..."
        )
        update_nodegroup_scaling(
            cluster,
            nodegroup,
            region,
            desired=desired,
            min_size=min_size,
            max_size=max_size,
        )
        wait_for_nodegroup_scaling(cluster, nodegroup, region, desired=desired)
        wait_for_k8s_nodegroup_state(nodegroup, expected_ready_schedulable=desired)

    if component_targets:
        if tikv_target is not None:
            console.print(
                f"[bold]Synchronizing PD replication max-replicas[/bold] to {tikv_target} before TiKV scale..."
            )
            set_pd_replication_max_replicas(namespace, tikv_target)

        run(tf_cmd("init", "-no-color"), cwd=env_dir)
        run(tf_cmd("validate", "-no-color"), cwd=env_dir)
        console.rule("[bold cyan]terraform apply (scale)[/bold cyan]")
        cmd = [
            *TERRAFORM_CMD,
            "apply",
            "-no-color",
            "-var-file=dev.tfvars",
            *terraform_vars,
            "-auto-approve",
        ]
        rc_apply, apply_log = run_tee(
            cmd,
            cwd=env_dir,
            check=False,
            log_prefix="terraform-scale-apply",
            announce="scale apply",
        )
        if rc_apply != 0:
            console.print(
                f"[bold red]Scale apply failed.[/bold red] Log saved to [bold]{apply_log}[/bold]"
            )
            sys.exit(rc_apply)

        wait_for_scaled_workload(
            namespace,
            component_targets,
            timeout_seconds=args.wait_timeout_seconds,
            stall_timeout_seconds=args.stall_timeout_seconds,
        )

    if node_shrink:
        console.print(
            f"[bold]Post-scaling node group '{nodegroup}'[/bold] to desired={desired}, min={min_size}, max={max_size} after workload shrink..."
        )
        update_nodegroup_scaling(
            cluster,
            nodegroup,
            region,
            desired=desired,
            min_size=min_size,
            max_size=max_size,
        )
        wait_for_nodegroup_scaling(cluster, nodegroup, region, desired=desired)
        wait_for_k8s_nodegroup_state(
            nodegroup,
            expected_registered=desired,
            expected_ready_schedulable=desired,
            exact_ready_schedulable=True,
        )
    elif node_target_requested and not node_expand:
        if (
            current["desired"] != desired
            or current["min"] != min_size
            or current["max"] != max_size
            or current["status"] != "ACTIVE"
        ):
            console.print(
                f"[bold]Scaling node group '{nodegroup}'[/bold] to desired={desired}, min={min_size}, max={max_size}..."
            )
            update_nodegroup_scaling(
                cluster,
                nodegroup,
                region,
                desired=desired,
                min_size=min_size,
                max_size=max_size,
            )
            wait_for_nodegroup_scaling(cluster, nodegroup, region, desired=desired)
            wait_for_requested_node_state()
        else:
            wait_for_requested_node_state()

    console.print("[bold green]Scale complete.[/bold green]")


def cmd_refresh_auth(args) -> None:
    region = args.region
    cluster = args.cluster_name
    profile = args.aws_profile
    sso_region = args.sso_role_region or region
    p = f"--profile {profile}" if profile else ""

    console.print("[bold][1/5] Checking current AWS identity...[/bold]")
    rc, _, _ = capture3(f"aws sts get-caller-identity {p}")
    if rc != 0:
        # Only attempt SSO login if the profile actually has SSO configured.
        # `default` profiles using externally-managed temporary creds will
        # fail with "Missing sso_start_url, sso_region" otherwise.
        rc_cfg, sso_url, _ = capture3(f"aws configure get sso_start_url {p}")
        if rc_cfg == 0 and sso_url:
            console.print("SSO session not valid — running aws sso login...")
            run(f"aws sso login {p}")
        else:
            console.print(
                "[bold red]AWS credentials invalid and profile is not SSO-configured.[/bold red]\n"
                f"Refresh the credentials for profile '{profile}' externally, then re-run."
            )
            sys.exit(1)

    console.print("[bold][2/5] Resolving stable IAM principal ARN...[/bold]")
    current_arn = capture(f"aws sts get-caller-identity {p} --query Arn --output text")
    account_id = capture(
        f"aws sts get-caller-identity {p} --query Account --output text"
    )

    if args.expected_account_id and account_id != args.expected_account_id:
        console.print(
            f"[bold red]Account mismatch: got '{account_id}', expected '{args.expected_account_id}'.[/bold red]\n"
            "Set --aws-profile to the correct profile or omit --expected-account-id."
        )
        sys.exit(1)

    principal_arn = resolve_principal_arn(sso_region, aws_profile=profile)
    console.print(f"  Principal  : {principal_arn}")

    console.print("[bold][3/5] Ensuring EKS access entry exists...[/bold]")
    rc2, _, _ = capture3(
        f"aws eks describe-access-entry --cluster-name {cluster} "
        f"--principal-arn '{principal_arn}' --region {region} {p}"
    )
    if rc2 != 0:
        run(
            f"aws eks create-access-entry --cluster-name {cluster} "
            f"--principal-arn '{principal_arn}' --region {region} {p}"
        )
    else:
        console.print("  Access entry already exists.")

    console.print("[bold][4/5] Ensuring cluster admin policy is associated...[/bold]")
    _, policies, _ = capture3(
        f"aws eks list-associated-access-policies --cluster-name {cluster} "
        f"--principal-arn '{principal_arn}' --region {region} {p} "
        f"--query 'associatedAccessPolicies[].policyArn' --output text"
    )
    if EKS_ADMIN_POLICY not in (policies or ""):
        run(
            f"aws eks associate-access-policy --cluster-name {cluster} "
            f"--principal-arn '{principal_arn}' --policy-arn {EKS_ADMIN_POLICY} "
            f"--access-scope type=cluster --region {region} {p}",
            check=False,
        )
    else:
        console.print("  Admin policy already associated.")

    console.print("[bold][5/5] Updating kubeconfig...[/bold]")
    run(f"aws eks update-kubeconfig --name {cluster} --region {region} {p}")
    run("kubectl auth can-i '*' '*' --all-namespaces")
    console.print(
        f"[bold green]EKS auth refreshed for cluster '{cluster}'.[/bold green]"
    )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description="TiDB one-box skill.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--repo-root",
        default=str(Path(__file__).parent),
        help="Path to the one-box repo root (default: directory containing this script)",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # bootstrap
    sub.add_parser("bootstrap", help="init + validate + plan (no apply)")

    # deploy
    p_deploy = sub.add_parser("deploy", help="Terraform apply for the full stack.")
    p_deploy.add_argument(
        "--auto-approve",
        action="store_true",
        help="Skip prompts (implied; deploy always auto-approves)",
    )

    # destroy
    p_destroy = sub.add_parser(
        "destroy", help="Destroy all resources with automatic remediations."
    )
    p_destroy.add_argument(
        "--auto-approve", action="store_true", help="Skip 'destroy' confirmation prompt"
    )
    p_destroy.add_argument(
        "--force-cleanup-elb",
        action="store_true",
        help="Delete classic ELBs before destroy",
    )
    p_destroy.add_argument(
        "--sso-role-region",
        default=None,
        help="SSO role region (default: same as aws_region)",
    )
    p_destroy.add_argument(
        "--also-delete-external-bucket",
        action="store_true",
        help="Also empty and delete the configured s3_bucket_name even if it "
        "is not Terraform-managed (default: leave external buckets intact).",
    )

    # validate
    p_val = sub.add_parser("validate", help="Check TiDB-X cluster health.")
    p_val.add_argument("--namespace", default="tidb-x-cluster")
    p_val.add_argument("--cluster", default="tidb-x-onebox")
    p_val.add_argument("--status-service", default="tidb-tidb")

    # verify-sql
    p_sql = sub.add_parser(
        "verify-sql", help="Verify SQL pushdown to TiKV via pymysql."
    )
    p_sql.add_argument("--namespace", default="tidb-x-cluster")
    p_sql.add_argument("--service", default="tidb-tidb")
    p_sql.add_argument(
        "--port",
        type=int,
        default=4000,
        help="Local port for port-forward (default: 4000; use 6000 for TiProxy)",
    )
    p_sql.add_argument("--user", default="root")
    p_sql.add_argument("--password", default="")

    # stop
    p_stop = sub.add_parser(
        "stop", help="Pause the environment by scaling node group to zero."
    )
    p_stop.add_argument(
        "--cluster",
        default=None,
        help="EKS cluster name (default: inferred from dev.tfvars)",
    )
    p_stop.add_argument(
        "--nodegroup-name",
        default=None,
        help="Managed node group name (default: auto-detect if exactly one exists)",
    )

    # start
    p_start = sub.add_parser(
        "start", help="Resume the environment by restoring node group size."
    )
    p_start.add_argument(
        "--cluster",
        default=None,
        help="EKS cluster name (default: inferred from dev.tfvars)",
    )
    p_start.add_argument(
        "--nodegroup-name",
        default=None,
        help="Managed node group name (default: auto-detect if exactly one exists)",
    )
    p_start.add_argument(
        "--desired-size",
        type=int,
        default=None,
        help="Desired node count (default: node_desired_size from dev.tfvars)",
    )
    p_start.add_argument(
        "--min-size",
        type=int,
        default=None,
        help="Minimum node count (default: node_min_size from dev.tfvars)",
    )
    p_start.add_argument(
        "--max-size",
        type=int,
        default=None,
        help="Maximum node count (default: node_max_size from dev.tfvars)",
    )

    # scale
    p_scale = sub.add_parser(
        "scale",
        help="Scale TiDB-X replicas and/or EKS node group size.",
    )
    p_scale.add_argument(
        "--cluster",
        default=None,
        help="EKS cluster name (default: inferred from dev.tfvars)",
    )
    p_scale.add_argument(
        "--nodegroup-name",
        default=None,
        help="Managed node group name (default: auto-detect if exactly one exists)",
    )
    p_scale.add_argument(
        "--desired-size",
        type=int,
        default=None,
        help="Desired node count override for this scale operation",
    )
    p_scale.add_argument(
        "--min-size",
        type=int,
        default=None,
        help="Minimum node count override for this scale operation",
    )
    p_scale.add_argument(
        "--max-size",
        type=int,
        default=None,
        help="Maximum node count override for this scale operation",
    )
    p_scale.add_argument(
        "--wait-timeout-seconds",
        type=int,
        default=900,
        help="Max time to wait for workload convergence (default: 900)",
    )
    p_scale.add_argument(
        "--stall-timeout-seconds",
        type=int,
        default=180,
        help="Fail early if scale status stops changing for this long (default: 180)",
    )
    p_scale.add_argument("--pd", type=int, default=None, help="PD replicas")
    p_scale.add_argument("--tikv", type=int, default=None, help="TiKV replicas")
    p_scale.add_argument("--tidb", type=int, default=None, help="TiDB replicas")
    p_scale.add_argument("--tso", type=int, default=None, help="TSO replicas")
    p_scale.add_argument(
        "--scheduling", type=int, default=None, help="Scheduling replicas"
    )
    p_scale.add_argument("--tiproxy", type=int, default=None, help="TiProxy replicas")
    p_scale.add_argument(
        "--tiflash-write",
        dest="tiflash_write",
        type=int,
        default=None,
        help="TiFlash write replicas",
    )
    p_scale.add_argument(
        "--tiflash-compute",
        dest="tiflash_compute",
        type=int,
        default=None,
        help="TiFlash compute replicas",
    )

    # refresh-auth
    p_auth = sub.add_parser(
        "refresh-auth", help="Refresh AWS SSO + EKS access entry + kubeconfig."
    )
    p_auth.add_argument("--aws-profile", default="default")
    p_auth.add_argument("--region", default="us-west-2")
    p_auth.add_argument("--cluster-name", default="tidb-onebox-dev")
    p_auth.add_argument("--sso-role-region", default=None)
    p_auth.add_argument(
        "--expected-account-id",
        default=None,
        help="Fail fast if the resolved account ID does not match",
    )

    args = parser.parse_args()

    {
        "bootstrap": cmd_bootstrap,
        "deploy": cmd_deploy,
        "destroy": cmd_destroy,
        "scale": cmd_scale,
        "validate": cmd_validate,
        "verify-sql": cmd_verify_sql,
        "stop": cmd_stop,
        "start": cmd_start,
        "refresh-auth": cmd_refresh_auth,
    }[args.command](args)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        console.print("\n[yellow]Interrupted.[/yellow]")
        sys.exit(1)
    except (RuntimeError, subprocess.CalledProcessError) as e:
        console.print(f"[bold red]Error:[/bold red] {e}")
        sys.exit(1)