#!/usr/bin/env python3
"""
fix_cde_ssl.py — Fix SSL PKIX error in the CDE dbus-wxm-client Kubernetes deployment.

In CDE the deployment already has a truststore secret (dbus-ozone-gateway-truststore)
and JVM_OPTS configured.  This script downloads the server certificate from the
telemetry.altus.url configured in the dbus-client.xml configmap, appends it to the
existing truststore inside the secret, and restarts the deployment so the pod picks
up the updated truststore.

Usage:
    export KUBECONFIG=/path/to/kubeconfig.yaml
    python3 fix_cde_ssl.py --namespace <namespace>

Requirements (runtime): Python 3.8+, kubectl on PATH, keytool (JDK) on PATH
"""

import argparse
import base64
import hashlib
import json
import os
import shutil
import socket
import ssl
import subprocess
import sys
import tempfile
import xml.etree.ElementTree as ET
from datetime import datetime
from typing import Callable, Dict, List, Optional
from urllib.parse import urlparse

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DEPLOYMENT_NAME = "dex-base-dbus-wxm-client"
CONTAINER_NAME = "dbus-wxm-client"
CONFIGMAP_NAME = "dex-base-dbus-wxm-client"
CONFIGMAP_XML_KEY = "dbus-client.xml"
SECRET_NAME = "dbus-ozone-gateway-truststore"
SECRET_JKS_KEY = "truststore.jks"
TRUSTSTORE_MOUNT_PATH = "/etc/dbus-wxm-client/trust/truststore.jks"
TRUSTSTORE_PASSWORD = "changeit"
ROLLOUT_TIMEOUT = "300s"

# ---------------------------------------------------------------------------
# Exception hierarchy
# ---------------------------------------------------------------------------


class CDESSLError(Exception):
    """Base error for all script failures."""


class PrerequisiteError(CDESSLError):
    """Missing tool on PATH, KUBECONFIG not set or file missing."""


class DeploymentError(CDESSLError):
    """Deployment/pod not found, JVM_OPTS misconfigured, or rollout failed."""


class ConfigMapError(CDESSLError):
    """ConfigMap missing, XML key missing, or telemetry.altus.url missing/empty."""


class CertificateError(CDESSLError):
    """Certificate download failure or cert already present in truststore."""


class TruststoreError(CDESSLError):
    """keytool failure while reading or modifying the truststore."""


class BackupError(CDESSLError):
    """Backup export failed — no K8s resources were mutated."""


class SecretError(CDESSLError):
    """Secret not found, missing expected key, or patch failed."""


# ---------------------------------------------------------------------------
# Pure functions — no I/O, fully unit-testable without mocks
# ---------------------------------------------------------------------------


def parse_telemetry_url(xml_content: str) -> str:
    """Parse telemetry.altus.url from dbus-client.xml content.

    Args:
        xml_content: Raw XML string from the configmap.

    Returns:
        The URL string (e.g. 'https://host:12022').

    Raises:
        ConfigMapError: If the property is missing or its value is empty.
    """
    root = ET.fromstring(xml_content)
    for prop in root.findall("property"):
        name_el = prop.find("name")
        if name_el is not None and name_el.text == "telemetry.altus.url":
            value_el = prop.find("value")
            if value_el is None or not (value_el.text or "").strip():
                raise ConfigMapError(
                    "telemetry.altus.url property found but value is empty"
                )
            return value_el.text.strip()
    raise ConfigMapError(
        "telemetry.altus.url property not found in dbus-client.xml"
    )


def parse_pem_certificates(pem_bundle: str) -> List[str]:
    """Split a PEM bundle string into individual certificate strings.

    Args:
        pem_bundle: String containing one or more PEM certificate blocks.

    Returns:
        List of individual PEM certificate strings (each includes header/footer).
        Returns empty list if pem_bundle contains no certificate blocks.
    """
    certs = []
    current: List[str] = []
    in_cert = False
    for line in pem_bundle.splitlines():
        if line.strip() == "-----BEGIN CERTIFICATE-----":
            in_cert = True
            current = [line]
        elif line.strip() == "-----END CERTIFICATE-----":
            current.append(line)
            certs.append("\n".join(current) + "\n")
            current = []
            in_cert = False
        elif in_cert:
            current.append(line)
    return certs


