import base64
import builtins
import io
import json
import shutil
import subprocess
import tarfile
import tempfile
import time
from pathlib import Path
from typing import Any, Optional, Union
from uuid import UUID, uuid4

import httpx
import pydantic
import srsly
from pydantic import AnyUrl
from pydantic_settings import BaseSettings
from radicli import Arg
from wasabi import msg

from ..cli import cli
from ..errors import (
    ConnectError,
    ConnectTimeout,
    HTTPStatusError,
    InvalidURL,
    RequestError,
)
from ..messages import Messages
from ..prodigy_teams_broker_sdk.models import CheckProgressRequest, CheckStartRequest
from ..prodigy_teams_pam_sdk import Client
from ..prodigy_teams_pam_sdk import models as prodigy_teams_pam_sdk_models
from ..prodigy_teams_pam_sdk.errors import BrokerNotFound
from ..prodigy_teams_pam_sdk.models import (
    BrokerCreating,
    BrokerCredentials,
    BrokerUpdating,
)
from ..query import resolve_broker, resolve_broker_id, resolve_package, resolve_recipe
from ..ty import BrokerStatusCheck
from ..ui import print_info_table, print_table_with_select
from ..util import URL
from ._state import get_auth_state

CLOUD_PROVIDERS = ("gcp", "aws", "k3s")
CONTAINER_REGISTRY = "europe-west1-docker.pkg.dev/pt-pam-prod/pam-external"
DEFAULT_NAMESPACE = "prodigy-teams"
DEFAULT_RELEASE = "prodigy-teams"

# Broker has wrong name on infra
# TODO: Basic crud stuff here for clusters


class Config(BaseSettings):
    class Terraform(BaseSettings):
        class TerraformCloud(BaseSettings):
            workspace: str
            organization: str

        class GCS(BaseSettings):
            bucket: str
            prefix: str

        terraform_cloud: TerraformCloud | None = None
        gcs: GCS | None = None

    class Deployment(BaseSettings):
        project: str
        zone: str
        domain: str
        public_ip: str = ""
        pam_domain: str
        container_registry: str
        builtin_recipes_version: str = ""
        cluster_name: str = ""
        image_tags: dict[str, str] = {}
        gateway_name: str = ""
        gateway_namespace: str = ""
        tls_secret_name: str = ""
        acme_email: str = ""
        prodigy_pypi_index_url: str = ""
        database_ip: str = ""
        database_user: str = ""
        database_name: str = ""
        nfs_pvc_name: str = ""
        infra_secret_name: str = ""

    class NodePool(BaseSettings):
        class GPU(BaseSettings):
            type: str
            count: int

        name: str
        machine_type: str
        target_size: int
        preemptible: bool
        node_class: str
        gpu: GPU | None = None
        per_job: dict

    name: str
    terraform: Terraform
    deployment: Deployment
    node_pools: builtins.list[NodePool] = []


def check_broker_status(
    broker: prodigy_teams_pam_sdk_models.BrokerSummary,
) -> BrokerStatusCheck:
    url = URL.parse(broker.address) / "api/v1/status"
    # Before querying the broker, verify that it has finished registration
    # TODO: also track when it is killed in PAM
    if broker.state == "creating":
        return BrokerStatusCheck.CREATING
    try:
        response = httpx.get(str(url))
        response.raise_for_status()
        payload = response.json()
        # Happy case: successfully connect to the broker and check "cluster" key
        if payload.get("cluster") == "Ready":
            return BrokerStatusCheck.RUNNING
        else:
            return BrokerStatusCheck.ISSUES
    except InvalidURL:
        return BrokerStatusCheck.INVALID_ADDRESS
    except (ConnectTimeout, ConnectError):
        # This could mean the broker is down, but could also mean we have no network access
        return BrokerStatusCheck.NOT_FOUND
    except RequestError as e:
        msg.warn(Messages.E041.format(name=broker.name), str(e))
        return BrokerStatusCheck.REQUEST_ERROR
    except HTTPStatusError as e:
        # We shouldn't really log here -- better to refactor this into the commands
        msg.warn(Messages.E041.format(name=broker.name), str(e.response.status_code))
        return BrokerStatusCheck.RESPONSE_ERROR


class BrokerSummary(prodigy_teams_pam_sdk_models.BrokerSummary):
    """
    Extends the data returned by the API with computed properties.
    This is a workaround for the models being auto-generated.
    """

    @property
    def status(self) -> str:
        status = check_broker_status(self)
        return str(status)

    @classmethod
    def get_properties(cls) -> builtins.list[str]:
        return [
            prop for prop in cls.__dict__ if isinstance(cls.__dict__[prop], property)
        ]

    def model_dump(self, **kwargs: Any) -> dict[str, Any]:  # type: ignore[override]
        include = kwargs.get("include")
        exclude = kwargs.get("exclude")
        data = super().model_dump(**kwargs)
        props = self.get_properties()
        if include:
            props = [prop for prop in props if prop in include]
        if exclude:
            props = [prop for prop in props if prop not in exclude]
        for key in props:
            data[key] = getattr(self, key)
        return data


class BrokerDetail(BrokerSummary, prodigy_teams_pam_sdk_models.BrokerDetail):
    pass


@cli.subcommand(
    "clusters",
    "list",
    # fmt: off
    select=Arg("--select", help=Messages.select.format(opts=builtins.list(BrokerSummary.model_fields))),
    as_json=Arg("--json", help=Messages.as_json),
    # fmt: on
)
def list(
    select: builtins.list[str] = ["id", "name", "status", "address"],
    as_json: bool = False,
) -> builtins.list[BrokerSummary]:
    """List resources on the cluster"""
    client = get_auth_state().client
    # TODO: this sequentially checks broker status which could be slow
    res = [
        BrokerSummary(**x.model_dump())
        for x in client.broker.all(page=1, size=100).items
    ]
    print_table_with_select(res, select=select, as_json=as_json)
    return res


