import base64
import json
import subprocess
from pathlib import Path
from typing import Optional
from uuid import UUID

import srsly
from radicli import Arg
from wasabi import msg

from prodigy_teams_pam_sdk.models import (
    BrokerUpdating,
)

from ...cli import cli
from ...helm import (
    HelmChartOverrides,
    preflight_test,
    upgrade,
    upgrade_from_ref,
)
from ...key_pair import generate_key_pair
from .._state import get_auth_state

DEFAULT_NAMESPACE = "prodigy-teams"
DEFAULT_RELEASE = "prodigy-teams"


@cli.subcommand(
    "infra",
    "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 Ellf 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 from PAM
    if latest:
        broker_id = values.get("cluster", {}).get("brokerId", "")
        if not broker_id:
            msg.fail("Cannot use --latest: values.yaml missing cluster.brokerId")
            raise SystemExit(1)
        auth = get_auth_state()
        creds = auth.client.broker.credentials(broker_id=UUID(broker_id))
        msg.info(
            f"Using versions from PAM: broker={creds.image.broker_version}"
            f" cpl={creds.image.cpl_version}"
        )
        values.setdefault("image", {})
        values["image"]["brokerVersion"] = creds.image.broker_version
        values["image"]["cplVersion"] = creds.image.cpl_version
        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)
    # and register the public key with PAM for token verification.
    infra_secret_name = values.get("secrets", {}).get("infraSecretName", "")
    if infra_secret_name:
        # Resolve the database password from postgresql subchart or explicit config.
        pg = values.get("postgresql", {})
        database_password = (
            pg.get("auth", {}).get("password", "")
            if pg.get("enabled")
            else values.get("secrets", {}).get("databasePassword", "")
        )
        public_key_pem = _ensure_broker_keypair_standalone(
            infra_secret_name, namespace, database_password=database_password
        )
        broker_id = values.get("cluster", {}).get("brokerId", "")
        if broker_id:
            auth = get_auth_state()
            broker = auth.client.broker.read(id=UUID(broker_id))
            auth.client.broker.update(
                BrokerUpdating(
                    id=broker.id,
                    name=broker.name,
                    address=broker.address,
                    public_key=public_key_pem,
                    cloud_provider=None,
                    cloud=None,
                    state=None,
                )
            )
            msg.good("Registered broker public key with PAM")
    # 4. Helm upgrade --install
    if chart is None:
        # Use OCI chart from the registry
        registry = values.get("image", {}).get("registry", "")
        if not registry:
            msg.fail("image.registry is empty in values.yaml")
            raise SystemExit(1)
        chart_ref = f"oci://{registry}/prodigy-teams"
        msg.info(f"Installing from {chart_ref}...")
        upgrade_from_ref(
            chart_ref, HelmChartOverrides(**values), namespace=namespace, wait=wait
        )
    else:
        msg.info(f"Installing from local chart {chart}...")
        upgrade(chart, HelmChartOverrides(**values), namespace=namespace, wait=wait)
    msg.good("Deployment complete")
    if wait:
        msg.info("Running verification checks...")
        preflight_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, database_password: str = ""
) -> str:
    """Ensure the infra K8s secret has a broker RSA keypair, return the public PEM.

    The broker uses ``PRODIGY_TEAMS_PRIVATE_KEY`` (base64-encoded PEM) for
    signing, so we read/write that key rather than the legacy
    ``broker_private_key`` field.  The public key returned here is registered
    with PAM so it can verify broker-signed tokens.

    When *database_password* is provided it is also stored in the secret as
    ``PRODIGY_TEAMS_DATABASE_PASSWORD`` so the broker and CPL containers can
    authenticate with the database.
    """

    check = subprocess.run(
        [
            "kubectl",
            "get",
            "secret",
            secret_name,
            "-n",
            namespace,
            "-o",
            "jsonpath={.data.PRODIGY_TEAMS_PUBLIC_KEY}",
        ],
        capture_output=True,
        text=True,
    )
    if check.returncode == 0 and check.stdout.strip():
        # PRODIGY_TEAMS_PUBLIC_KEY is base64-encoded PEM (double-encoded in k8s)
        public_pem = base64.b64decode(base64.b64decode(check.stdout.strip())).decode(
            "utf-8"
        )
        msg.info(f"Using existing broker keypair from secret {secret_name}")
        # Ensure the database password is present even if the keypair already
        # exists (it may have been created before this fix).
        if database_password:
            _patch_secret_field(
                secret_name,
                namespace,
                "PRODIGY_TEAMS_DATABASE_PASSWORD",
                database_password,
            )
        return public_pem
    msg.info("Generating new broker RSA keypair...")
    kp = generate_key_pair()
    # The broker expects base64-encoded PEM values.
    private_key_b64 = base64.b64encode(kp.private_pem.encode()).decode()
    public_key_b64 = base64.b64encode(kp.public_pem.encode()).decode()
    cmd = [
        "kubectl",
        "create",
        "secret",
        "generic",
        secret_name,
        f"--namespace={namespace}",
        f"--from-literal=PRODIGY_TEAMS_PRIVATE_KEY={private_key_b64}",
        f"--from-literal=PRODIGY_TEAMS_PUBLIC_KEY={public_key_b64}",
    ]
    if database_password:
        cmd.append(
            f"--from-literal=PRODIGY_TEAMS_DATABASE_PASSWORD={database_password}"
        )
    cmd += ["--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 _patch_secret_field(secret_name: str, namespace: str, key: str, value: str) -> None:
    """Ensure a single field exists in an existing K8s secret."""
    patch = json.dumps({"stringData": {key: value}})
    subprocess.run(
        [
            "kubectl",
            "patch",
            "secret",
            secret_name,
            f"--namespace={namespace}",
            "--type",
            "merge",
            "-p",
            patch,
        ],
        check=True,
        text=True,
        capture_output=True,
    )