def get_cert_fingerprint_sha256(pem: str) -> str:
    """Compute SHA-256 fingerprint of a PEM certificate.

    Args:
        pem: PEM-encoded certificate string.

    Returns:
        Colon-separated uppercase hex string, e.g. "AA:BB:CC:...".

    Raises:
        CertificateError: If the PEM is malformed and cannot be decoded.
    """
    try:
        der = ssl.PEM_cert_to_DER_cert(pem)
    except Exception as exc:
        raise CertificateError(
            f"Cannot decode PEM certificate: {exc}"
        ) from exc
    digest = hashlib.sha256(der).digest()
    return ":".join(f"{b:02X}" for b in digest)


def get_existing_aliases(keytool_list_output: str) -> List[str]:
    """Parse alias names from the output of `keytool -list`.

    Each entry line produced by keytool -list (non-verbose) looks like:
        alias, <date>, trustedCertEntry,

    Args:
        keytool_list_output: stdout text from `keytool -list -keystore ...`.

    Returns:
        List of alias name strings in the order they appear.
        Returns empty list if the keystore is empty.
    """
    aliases = []
    for line in keytool_list_output.splitlines():
        # keytool entry lines: "<alias>, <date>, <type>,"
        parts = line.split(",")
        if len(parts) >= 3 and parts[-1].strip() == "":
            candidate = parts[0].strip()
            # Filter out lines that are clearly not alias entries
            # (e.g. "Keystore type: JKS", "Your keystore contains N entries")
            if candidate and " " not in candidate and ":" not in candidate:
                aliases.append(candidate)
    return aliases


def find_next_alias(existing_aliases: List[str]) -> str:
    """Compute the next certN alias given a list of existing alias strings.

    Uses max(N) + 1 strategy to avoid overwriting any existing alias even if
    there are gaps in the sequence (e.g. ["cert0","cert5"] → "cert6").

    Args:
        existing_aliases: List of alias strings, possibly including non-certN names.

    Returns:
        "certN" where N = max existing certN index + 1, or "cert0" if none exist.
    """
    max_n = -1
    for alias in existing_aliases:
        if alias.startswith("cert"):
            suffix = alias[4:]
            if suffix.isdigit():
                max_n = max(max_n, int(suffix))
    return f"cert{max_n + 1}"


def check_cert_not_in_truststore(keytool_output: str, cert_pem: str) -> None:
    """Verify the certificate is not already present in the truststore.

    Computes the SHA-256 fingerprint of cert_pem dynamically (never hardcoded)
    and searches keytool_output for that fingerprint.

    Args:
        keytool_output: stdout text from `keytool -list -keystore ...`.
        cert_pem: PEM string of the certificate to check.

    Raises:
        CertificateError: If the fingerprint is already found in keytool_output.
    """
    fingerprint = get_cert_fingerprint_sha256(cert_pem)
    if fingerprint in keytool_output:
        raise CertificateError(
            "Certificate is already present in the truststore "
            f"(fingerprint: {fingerprint}). "
            "The script has likely been run before — no action needed."
        )


# ---------------------------------------------------------------------------
# I/O boundary functions — all external calls injectable via `run` parameter
# ---------------------------------------------------------------------------


def validate_prerequisites(env: Dict[str, str], run: Callable) -> None:
    """Validate all prerequisites before any K8s or network calls.

    Checks:
    - KUBECONFIG env var is set
    - KUBECONFIG file exists on disk
    - kubectl is on PATH
    - keytool is on PATH

    Args:
        env: Environment variable dict (pass os.environ in production).
        run: Callable matching subprocess.run signature (injectable for tests).

    Raises:
        PrerequisiteError: If any prerequisite is not met.
    """
    kubeconfig = env.get("KUBECONFIG")
    if not kubeconfig:
        raise PrerequisiteError(
            "KUBECONFIG environment variable is not set. "
            "Run: export KUBECONFIG=/path/to/kubeconfig.yaml"
        )
    if not os.path.isfile(kubeconfig):
        raise PrerequisiteError(
            f"KUBECONFIG file not found: {kubeconfig}"
        )
    if not shutil.which("kubectl"):
        raise PrerequisiteError(
            "kubectl not found on PATH. Install kubectl and ensure it is on your PATH."
        )
    if not shutil.which("keytool"):
        raise PrerequisiteError(
            "keytool not found on PATH. Install a JDK and ensure it is on your PATH."
        )