@cli.subcommand(
    "clusters",
    "info",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="cluster")),
    select=Arg(
        "--select",
        help=Messages.select.format(opts=builtins.list(BrokerDetail.model_fields)),
    ),
    as_json=Arg("--json", help=Messages.as_json),
)
def info(
    name_or_id: Union[str, UUID],
    select: Optional[builtins.list[str]] = None,
    as_json: bool = False,
) -> BrokerDetail:
    """Get detailed info for a cluster"""
    cluster = BrokerDetail(**resolve_broker(name_or_id).model_dump())
    select = [*BrokerDetail.model_fields.keys(), *BrokerDetail.get_properties()]
    print_info_table(cluster, select=select, as_json=as_json)
    return cluster


@cli.subcommand(
    "clusters",
    "update",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="cluster")),
    name=Arg("--new-name", help=Messages.new_name.format(noun="cluster")),
    address=Arg("--address", help=Messages.new_address.format(noun="cluster")),
    as_json=Arg("--json", help=Messages.as_json),
)
def update(
    name_or_id: Union[str, UUID],
    name: Optional[str] = None,
    address: Optional[str] = None,
    as_json: bool = False,
) -> UUID:
    """Update the cluster info"""
    cluster = resolve_broker(name_or_id)
    # TODO: The BrokerUpdating model is currently wrong on pam, making these
    # fields required. Change the query when this is fixed.
    body = BrokerUpdating(
        id=cluster.id,
        name=name or cluster.name,
        address=address or cluster.address,
        public_key=None,
        cloud_provider=None,
        cloud=None,
        state=None,
    )
    auth = get_auth_state()
    res = auth.client.broker.update(body)
    print_info_table(res, as_json=as_json)
    return cluster.id


@cli.subcommand(
    "clusters",
    "delete",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="cluster")),
)
def delete(name_or_id: Union[str, UUID]) -> UUID:
    """
    Delete a cluster from PAM. This only removes PAM's record
    of it. The cluster itself will continue to exist - you need
    to shut it down separately.
    """
    auth = get_auth_state()
    client = auth.client
    cluster_id = resolve_broker_id(name_or_id)
    client.broker.delete(id=cluster_id)
    msg.good(Messages.T003.format(noun="cluster", name=cluster_id))
    return cluster_id


def _check_result(label: str, ok: bool, detail: str = "") -> tuple[str, bool]:
    """Print a check result line and return (label, ok)."""
    symbol = "\u2714" if ok else "\u2718"
    suffix = f" ({detail})" if detail else ""
    if ok:
        msg.good(f"{symbol} {label}{suffix}")
    else:
        msg.fail(f"{symbol} {label}{suffix}")
    return (label, ok)


@cli.subcommand(
    "clusters",
    "check",
    name_or_id=Arg("--cluster", help=Messages.name_or_id.format(noun="cluster")),
    s3_bucket=Arg("--s3-bucket", help=Messages.check_s3_bucket),
    nfs_path=Arg("--nfs-path", help=Messages.check_nfs_path),
    recipe_name_or_id=Arg("--recipe", help=Messages.check_recipe),
    recipe_args=Arg("--recipe-args", help=Messages.check_recipe_args),
    deep=Arg("--deep", help="Run broker-side deployment checks (K8s, NFS, DB)"),
)
def check(
    name_or_id: Optional[Union[str, UUID]] = None,
    s3_bucket: Optional[str] = None,
    nfs_path: Optional[str] = None,
    recipe_name_or_id: Optional[Union[str, UUID]] = None,
    recipe_args: Optional[str] = None,
    deep: bool = False,
) -> None:
    """Check the health of a cluster deployment.

    Runs CLI-side connectivity checks against the broker and PAM.
    Use --deep to also trigger broker-side deployment checks (K8s API,
    NFS, database, etc.).
    """
    auth = get_auth_state()
    results: builtins.list[tuple[str, bool]] = []
    # Resolve the broker to check
    if name_or_id:
        broker = resolve_broker(name_or_id)
    else:
        brokers = auth.client.broker.all(page=1, size=1).items
        if not brokers:
            msg.fail("No clusters found. Register one with 'ptc clusters register'.")
            raise SystemExit(1)
        broker = brokers[0]
    broker_url = broker.address.rstrip("/")
    # 1. Broker API responding
    try:
        r = httpx.get(f"{broker_url}/api/v1/status", timeout=10)
        r.raise_for_status()
        payload = r.json()
        api_ok = payload.get("status") == "Ready"
        cluster_ok = payload.get("cluster") == "Ready"
        results.append(_check_result(f"Broker API responding on {broker_url}", api_ok))
        # 2. Cluster status (K8s connectivity from broker)
        if api_ok:
            results.append(
                _check_result(
                    "Cluster healthy (K8s API accessible from broker)",
                    cluster_ok,
                    detail=payload.get("cluster", "unknown"),
                )
            )
    except (httpx.ConnectError, httpx.ConnectTimeout):
        results.append(
            _check_result(
                f"Broker API responding on {broker_url}",
                False,
                detail="connection failed",
            )
        )
    except httpx.HTTPStatusError as exc:
        results.append(
            _check_result(
                f"Broker API responding on {broker_url}",
                False,
                detail=f"HTTP {exc.response.status_code}",
            )
        )
    except Exception as exc:
        results.append(
            _check_result(
                f"Broker API responding on {broker_url}", False, detail=str(exc)
            )
        )
    # 3. Database connected — check via kubectl if available
    db_ok = False
    try:
        pod_result = subprocess.run(
            [
                "kubectl",
                "get",
                "pods",
                "--namespace",
                DEFAULT_NAMESPACE,
                "-l",
                "app.kubernetes.io/name=prodigy-teams-broker",
                "-o",
                "jsonpath={.items[0].status.containerStatuses[0].ready}",
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if pod_result.returncode == 0 and pod_result.stdout.strip() == "true":
            db_ok = True
        results.append(_check_result("Database connected", db_ok))
    except (FileNotFoundError, subprocess.TimeoutExpired):
        # kubectl not available — skip this check
        pass
    # 4. NFS storage — check PVC is bound
    try:
        pvc_result = subprocess.run(
            [
                "kubectl",
                "get",
                "pvc",
                "--namespace",
                DEFAULT_NAMESPACE,
                "-o",
                "jsonpath={.items[*].status.phase}",
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if pvc_result.returncode == 0 and pvc_result.stdout.strip():
            phases = pvc_result.stdout.strip().split()
            all_bound = all(p == "Bound" for p in phases)
            results.append(
                _check_result(
                    "NFS storage mounted",
                    all_bound,
                    detail=f"{len(phases)} PVC(s) {'all bound' if all_bound else ', '.join(phases)}",
                )
            )
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    # 5. PAM sync active — check broker state in PAM
    pam_ok = broker.state not in ("creating", "error")
    results.append(_check_result("PAM sync active", pam_ok, detail=broker.state))
    # 6. Provision job — check most recent job status
    try:
        job_result = subprocess.run(
            [
                "kubectl",
                "get",
                "jobs",
                "--namespace",
                DEFAULT_NAMESPACE,
                "-l",
                "app=provision",
                "-o",
                "jsonpath={.items[-1:].status.conditions[0].type}",
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if job_result.returncode == 0 and job_result.stdout.strip():
            job_status = job_result.stdout.strip()
            provision_ok = job_status == "Complete"
            results.append(
                _check_result(
                    "Provision job completed", provision_ok, detail=job_status
                )
            )
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    # Deep checks: delegate to broker-side check API
    if deep:
        msg.divider("Broker-side deployment checks")
        if recipe_args:
            recipe_args_dict = json.loads(recipe_args)
            assert isinstance(recipe_args_dict, dict)
        else:
            recipe_args_dict = {}
        if recipe_name_or_id:
            recipe = resolve_recipe(recipe_name_or_id, broker_id=None)
            package = resolve_package(recipe.package_id, broker_id=None)
            recipe_name = recipe.name
            package_environment = package.environment
        else:
            recipe_name = None
            package_environment = None
        start_body = CheckStartRequest(
            s3_bucket=s3_bucket,
            nfs_path=nfs_path,
            package_environment=package_environment,
            recipe_name=recipe_name,
            recipe_args=recipe_args_dict,
        )
        start_response = auth.broker_client.check.start(start_body)
        while True:
            progress_body = CheckProgressRequest(id=start_response.id)
            progress_response = auth.broker_client.check.progress(progress_body)
            if progress_response.status == "done":
                assert progress_response.report
                for check_name, result in progress_response.report.items():
                    ok = result == "success"
                    skipped = result == "skip"
                    if skipped:
                        msg.info(f"- {check_name}: skipped")
                    else:
                        results.append(_check_result(check_name, ok))
                break
            time.sleep(1)
    # Summary
    failed = [label for label, ok in results if not ok]
    if failed:
        raise SystemExit(1)


@cli.subcommand(
    "clusters",
    "nodes",
    # fmt: off
    select=Arg("--select", help=Messages.select.format(opts=["name", "status", "cpu", "memory", "gpu", "pod_count"])),
    as_json=Arg("--json", help=Messages.as_json),
    # fmt: on
)
def nodes(
    select: builtins.list[str] = [
        "name",
        "status",
        "cpu",
        "memory",
        "gpu",
        "pod_count",
    ],
    as_json: bool = False,
) -> builtins.list[dict[str, Any]]:
    """List cluster nodes with capacity and utilization."""
    auth = get_auth_state()
    node_list = auth.broker_client.nodes.list()
    rows = [
        {
            "name": n.name,
            "status": n.status,
            "cpu": n.capacity.cpu,
            "memory": n.capacity.memory,
            "gpu": str(n.capacity.gpu) if n.capacity.gpu is not None else "-",
            "pod_count": str(n.pod_count),
        }
        for n in node_list
    ]
    print_table_with_select(rows, select=select, as_json=as_json)
    return rows


@cli.subcommand(
    "clusters",
    "start",
    src=Arg(help="Path to cluster definition files"),
    config_path=Arg(help="Path to config yaml"),
    config_key=Arg(
        "--config-key",
        help="Top-level key for the cluster config. Allows multiple cluster configurations to be kept in one file. If not set, the whole config file is used.",
    ),
    offline=Arg(
        "--offline",
        help="Path to a credentials JSON file (container_sa_key, prodigy_license_key). Deploys without connecting to the management server.",
    ),
)
def gcp_cluster_start(
    src: Union[Path, str] = ".",
    config_path: Optional[Path] = None,
    config_key: Optional[str] = None,
    offline: Optional[Path] = None,
) -> None:
    """Start a cluster on infrastructure configured under your own cloud account.

    See 'ptc gcp' to create the infrastructure.
    """
    if src == ".":
        src = Path.cwd()
    elif isinstance(src, str):
        src = Path(src)
    if config_path is None:
        config_path = src / "config.yaml"
    config_dict = srsly.read_yaml(config_path)
    assert isinstance(config_dict, dict)
    if config_key is not None:
        config_dict = config_dict[config_key]
    config = Config(**config_dict)
    if config.deployment.builtin_recipes_version:
        recipes_version = config.deployment.builtin_recipes_version
    else:
        recipes_version = _get_latest_version(
            "https://pypi.org/pypi/prodigy-teams-recipes/json"
        )
    print("Using recipes version", recipes_version)
    cluster_name = config.deployment.cluster_name
    assert cluster_name, "cluster_name must be set in deployment config"
    _get_cluster_credentials(
        config.deployment.project, config.deployment.zone, cluster_name
    )
    broker_public_key_pem = _ensure_broker_keypair(config)
    if offline is not None:
        creds_data = srsly.read_json(offline)
        assert isinstance(creds_data, dict)
        cluster_auth = BrokerCredentials(
            broker_id=(
                UUID(creds_data["broker_id"]) if "broker_id" in creds_data else uuid4()
            ),
            container_sa_key=creds_data.get("container_sa_key", ""),
            prodigy_license_key=creds_data.get("prodigy_license_key", ""),
            public_key=broker_public_key_pem,
        )
    else:
        auth = get_auth_state()
        ensure_cluster_record(auth.client, auth.org_id, config, broker_public_key_pem)
        cluster_auth = get_cluster_auth(auth.client, config.deployment.domain)
    values = _build_helm_values(
        config,
        creds=cluster_auth,
        builtin_recipes_version=recipes_version,
    )
    msg.info("Running preflight checks...")
    if not _helm_test():
        msg.fail(
            "Preflight checks failed — fix infrastructure issues before deploying."
        )
        raise SystemExit(1)
    if config.deployment.gateway_name:
        _ensure_gateway_crds()
        _ensure_traefik(public_ip=config.deployment.public_ip)
        if config.deployment.acme_email:
            _ensure_cert_manager()
    _helm_upgrade(src / "k8s" / "chart", values=values)


@cli.subcommand(
    "clusters",
    "provision",
    src=Arg(help="Path to cluster definition files"),
    config_path=Arg(help="Path to config yaml"),
    config_key=Arg(
        "--config-key",
        help="Top-level key for the cluster config. Allows multiple cluster configurations to be kept in one file. If not set, the whole config file is used.",
    ),
)
def gcp_cluster_provision(
    src: Union[Path, str] = ".",
    config_path: Optional[Path] = None,
    config_key: Optional[str] = None,
) -> None:
    """Provision NFS wheelhouse on a GKE cluster.

    See 'ptc gcp' to create the infrastructure.
    """
    if src == ".":
        src = Path.cwd()
    elif isinstance(src, str):
        src = Path(src)
    auth = get_auth_state()
    if config_path is None:
        config_path = src / "config.yaml"
    config_dict = srsly.read_yaml(config_path)
    assert isinstance(config_dict, dict)
    if config_key is not None:
        config_dict = config_dict[config_key]
    config = Config(**config_dict)
    recipes_version = _get_latest_version(
        "https://pypi.org/pypi/prodigy-teams-recipes/json"
    )
    print("Using recipes version", recipes_version)
    cluster_name = config.deployment.cluster_name
    assert cluster_name, "cluster_name must be set in deployment config"
    _get_cluster_credentials(
        config.deployment.project, config.deployment.zone, cluster_name
    )
    _ensure_broker_keypair(config)
    cluster_auth = get_cluster_auth(auth.client, config.deployment.domain)
    values = _build_helm_values(
        config,
        creds=cluster_auth,
        builtin_recipes_version=recipes_version,
    )
    msg.info("Running preflight checks...")
    if not _helm_test():
        msg.fail(
            "Preflight checks failed — fix infrastructure issues before deploying."
        )
        raise SystemExit(1)
    if config.deployment.gateway_name:
        _ensure_gateway_crds()
        _ensure_traefik(public_ip=config.deployment.public_ip)
        if config.deployment.acme_email:
            _ensure_cert_manager()
    msg.info("Running helm upgrade (provision hook will execute)...")
    _helm_upgrade(src / "k8s" / "chart", values=values, wait=True)


def _get_latest_version(url: str) -> str:
    r = httpx.get(url)
    r.raise_for_status()
    return r.json()["info"]["version"]


def _ensure_broker_keypair(config: Config) -> str:
    """Ensure the infra K8s secret has a broker RSA keypair, return the public PEM.

    If the secret already exists and contains a private key, the existing public
    key is returned.  Otherwise a new keypair is generated, and the secret is
    created/patched with the keypair plus a placeholder database password (the
    real password is managed by Helm values).
    """
    from ..key_pair import generate_key_pair

    secret_name = config.deployment.infra_secret_name
    namespace = "prodigy-teams"

    # Try to read the existing public key from the secret.
    check = subprocess.run(
        [
            "kubectl",
            "get",
            "secret",
            secret_name,
            "-n",
            namespace,
            "-o",
            "jsonpath={.data.broker_public_key}",
        ],
        capture_output=True,
        text=True,
    )
    if check.returncode == 0 and check.stdout.strip():
        public_pem = base64.b64decode(check.stdout.strip()).decode("utf-8")
        msg.info(f"Using existing broker keypair from secret {secret_name}")
        return public_pem

    # Generate a new keypair and create/patch the secret.
    msg.info("Generating new broker RSA keypair...")
    kp = generate_key_pair()
    cmd = [
        "kubectl",
        "create",
        "secret",
        "generic",
        secret_name,
        f"--namespace={namespace}",
        f"--from-literal=broker_private_key={kp.private_pem}",
        f"--from-literal=broker_public_key={kp.public_pem}",
        "--dry-run=client",
        "-o",
        "yaml",
    ]
    create = subprocess.run(cmd, capture_output=True, text=True, check=True)
    subprocess.run(
        ["kubectl", "apply", "-f", "-"],
        input=create.stdout,
        check=True,
        text=True,
    )
    msg.good(f"Created broker keypair in secret {secret_name}")
    return kp.public_pem


def _build_helm_values(
    config: Config,
    creds: BrokerCredentials,
    builtin_recipes_version: str,
) -> dict:
    """Build a Helm values dict from config, credentials and terraform outputs."""
    image_values: dict = {
        "registry": config.deployment.container_registry,
    }
    for service, tag in config.deployment.image_tags.items():
        image_values[service] = {"tag": tag}
    values: dict = {
        "image": image_values,
        "cluster": {
            "fqdn": config.deployment.domain,
            "pamFqdn": config.deployment.pam_domain,
        },
        "gateway": {
            "name": config.deployment.gateway_name,
            "namespace": config.deployment.gateway_namespace,
            "tls": (
                [{"secretName": config.deployment.tls_secret_name}]
                if config.deployment.tls_secret_name
                else []
            ),
        },
    }
    # Only set GCP values when a GCP project is configured
    if config.deployment.project:
        values["gcp"] = {
            "project": config.deployment.project,
            "zone": config.deployment.zone,
        }
        values["serviceAccount"] = {
            "annotations": {
                "iam.gke.io/gcp-service-account": f"cluster-gke-sa@{config.deployment.project}.iam.gserviceaccount.com",
            },
        }
    values["workerTypes"] = json.dumps(
        [
            {
                "name": w.node_class,
                "node_class": w.node_class,
                "description": w.name,
                "image_id": "gpu" if w.gpu else "cpu",
                "per_job": w.per_job,
            }
            for w in config.node_pools
        ]
    )
    if config.deployment.prodigy_pypi_index_url:
        values["prodigyPypiIndexUrl"] = config.deployment.prodigy_pypi_index_url
    values["builtinRecipesVersion"] = builtin_recipes_version
    values["cluster"]["brokerId"] = str(creds.broker_id)
    values.setdefault("secrets", {})
    values["secrets"]["containerSaKey"] = creds.container_sa_key
    values["secrets"]["prodigyLicenseKey"] = creds.prodigy_license_key
    values["database"] = {
        "ip": config.deployment.database_ip,
        "user": config.deployment.database_user,
        "name": config.deployment.database_name,
    }
    values["nfs"] = {"claimName": config.deployment.nfs_pvc_name}
    values["secrets"]["infraSecretName"] = config.deployment.infra_secret_name
    if config.deployment.tls_secret_name and config.deployment.acme_email:
        values["certManager"] = {
            "enabled": True,
            "secretName": config.deployment.tls_secret_name,
            "acmeEmail": config.deployment.acme_email,
        }
    return values


def _get_cluster_credentials(project: str, zone: str, cluster_name: str) -> None:
    """Run gcloud container clusters get-credentials to configure kubectl."""
    cmd = [
        "gcloud",
        "container",
        "clusters",
        "get-credentials",
        cluster_name,
        "--zone",
        zone,
        "--project",
        project,
    ]
    subprocess.run(cmd, check=True, text=True)


def _helm_test(
    release: str = "prodigy-teams",
    namespace: str = "prodigy-teams",
    timeout: str = "2m",
) -> bool:
    """Run helm test preflight checks. Returns True if all pass."""
    # Skip if release doesn't exist yet (first install) or is in a failed state
    # (tests run against the current release values, which may be broken)
    check = subprocess.run(
        ["helm", "status", release, "--namespace", namespace, "-o", "json"],
        capture_output=True,
        text=True,
    )
    if check.returncode != 0:
        return True  # no release yet, skip preflight
    status = json.loads(check.stdout).get("info", {}).get("status", "")
    if status == "failed":
        msg.warn("Current release is in failed state — skipping preflight checks")
        return True
    cmd = [
        "helm",
        "test",
        release,
        "--namespace",
        namespace,
        "--timeout",
        timeout,
    ]
    result = subprocess.run(cmd, text=True)
    return result.returncode == 0


def _ensure_gateway_crds() -> None:
    """Install Gateway API CRDs if not already present."""
    check = subprocess.run(
        ["kubectl", "get", "crd", "gateways.gateway.networking.k8s.io"],
        capture_output=True,
        text=True,
    )
    if check.returncode == 0:
        return
    msg.info("Installing Gateway API CRDs...")
    subprocess.run(
        [
            "kubectl",
            "apply",
            "-f",
            "https://github.com/kubernetes-sigs/gateway-api/releases/download/v1.2.1/standard-install.yaml",
        ],
        check=True,
        text=True,
    )


def _ensure_traefik(public_ip: str = "") -> None:
    """Install Traefik via Helm if not already present."""
    check = subprocess.run(
        ["helm", "status", "traefik", "-n", "traefik"],
        capture_output=True,
        text=True,
    )
    if check.returncode == 0:
        return
    msg.info("Installing Traefik...")
    subprocess.run(
        ["helm", "repo", "add", "traefik", "https://traefik.github.io/charts"],
        check=True,
        text=True,
    )
    subprocess.run(["helm", "repo", "update", "traefik"], check=True, text=True)
    values: dict = {
        "providers": {
            "kubernetesGateway": {"enabled": True},
            "kubernetesIngress": {"enabled": False},
        },
        "ingressRoute": {"dashboard": {"enabled": False}},
        "gateway": {"enabled": False},
        "ports": {"web": {"port": 80}, "websecure": {"port": 443}},
        "securityContext": {
            "capabilities": {"drop": [], "add": ["NET_BIND_SERVICE"]},
        },
    }
    if public_ip:
        values["service"] = {"spec": {"loadBalancerIP": public_ip}}
    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
        tmp.write(json.dumps(values))
        tmp_path = tmp.name
    try:
        subprocess.run(
            [
                "helm",
                "upgrade",
                "--install",
                "traefik",
                "traefik/traefik",
                "-n",
                "traefik",
                "--create-namespace",
                "--skip-crds",
                "-f",
                tmp_path,
            ],
            check=True,
            text=True,
        )
    finally:
        Path(tmp_path).unlink(missing_ok=True)
    if public_ip:
        msg.good(f"Traefik installed with static IP {public_ip}")
    else:
        msg.info(
            "Traefik installed. Get the LoadBalancer IP with:\n"
            "  kubectl get svc -n traefik traefik"
            " -o jsonpath='{.status.loadBalancer.ingress[0].ip}'"
        )


def _ensure_cert_manager() -> None:
    """Install cert-manager via Helm if not already present."""
    check = subprocess.run(
        ["helm", "status", "cert-manager", "-n", "cert-manager"],
        capture_output=True,
        text=True,
    )
    if check.returncode == 0:
        return
    msg.info("Installing cert-manager...")
    subprocess.run(
        ["helm", "repo", "add", "jetstack", "https://charts.jetstack.io"],
        check=True,
        text=True,
    )
    subprocess.run(["helm", "repo", "update", "jetstack"], check=True, text=True)
    subprocess.run(
        [
            "helm",
            "upgrade",
            "--install",
            "cert-manager",
            "jetstack/cert-manager",
            "-n",
            "cert-manager",
            "--create-namespace",
            "--set",
            "crds.enabled=true",
            "--set",
            "crds.keep=true",
            "--set",
            "config.enableGatewayAPI=true",
        ],
        check=True,
        text=True,
    )


def _helm_upgrade(
    chart_path: Path,
    values: dict,
    release: str = "prodigy-teams",
    namespace: str = "prodigy-teams",
    wait: bool = False,
    timeout: str = "10m",
) -> None:
    """Run helm upgrade --install for the prodigy-teams chart."""
    subprocess.run(
        ["helm", "dependency", "build", str(chart_path)],
        check=True,
        text=True,
    )
    cmd = [
        "helm",
        "upgrade",
        "--install",
        release,
        str(chart_path),
        "--namespace",
        namespace,
        "--create-namespace",
        "--server-side=true",
        "--force-conflicts",
    ]
    cmd += ["--timeout", timeout]
    if wait:
        cmd += ["--wait"]
    values_path = chart_path / ".helm-values-override.json"
    values_path.write_text(json.dumps(values, indent=2))
    cmd += ["-f", str(values_path)]
    subprocess.run(cmd, check=True, text=True)


def get_cluster_auth(client: Client, domain: str) -> BrokerCredentials:
    broker = client.broker.read(address=domain)
    return client.broker.credentials(broker_id=broker.id)


def ensure_cluster_record(
    client: Client,
    org_id: UUID,
    config: Config,
    public_key: str,
    *,
    cluster_name: str = "Test Cluster",
    cloud_provider: str = "gcp",
) -> None:
    domain = config.deployment.domain
    try:
        broker = client.broker.read(address=domain)
    except BrokerNotFound:
        body = BrokerCreating(
            name=cluster_name,
            org_id=org_id,
            address=pydantic.TypeAdapter(AnyUrl).validate_python(URL.parse(domain).url),
            cloud_provider=cloud_provider,
            public_key=public_key,
            state="creating",
            cloud={},
        )
        _ = client.broker.create(body)
    else:
        body = BrokerUpdating(
            id=broker.id,
            name=cluster_name,
            address=domain,
            cloud_provider=cloud_provider,
            public_key=public_key,
            state="creating",
            cloud={},
        )
        _ = client.broker.update(body)


# ---------------------------------------------------------------------------
# NEW COMMANDS — cloud-agnostic cluster lifecycle
# ---------------------------------------------------------------------------


@cli.subcommand(
    "clusters",
    "register",
    name=Arg("--name", help="Human-readable name for the cluster"),
    domain=Arg(
        "--domain", help="Public FQDN of the cluster (e.g. cluster.example.com)"
    ),
    cloud_provider=Arg(
        "--cloud-provider",
        help=f"Cloud provider: {', '.join(CLOUD_PROVIDERS)}",
    ),
    as_json=Arg("--json", help=Messages.as_json),
)
def register(
    name: str,
    domain: str,
    cloud_provider: str,
    as_json: bool = False,
) -> None:
    """Register a new cluster with PAM and write credentials to cluster-creds.json."""
    auth = get_auth_state()
    client = auth.client
    org_id = auth.org_id
    # Create/update the cluster record in PAM
    try:
        broker = client.broker.read(address=domain)
    except BrokerNotFound:
        body = BrokerCreating(
            name=name,
            org_id=org_id,
            address=pydantic.TypeAdapter(AnyUrl).validate_python(URL.parse(domain).url),
            cloud_provider=cloud_provider,
            public_key="",
            state="creating",
            cloud={},
        )
        broker = client.broker.create(body)
    else:
        update_body = BrokerUpdating(
            id=broker.id,
            name=name,
            address=domain,
            cloud_provider=cloud_provider,
            state="creating",
            cloud={},
        )
        _ = client.broker.update(update_body)
    # Fetch credentials
    creds = client.broker.credentials(broker_id=broker.id)
    creds_data = {
        "broker_id": str(creds.broker_id),
        "container_sa_key": creds.container_sa_key,
        "prodigy_license_key": creds.prodigy_license_key,
    }
    creds_path = Path("cluster-creds.json")
    srsly.write_json(creds_path, creds_data)
    if as_json:
        print(json.dumps(creds_data, indent=2))
    else:
        msg.good("Cluster registered with PAM")
        msg.info(f"  Broker ID:    {creds.broker_id}")
        msg.info(f"  Wrote credentials to {creds_path}")


@cli.subcommand(
    "clusters",
    "get-chart",
    dest=Arg("--dest", help="Directory to save chart and terraform files to"),
    cloud_provider=Arg(
        "--cloud-provider",
        help=f"Include terraform for this cloud provider: {', '.join(CLOUD_PROVIDERS)}",
    ),
)
def get_chart(
    dest: Union[Path, str] = ".",
    cloud_provider: Optional[str] = None,
) -> None:
    """Download the Helm chart (and optionally terraform files) for cluster deployment."""
    if isinstance(dest, str):
        dest = Path(dest)
    auth = get_auth_state()
    dest.mkdir(exist_ok=True, parents=True)
    # Download the chart + terraform bundle from PAM
    tar_stream = auth.client.broker.download_gcloud_files()
    tar_bytes = tar_stream.read()
    with tarfile.open(fileobj=io.BytesIO(tar_bytes)) as tf:
        tf.extractall(dest)
    # Restructure: move chart to top-level
    cluster_dir = dest / "cluster"
    if cluster_dir.exists():
        chart_src = cluster_dir / "k8s" / "chart"
        chart_dest = dest / "chart"
        if chart_src.exists() and not chart_dest.exists():
            shutil.copytree(chart_src, chart_dest)
        # Copy provider-specific terraform if requested
        if cloud_provider:
            provider_src = cluster_dir / cloud_provider
            provider_dest = dest / cloud_provider
            if provider_src.exists() and not provider_dest.exists():
                shutil.copytree(provider_src, provider_dest)
        # Copy tasks if present
        tasks_src = cluster_dir / "tasks"
        tasks_dest = dest / "tasks"
        if tasks_src.exists() and not tasks_dest.exists():
            shutil.copytree(tasks_src, tasks_dest)
        shutil.rmtree(cluster_dir)
    msg.good(f"Chart files saved to {dest}")
    if cloud_provider:
        msg.info(
            f"  Terraform for {cloud_provider} included in {dest / cloud_provider}"
        )


@cli.subcommand(
    "clusters",
    "setup-infra",
    gateway=Arg("--gateway", help="Install Gateway API CRDs and Traefik"),
    cert_manager=Arg("--cert-manager", help="Install cert-manager with Let's Encrypt"),
    acme_email=Arg(
        "--acme-email", help="Email for ACME/Let's Encrypt certificate registration"
    ),
    public_ip=Arg("--public-ip", help="Static IP for the Traefik load balancer"),
)
def setup_infra(
    gateway: bool = False,
    cert_manager: bool = False,
    acme_email: Optional[str] = None,
    public_ip: Optional[str] = None,
) -> None:
    """Install cluster infrastructure prerequisites (Gateway API, Traefik, cert-manager).

    These operations are idempotent and safe to re-run.
    """
    if gateway:
        _ensure_gateway_crds()
        _ensure_traefik(public_ip=public_ip or "")
        msg.good("Gateway API and Traefik installed")
    if cert_manager or acme_email:
        _ensure_cert_manager()
        msg.good("cert-manager installed")


@cli.subcommand(
    "clusters",
    "init-values",
    creds=Arg("--creds", help="Path to cluster-creds.json from 'clusters register'"),
    cloud_provider=Arg(
        "--cloud-provider",
        help=f"Cloud provider: {', '.join(CLOUD_PROVIDERS)}",
    ),
    domain=Arg("--domain", help="Public FQDN of the cluster"),
    out=Arg("--out", help="Output path for generated values.yaml"),
    version=Arg("--version", help="Image version tag to use"),
    acme_email=Arg("--acme-email", help="Email for Let's Encrypt certificates"),
)
def init_values(
    creds: Path,
    cloud_provider: str,
    domain: str,
    out: Union[Path, str] = "values.yaml",
    version: Optional[str] = None,
    acme_email: Optional[str] = None,
) -> None:
    """Generate a values.yaml for Helm deployment from credentials and provider defaults."""
    if isinstance(out, str):
        out = Path(out)
    creds_data = srsly.read_json(creds)
    assert isinstance(creds_data, dict)
    broker_id = creds_data["broker_id"]
    container_sa_key = creds_data.get("container_sa_key", "")
    prodigy_license_key = creds_data.get("prodigy_license_key", "")
    if version is None:
        version = _get_latest_version(
            "https://pypi.org/pypi/prodigy-teams-recipes/json"
        )
    auth = get_auth_state()
    pam_fqdn = auth.pam_url.netloc
    values: dict[str, Any] = {
        "image": {
            "registry": CONTAINER_REGISTRY,
            "broker": {"tag": version},
            "recipes": {"tag": version},
            "cpl": {"tag": version},
        },
        "cluster": {
            "brokerId": broker_id,
            "fqdn": domain,
            "pamFqdn": pam_fqdn,
        },
        "gateway": {
            "name": "prodigy-teams-gateway",
            "namespace": DEFAULT_NAMESPACE,
            "tls": [{"secretName": "prodigy-teams-tls"}],
        },
        "builtinRecipesVersion": version,
    }
    # cert-manager
    if acme_email:
        values["certManager"] = {
            "enabled": True,
            "secretName": "prodigy-teams-tls",
            "acmeEmail": acme_email,
        }
    # Cloud-specific defaults
    if cloud_provider == "k3s":
        values["postgresql"] = {
            "enabled": True,
            "auth": {
                "database": "brokerdb",
                "username": "broker_user",
                "password": "CHANGE-ME",
            },
        }
        values["storage"] = {
            "type": "hostPath",
            "hostPath": "/srv/prodigy-data",
        }
        values["imagePullSecrets"] = [{"name": "explosion-registry"}]
        values["secrets"] = {
            "infraSecretName": "prodigy-teams-infra",
            "containerSaKey": container_sa_key,
            "prodigyLicenseKey": prodigy_license_key,
        }
    elif cloud_provider == "gcp":
        values["gcp"] = {
            "project": "your-gcp-project",
            "zone": "europe-west1-d",
        }
        values["serviceAccount"] = {
            "annotations": {
                "iam.gke.io/gcp-service-account": "cluster-gke-sa@your-project.iam.gserviceaccount.com",
            },
        }
        values["database"] = {
            "ip": "FILL-IN",
            "user": "broker_user",
            "name": "brokerdb",
        }
        values["secrets"] = {
            "infraSecretName": "prodigy-teams-infra",
            "containerSaKey": container_sa_key,
            "prodigyLicenseKey": prodigy_license_key,
        }
        values["nfs"] = {"claimName": "prodigy-nfs"}
    elif cloud_provider == "aws":
        values["database"] = {
            "ip": "FILL-IN",
            "user": "broker_user",
            "name": "brokerdb",
        }
        values["imagePullSecrets"] = [{"name": "explosion-registry"}]
        values["secrets"] = {
            "infraSecretName": "prodigy-teams-infra",
            "containerSaKey": container_sa_key,
            "prodigyLicenseKey": prodigy_license_key,
        }
        values["nfs"] = {"claimName": "prodigy-nfs"}
    srsly.write_yaml(out, values)
    msg.good(f"Generated {out}")
    msg.info(
        f"  Review and edit before deploying with 'ptc clusters deploy --values {out}'"
    )


@cli.subcommand(
    "clusters",
    "deploy",
    values_file=Arg("--values", help="Path to values.yaml for Helm deployment"),
    chart=Arg(
        "--chart",
        help="Path to local chart directory (default: OCI chart from registry)",
    ),
    latest=Arg(
        "--latest",
        help="Fetch and use the latest image tags before deploying",
    ),
    namespace=Arg("--namespace", help="Kubernetes namespace"),
    wait=Arg("--wait", help="Wait for pods to be ready after deploying"),
)
def deploy(
    values_file: Path,
    chart: Optional[Path] = None,
    latest: bool = False,
    namespace: str = DEFAULT_NAMESPACE,
    wait: bool = False,
) -> None:
    """Deploy or upgrade the Prodigy Teams Helm chart.

    Creates the namespace, image pull secret, broker keypair secret,
    then runs helm upgrade --install with the provided values.
    """
    values = srsly.read_yaml(values_file)
    assert isinstance(values, dict)
    # Optionally update to latest version
    if latest:
        latest_version = _get_latest_version(
            "https://pypi.org/pypi/prodigy-teams-recipes/json"
        )
        msg.info(f"Using latest version: {latest_version}")
        values.setdefault("image", {})
        for service in ("broker", "recipes", "cpl"):
            values["image"].setdefault(service, {})["tag"] = latest_version
        values["builtinRecipesVersion"] = latest_version
        # Write back updated values
        srsly.write_yaml(values_file, values)
    # 1. Create namespace
    msg.info(f"Ensuring namespace {namespace}...")
    subprocess.run(
        ["kubectl", "create", "namespace", namespace, "--dry-run=client", "-o", "yaml"],
        capture_output=True,
        text=True,
        check=True,
    )
    subprocess.run(
        ["kubectl", "create", "namespace", namespace],
        capture_output=True,
        text=True,
    )
    # 2. Create image pull secret (for non-GKE clusters using SA key)
    image_pull_secrets = values.get("imagePullSecrets", [])
    container_sa_key = values.get("secrets", {}).get("containerSaKey", "")
    if image_pull_secrets and container_sa_key:
        secret_name = image_pull_secrets[0]["name"]
        _create_image_pull_secret(
            secret_name=secret_name,
            sa_key_b64=container_sa_key,
            namespace=namespace,
        )
    # 3. Create broker RSA keypair secret (if infra secret configured)
    infra_secret_name = values.get("secrets", {}).get("infraSecretName", "")
    if infra_secret_name:
        _ensure_broker_keypair_standalone(infra_secret_name, namespace)
    # 4. Helm upgrade --install
    if chart is None:
        # Use OCI chart from the registry
        registry = values.get("image", {}).get("registry", CONTAINER_REGISTRY)
        chart_ref = f"oci://{registry}/prodigy-teams-chart"
        msg.info(f"Installing from {chart_ref}...")
        _helm_upgrade_from_ref(chart_ref, values, namespace=namespace, wait=wait)
    else:
        msg.info(f"Installing from local chart {chart}...")
        _helm_upgrade(chart, values, namespace=namespace, wait=wait)
    msg.good("Deployment complete")
    if wait:
        msg.info("Running verification checks...")
        _helm_test(release=DEFAULT_RELEASE, namespace=namespace)


def _create_image_pull_secret(
    secret_name: str, sa_key_b64: str, namespace: str
) -> None:
    """Create a docker-registry pull secret from a base64-encoded SA key."""
    try:
        sa_key_json = base64.b64decode(sa_key_b64).decode("utf-8")
    except Exception:
        sa_key_json = sa_key_b64
    cmd = [
        "kubectl",
        "create",
        "secret",
        "docker-registry",
        secret_name,
        f"--namespace={namespace}",
        "--docker-server=europe-west1-docker.pkg.dev",
        "--docker-username=_json_key",
        f"--docker-password={sa_key_json}",
        "--dry-run=client",
        "-o",
        "yaml",
    ]
    create = subprocess.run(cmd, capture_output=True, text=True, check=True)
    subprocess.run(
        ["kubectl", "apply", "-f", "-"],
        input=create.stdout,
        check=True,
        text=True,
    )
    msg.info(f"Image pull secret '{secret_name}' created/updated")


def _ensure_broker_keypair_standalone(secret_name: str, namespace: str) -> str:
    """Ensure the infra K8s secret has a broker RSA keypair, return the public PEM."""
    from ..key_pair import generate_key_pair

    check = subprocess.run(
        [
            "kubectl",
            "get",
            "secret",
            secret_name,
            "-n",
            namespace,
            "-o",
            "jsonpath={.data.broker_public_key}",
        ],
        capture_output=True,
        text=True,
    )
    if check.returncode == 0 and check.stdout.strip():
        msg.info(f"Using existing broker keypair from secret {secret_name}")
        return base64.b64decode(check.stdout.strip()).decode("utf-8")
    msg.info("Generating new broker RSA keypair...")
    kp = generate_key_pair()
    cmd = [
        "kubectl",
        "create",
        "secret",
        "generic",
        secret_name,
        f"--namespace={namespace}",
        f"--from-literal=broker_private_key={kp.private_pem}",
        f"--from-literal=broker_public_key={kp.public_pem}",
        "--dry-run=client",
        "-o",
        "yaml",
    ]
    create = subprocess.run(cmd, capture_output=True, text=True, check=True)
    subprocess.run(
        ["kubectl", "apply", "-f", "-"],
        input=create.stdout,
        check=True,
        text=True,
    )
    msg.good(f"Created broker keypair in secret {secret_name}")
    return kp.public_pem


def _helm_upgrade_from_ref(
    chart_ref: str,
    values: dict,
    release: str = DEFAULT_RELEASE,
    namespace: str = DEFAULT_NAMESPACE,
    wait: bool = False,
    timeout: str = "10m",
) -> None:
    """Run helm upgrade --install using an OCI chart reference."""
    cmd = [
        "helm",
        "upgrade",
        "--install",
        release,
        chart_ref,
        "--namespace",
        namespace,
        "--create-namespace",
        "--timeout",
        timeout,
    ]
    if wait:
        cmd += ["--wait"]
    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
        tmp.write(json.dumps(values, indent=2))
        tmp_path = tmp.name
    try:
        cmd += ["-f", tmp_path]
        subprocess.run(cmd, check=True, text=True)
    finally:
        Path(tmp_path).unlink(missing_ok=True)