def find_deployment(namespace: str, run: Callable) -> dict:
    """Retrieve the dex-base-dbus-wxm-client Deployment as a parsed JSON dict.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Returns:
        Parsed deployment JSON dict.

    Raises:
        DeploymentError: If the deployment is not found.
    """
    result = run(
        ["kubectl", "get", "deployment", DEPLOYMENT_NAME,
         "-n", namespace, "-o", "json"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DeploymentError(
            f"{DEPLOYMENT_NAME} deployment not found in namespace '{namespace}'. "
            f"Error: {result.stderr.strip()}"
        )
    return json.loads(result.stdout)


def check_jvm_opts_is_set(deployment: dict) -> None:
    """Verify JVM_OPTS is correctly set in the dbus-wxm-client container.

    In CDE the JVM_OPTS env var and truststore mount are already in place —
    this is a prerequisite, not something this script configures.

    Raises DeploymentError if:
    (a) JVM_OPTS is absent from the dbus-wxm-client container env, OR
    (b) JVM_OPTS does not reference TRUSTSTORE_MOUNT_PATH.

    Args:
        deployment: Parsed deployment JSON dict (from find_deployment).

    Raises:
        DeploymentError: If JVM_OPTS is missing or points to wrong path.
    """
    containers = (
        deployment.get("spec", {})
        .get("template", {})
        .get("spec", {})
        .get("containers", [])
    )
    for container in containers:
        if container.get("name") == CONTAINER_NAME:
            for env_var in container.get("env", []):
                if env_var.get("name") == "JVM_OPTS":
                    value = env_var.get("value", "")
                    if TRUSTSTORE_MOUNT_PATH not in value:
                        raise DeploymentError(
                            f"JVM_OPTS is set but does not reference the expected "
                            f"truststore path '{TRUSTSTORE_MOUNT_PATH}'. "
                            f"Current value: {value!r}"
                        )
                    return
            raise DeploymentError(
                f"JVM_OPTS env var is not set on the '{CONTAINER_NAME}' container. "
                "CDE is expected to have JVM_OPTS pre-configured. "
                "Check that the deployment is the correct CDE version."
            )
    raise DeploymentError(
        f"Container '{CONTAINER_NAME}' not found in deployment '{DEPLOYMENT_NAME}'."
    )


def get_existing_secret_jks(namespace: str, temp_dir: str, run: Callable) -> str:
    """Download the existing truststore.jks from the K8s secret.

    Downloads the dbus-ozone-gateway-truststore secret, base64-decodes the
    truststore.jks key, and writes it to temp_dir/existing-truststore.jks.

    Args:
        namespace: Kubernetes namespace.
        temp_dir: Directory where the JKS file will be written.
        run: subprocess.run-compatible callable.

    Returns:
        Absolute path to the written JKS file.

    Raises:
        SecretError: If the secret is not found or has no truststore.jks key.
    """
    result = run(
        ["kubectl", "get", "secret", SECRET_NAME, "-n", namespace, "-o", "json"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise SecretError(
            f"Secret '{SECRET_NAME}' not found in namespace '{namespace}'. "
            "This script expects the secret to already exist. "
            f"Error: {result.stderr.strip()}"
        )
    secret = json.loads(result.stdout)
    jks_b64 = secret.get("data", {}).get(SECRET_JKS_KEY)
    if jks_b64 is None:
        raise SecretError(
            f"Key '{SECRET_JKS_KEY}' not found in secret '{SECRET_NAME}'. "
            f"Available keys: {list(secret.get('data', {}).keys())}"
        )
    jks_bytes = base64.b64decode(jks_b64)
    out_path = os.path.join(temp_dir, "existing-truststore.jks")
    with open(out_path, "wb") as f:
        f.write(jks_bytes)
    return out_path


def get_keytool_list_output(truststore_path: str, run: Callable) -> str:
    """Run `keytool -list` and return its stdout.

    Args:
        truststore_path: Path to the JKS truststore file.
        run: subprocess.run-compatible callable.

    Returns:
        stdout text from keytool.

    Raises:
        TruststoreError: If keytool returns non-zero exit code.
    """
    result = run(
        [
            "keytool", "-list",
            "-keystore", truststore_path,
            "-storepass", TRUSTSTORE_PASSWORD,
        ],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise TruststoreError(
            f"keytool -list failed on '{truststore_path}': {result.stderr.strip()}"
        )
    return result.stdout


def backup_resources(
    namespace: str, backup_dir: str, original_jks_path: str, run: Callable
) -> None:
    """Export deployment YAML, secret YAML, and original JKS binary to backup_dir.

    Writes:
    - <backup_dir>/deployment-dex-base-dbus-wxm-client.yaml
    - <backup_dir>/secret-dbus-ozone-gateway-truststore.yaml
    - <backup_dir>/truststore.jks  (binary copy of original)

    Args:
        namespace: Kubernetes namespace.
        backup_dir: Path to (already created) backup directory.
        original_jks_path: Path to the existing truststore.jks before modification.
        run: subprocess.run-compatible callable.

    Raises:
        BackupError: If any export fails.
    """
    resources = [
        (
            ["kubectl", "get", "deployment", DEPLOYMENT_NAME,
             "-n", namespace, "-o", "yaml"],
            f"deployment-{DEPLOYMENT_NAME}.yaml",
        ),
        (
            ["kubectl", "get", "secret", SECRET_NAME,
             "-n", namespace, "-o", "yaml"],
            f"secret-{SECRET_NAME}.yaml",
        ),
    ]
    for cmd, filename in resources:
        result = run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            raise BackupError(
                f"Failed to back up {filename}: {result.stderr.strip()}. "
                "No Kubernetes resources have been modified."
            )
        out_path = os.path.join(backup_dir, filename)
        with open(out_path, "w") as f:
            f.write(result.stdout)

    # Copy original JKS binary to backup dir for manual rollback
    jks_backup_path = os.path.join(backup_dir, SECRET_JKS_KEY)
    shutil.copy2(original_jks_path, jks_backup_path)


def get_configmap(namespace: str, run: Callable) -> str:
    """Retrieve the dbus-client.xml content from the configmap.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Returns:
        Raw XML string from the dbus-client.xml key.

    Raises:
        ConfigMapError: If the configmap or the key is not found.
    """
    result = run(
        ["kubectl", "get", "configmap", CONFIGMAP_NAME,
         "-n", namespace, "-o", "json"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise ConfigMapError(
            f"ConfigMap '{CONFIGMAP_NAME}' not found in namespace '{namespace}'. "
            f"Error: {result.stderr.strip()}"
        )
    cm = json.loads(result.stdout)
    xml_content = cm.get("data", {}).get(CONFIGMAP_XML_KEY)
    if xml_content is None:
        raise ConfigMapError(
            f"Key '{CONFIGMAP_XML_KEY}' not found in ConfigMap '{CONFIGMAP_NAME}'."
        )
    return xml_content


def check_pod_running(namespace: str, run: Callable) -> str:
    """Find a Running dbus-wxm-client pod and return its name.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Returns:
        Pod name string.

    Raises:
        DeploymentError: If no dbus-wxm-client pod is in Running state.
    """
    result = run(
        ["kubectl", "get", "pods", "-n", namespace,
         "-l", f"app.kubernetes.io/name={CONTAINER_NAME}", "-o", "json"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DeploymentError(
            f"Failed to list pods in namespace '{namespace}': {result.stderr.strip()}"
        )
    pods_data = json.loads(result.stdout)
    for pod in pods_data.get("items", []):
        phase = pod.get("status", {}).get("phase", "")
        if phase == "Running":
            return pod["metadata"]["name"]
    raise DeploymentError(
        f"No {CONTAINER_NAME} pod found in Running state in namespace '{namespace}'. "
        "Ensure the pod is running before executing this script."
    )


def _fetch_pem_certificate(host: str, port: int) -> str:
    """Open an SSL connection and return the server certificate as PEM string.

    Separated from download_server_certificate so tests can patch it directly.
    """
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    with ctx.wrap_socket(
        socket.create_connection((host, port), timeout=10), server_hostname=host
    ) as ssock:
        der_cert = ssock.getpeercert(binary_form=True)
    return ssl.DER_cert_to_PEM_cert(der_cert)


def download_server_certificate(host: str, port: int, temp_dir: str) -> str:
    """Download the server's TLS certificate via an SSL socket connection.

    Args:
        host: Hostname to connect to.
        port: TCP port.
        temp_dir: Directory where the PEM file will be written.

    Returns:
        Absolute path to the written PEM file.

    Raises:
        CertificateError: If the connection fails or no certificate is returned.
    """
    try:
        pem = _fetch_pem_certificate(host, port)
    except (OSError, socket.error) as exc:
        raise CertificateError(
            f"Could not connect to {host}:{port} to download the server certificate. "
            f"Ensure you have network/VPN access to this host. Error: {exc}"
        ) from exc

    if not pem or "BEGIN CERTIFICATE" not in pem:
        raise CertificateError(
            f"No certificate returned from {host}:{port}."
        )

    cert_path = os.path.join(temp_dir, "server.pem")
    with open(cert_path, "w") as f:
        f.write(pem)
    return cert_path


def append_cert_to_truststore(
    truststore_path: str, cert_pem_path: str, alias: str, run: Callable
) -> None:
    """Import a certificate into an existing JKS truststore.

    Args:
        truststore_path: Path to the existing truststore.jks.
        cert_pem_path: Path to the PEM file of the cert to import.
        alias: Alias to use for this cert entry (e.g. "cert159").
        run: subprocess.run-compatible callable.

    Raises:
        TruststoreError: If keytool returns non-zero exit code.
    """
    result = run(
        [
            "keytool", "-importcert",
            "-noprompt",
            "-keystore", truststore_path,
            "-storepass", TRUSTSTORE_PASSWORD,
            "-alias", alias,
            "-file", cert_pem_path,
        ],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise TruststoreError(
            f"keytool failed to import certificate with alias '{alias}': "
            f"{result.stderr.strip()}"
        )


def update_k8s_secret(namespace: str, truststore_path: str, run: Callable) -> None:
    """Update the existing K8s secret with the modified truststore.jks.

    Uses kubectl create secret --dry-run=client -o yaml | kubectl apply -f -
    to perform an idempotent in-place update of the secret data.

    Args:
        namespace: Kubernetes namespace.
        truststore_path: Path to the updated truststore.jks file.
        run: subprocess.run-compatible callable.

    Raises:
        SecretError: If kubectl fails.
    """
    # Generate the secret YAML with the new JKS
    dry_run = run(
        [
            "kubectl", "create", "secret", "generic", SECRET_NAME,
            f"--from-file={SECRET_JKS_KEY}={truststore_path}",
            "-n", namespace,
            "--dry-run=client", "-o", "yaml",
        ],
        capture_output=True, text=True,
    )
    if dry_run.returncode != 0:
        raise SecretError(
            f"Failed to generate updated secret YAML: {dry_run.stderr.strip()}"
        )

    # Apply the generated YAML to update the secret in-place
    apply = run(
        ["kubectl", "apply", "-f", "-", "-n", namespace],
        input=dry_run.stdout,
        capture_output=True, text=True,
    )
    if apply.returncode != 0:
        raise SecretError(
            f"Failed to apply updated secret '{SECRET_NAME}': {apply.stderr.strip()}"
        )


def restart_deployment(namespace: str, run: Callable) -> None:
    """Trigger a rolling restart of the dex-base-dbus-wxm-client deployment.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Raises:
        DeploymentError: If kubectl rollout restart fails.
    """
    result = run(
        ["kubectl", "rollout", "restart",
         f"deployment/{DEPLOYMENT_NAME}", "-n", namespace],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DeploymentError(
            f"Failed to restart deployment '{DEPLOYMENT_NAME}': {result.stderr.strip()}"
        )


def wait_for_rollout(namespace: str, run: Callable) -> None:
    """Wait for the deployment rollout to complete.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Raises:
        DeploymentError: If the rollout does not complete within ROLLOUT_TIMEOUT.
    """
    result = run(
        [
            "kubectl", "rollout", "status",
            f"deployment/{DEPLOYMENT_NAME}",
            "-n", namespace,
            f"--timeout={ROLLOUT_TIMEOUT}",
        ],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DeploymentError(
            f"Rollout of '{DEPLOYMENT_NAME}' did not complete within {ROLLOUT_TIMEOUT}. "
            f"Output: {result.stderr.strip() or result.stdout.strip()}"
        )


def rollback_secret(
    namespace: str, original_jks_path: str, run: Callable
) -> None:
    """Restore the original truststore.jks to the K8s secret.

    Best-effort: logs a warning if restoration fails but does not re-raise,
    so the caller can still print full rollback instructions.

    Args:
        namespace: Kubernetes namespace.
        original_jks_path: Path to the saved original JKS file.
        run: subprocess.run-compatible callable.
    """
    dry_run = run(
        [
            "kubectl", "create", "secret", "generic", SECRET_NAME,
            f"--from-file={SECRET_JKS_KEY}={original_jks_path}",
            "-n", namespace,
            "--dry-run=client", "-o", "yaml",
        ],
        capture_output=True, text=True,
    )
    if dry_run.returncode != 0:
        print(
            f"WARNING: Failed to generate rollback secret YAML for '{SECRET_NAME}': "
            f"{dry_run.stderr.strip()}",
            file=sys.stderr,
        )
        return

    apply = run(
        ["kubectl", "apply", "-f", "-", "-n", namespace],
        input=dry_run.stdout,
        capture_output=True, text=True,
    )
    if apply.returncode != 0:
        print(
            f"WARNING: Rollback failed — could not restore '{SECRET_NAME}': "
            f"{apply.stderr.strip()}",
            file=sys.stderr,
        )


def print_success_summary(namespace: str, backup_dir: str) -> None:
    """Print the backup path and rollback instructions to stdout.

    Args:
        namespace: Kubernetes namespace (used in rollback commands).
        backup_dir: Path to the timestamped backup directory.
    """
    print("\n" + "=" * 60)
    print("SUCCESS: CDE dbus-wxm-client SSL truststore updated.")
    print("=" * 60)
    print(f"\nBackup saved to: {backup_dir}")
    print(
        "\nTo revert all changes (restore original truststore):\n"
        f"  kubectl create secret generic {SECRET_NAME} \\\n"
        f"    --from-file={SECRET_JKS_KEY}={backup_dir}/{SECRET_JKS_KEY} \\\n"
        f"    -n {namespace} --dry-run=client -o yaml | kubectl apply -f -\n"
        f"  kubectl rollout restart deployment/{DEPLOYMENT_NAME} -n {namespace}"
    )
    print()


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
    """Parse command-line arguments.

    Args:
        argv: Argument list (defaults to sys.argv[1:]).

    Returns:
        Parsed namespace with .namespace attribute.
    """
    parser = argparse.ArgumentParser(
        description=(
            "Fix SSL PKIX error in the CDE dex-base-dbus-wxm-client deployment.\n"
            "Appends the telemetry endpoint certificate to the existing truststore."
        )
    )
    parser.add_argument(
        "--namespace",
        "-n",
        required=True,
        help="Kubernetes namespace containing the dex-base-dbus-wxm-client deployment.",
    )
    return parser.parse_args(argv)


def main() -> None:
    """Orchestrate all steps in order with full error handling and rollback."""
    args = parse_args()
    namespace = args.namespace

    validate_prerequisites(env=os.environ, run=subprocess.run)

    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup_dir = os.path.abspath(f"cde-ssl-backup-{timestamp}")
    original_jks_path: Optional[str] = None

    try:
        with tempfile.TemporaryDirectory() as temp_dir:
            # --- Guard checks (read-only) ---
            deployment = find_deployment(namespace, run=subprocess.run)
            check_jvm_opts_is_set(deployment)

            # --- Download existing JKS from secret ---
            print(f"[..] Fetching existing truststore from secret '{SECRET_NAME}' ...")
            original_jks_path = get_existing_secret_jks(
                namespace, temp_dir, run=subprocess.run
            )
            print(f"[OK] Existing truststore written to: {original_jks_path}")

            # --- Read configmap and extract URL ---
            xml_content = get_configmap(namespace, run=subprocess.run)
            telemetry_url = parse_telemetry_url(xml_content)

            parsed = urlparse(telemetry_url)
            host = parsed.hostname
            port = parsed.port or 443

            # --- Check if cert is already in the truststore ---
            keytool_out = get_keytool_list_output(original_jks_path, run=subprocess.run)
            print(f"[..] Downloading server certificate from {host}:{port} to check ...")
            # Download first so we can fingerprint it for the duplicate check
            server_cert_path = download_server_certificate(host, port, temp_dir)
            server_cert_pem = open(server_cert_path).read()
            check_cert_not_in_truststore(keytool_out, server_cert_pem)
            print(f"[OK] Certificate is not yet in the truststore — proceeding.")

            # --- Backup before any mutations ---
            os.makedirs(backup_dir, exist_ok=True)
            backup_resources(
                namespace, backup_dir, original_jks_path, run=subprocess.run
            )
            print(f"[OK] Backup saved to: {backup_dir}")

            # --- Verify pod is running ---
            pod_name = check_pod_running(namespace, run=subprocess.run)
            print(f"[OK] Found running pod: {pod_name}")

            # --- Compute next alias ---
            existing_aliases = get_existing_aliases(keytool_out)
            next_alias = find_next_alias(existing_aliases)
            print(f"[OK] Will import certificate as alias '{next_alias}'.")

            # --- Append cert to truststore ---
            # Work on a copy so we never touch the original in temp_dir
            import shutil as _shutil
            updated_jks = os.path.join(temp_dir, "updated-truststore.jks")
            _shutil.copy2(original_jks_path, updated_jks)
            print(f"[..] Appending certificate to truststore ...")
            append_cert_to_truststore(
                updated_jks, server_cert_path, next_alias, run=subprocess.run
            )
            print(f"[OK] Certificate appended with alias '{next_alias}'.")

            # --- Update K8s secret ---
            print(f"[..] Updating secret '{SECRET_NAME}' ...")
            try:
                update_k8s_secret(namespace, updated_jks, run=subprocess.run)
            except SecretError as exc:
                print(f"\nERROR: Secret update failed: {exc}", file=sys.stderr)
                _print_rollback_hint(namespace, backup_dir)
                raise

            print(f"[OK] Secret updated.")

            # --- Restart deployment ---
            print(f"[..] Restarting deployment '{DEPLOYMENT_NAME}' ...")
            try:
                restart_deployment(namespace, run=subprocess.run)
                print(f"[..] Waiting for rollout to complete ...")
                wait_for_rollout(namespace, run=subprocess.run)
            except DeploymentError as exc:
                print(f"\nERROR: Restart/rollout failed: {exc}", file=sys.stderr)
                print("Rolling back: restoring original truststore ...", file=sys.stderr)
                rollback_secret(namespace, original_jks_path, run=subprocess.run)
                _print_rollback_hint(namespace, backup_dir)
                raise

            print(f"[OK] Deployment restarted successfully.")
            print_success_summary(namespace, backup_dir)

    except CDESSLError as exc:
        print(f"\nERROR: {exc}", file=sys.stderr)
        sys.exit(1)


def _print_rollback_hint(namespace: str, backup_dir: str) -> None:
    print(
        f"\nTo manually restore the original truststore:\n"
        f"  kubectl create secret generic {SECRET_NAME} \\\n"
        f"    --from-file={SECRET_JKS_KEY}={backup_dir}/{SECRET_JKS_KEY} \\\n"
        f"    -n {namespace} --dry-run=client -o yaml | kubectl apply -f -",
        file=sys.stderr,
    )


if __name__ == "__main__":
    main()