from __future__ import annotations

import json
import subprocess
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Iterator

from pydantic import BaseModel, Field
from wasabi import msg

if TYPE_CHECKING:
    from prodigy_teams_pam_sdk.models import BrokerCredentials

    from .cluster_config import ClusterConfig

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


class ResourceSpec(BaseModel):
    requests: dict[str, str]
    limits: dict[str, str]


class BridgeService(BaseModel):
    name: str
    port: int
    ip: str


class HostAlias(BaseModel):
    ip: str
    hostnames: list[str]


class DevMounts(BaseModel):
    broker_src_host_path: Annotated[str, Field(alias="brokerSrcHostPath")]
    cpl_src_host_path: Annotated[str, Field(alias="cplSrcHostPath")]
    prodigy_src_host_path: Annotated[str, Field(alias="prodigySrcHostPath")]
    recipes_src_host_path: Annotated[str, Field(alias="recipesSrcHostPath")]
    recipes_sdk_src_host_path: Annotated[str, Field(alias="recipesSdkSrcHostPath")]
    model_config = {"populate_by_name": True}


class PerJob(BaseModel):
    cores: int
    memory_mb: int
    memory_max_mb: int


class WorkerType(BaseModel):
    name: str
    node_class: str
    description: str
    per_job: PerJob


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

    class GCPConfig(BaseModel):
        preemptible: bool = False
        spot: bool = False

    name: str
    node_class: str
    machine_type: str
    target_size: int
    min_size: int | None = None
    max_size: int | None = None
    gpu: GPU | None = None
    gcp: GCPConfig = Field(default_factory=GCPConfig)
    per_job: PerJob

    @property
    def preemptible(self) -> bool:
        return self.gcp.preemptible

    @property
    def spot(self) -> bool:
        return self.gcp.spot

    def to_worker_type(self) -> WorkerType:
        return WorkerType(
            name=self.node_class,
            node_class=self.node_class,
            description=self.name,
            per_job=self.per_job,
        )

    def to_terraform_var(self) -> dict:
        """Convert to a terraform worker_types variable entry."""
        return {
            "name": self.name,
            "node_class": self.node_class,
            "machine_type": self.machine_type,
            "preemptible": self.preemptible,
            "spot": self.spot,
            "min_size": self.min_size if self.min_size is not None else 0,
            "max_size": self.max_size
            if self.max_size is not None
            else max(self.target_size, 1),
            "guest_accelerator": self.gpu.model_dump() if self.gpu else None,
        }


class LocalTls(BaseModel):
    crt: str
    key: str


class HelmChartValues(BaseModel):
    """Schema for our Helm chart overrides.

    It's extremely important that this model functions as a schema. All
    instantiations of this model MUST be semantically valid values to
    use in the rest of the program. Never set a default here that would
    be an invalid value.
    """

    class Image(BaseModel):
        registry: str
        broker_version: Annotated[str, Field(alias="brokerVersion")]
        cpl_version: Annotated[str, Field(alias="cplVersion")]
        pull_policy: Annotated[str, Field(alias="pullPolicy")]
        model_config = {"populate_by_name": True}

    class ServiceAccount(BaseModel):
        annotations: dict[str, str]

    class Broker(BaseModel):
        class BrokerProbe(BaseModel):
            initial_delay_seconds: Annotated[int, Field(alias="initialDelaySeconds")]
            period_seconds: Annotated[int, Field(alias="periodSeconds")]
            model_config = {"populate_by_name": True}

        replicas: int
        resources: dict[str, ResourceSpec]
        readiness_probe: Annotated[BrokerProbe, Field(alias="readinessProbe")]
        liveness_probe: Annotated[BrokerProbe, Field(alias="livenessProbe")]
        cpl_readiness_probe: Annotated[dict, Field(alias="cplReadinessProbe")]
        cpl_liveness_probe: Annotated[dict, Field(alias="cplLivenessProbe")]
        node_selector: Annotated[dict[str, str], Field(alias="nodeSelector")]
        model_config = {"populate_by_name": True}

    class Ingress(BaseModel):
        class TLS(BaseModel):
            secret_name: Annotated[str, Field(alias="secretName")]
            model_config = {"populate_by_name": True}

        class_name: Annotated[str, Field(alias="className")]
        tls: TLS
        model_config = {"populate_by_name": True}

    class Cluster(BaseModel):
        broker_id: Annotated[str, Field(alias="brokerId")]
        fqdn: str
        pam_fqdn: Annotated[str, Field(alias="pamFqdn")]
        container_registry: Annotated[str, Field(alias="containerRegistry")]
        stream_tracebacks: Annotated[bool, Field(alias="streamTracebacks")]
        model_config = {"populate_by_name": True}

    class Database(BaseModel):
        ip: str
        user: str
        name: str

    class CertManager(BaseModel):
        cluster_issuer_name: Annotated[str, Field(alias="clusterIssuerName")]
        acme_email: Annotated[str, Field(alias="acmeEmail")]
        model_config = {"populate_by_name": True}

    class Development(BaseModel):
        verbose: bool
        json_logs: Annotated[bool, Field(alias="jsonLogs")]
        model_config = {"populate_by_name": True}

    class NFS(BaseModel):
        claim_name: Annotated[str, Field(alias="claimName")]
        model_config = {"populate_by_name": True}

    class Storage(BaseModel):
        type: str
        host_path: Annotated[str, Field(alias="hostPath")]
        capacity: str
        model_config = {"populate_by_name": True}

    class Secrets(BaseModel):
        infra_secret_name: Annotated[str, Field(alias="infraSecretName")]
        rollout_checksum: Annotated[str, Field(alias="rolloutChecksum")]
        container_sa_key: Annotated[str, Field(alias="containerSaKey")]
        prodigy_license_key: Annotated[str, Field(alias="prodigyLicenseKey")]
        database_password: Annotated[str, Field(alias="databasePassword")]
        private_key: Annotated[str, Field(alias="privateKey")]
        public_key: Annotated[str, Field(alias="publicKey")]
        model_config = {"populate_by_name": True}

    class GCP(BaseModel):
        project: str
        zone: str

    class ConfigFile(BaseModel):
        broker_config: Annotated[str, Field(alias="brokerConfig")]
        cpl_config: Annotated[str, Field(alias="cplConfig")]
        model_config = {"populate_by_name": True}

    class Logging(BaseModel):
        class Loki(BaseModel):
            class Storage(BaseModel):
                nfs_sub_path: Annotated[str, Field(alias="nfsSubPath")]
                retention_period: Annotated[str, Field(alias="retentionPeriod")]
                model_config = {"populate_by_name": True}

            image: str
            storage: Storage
            resources: ResourceSpec
            node_selector: Annotated[dict[str, str], Field(alias="nodeSelector")]
            model_config = {"populate_by_name": True}

        class Vector(BaseModel):
            image: str
            resources: ResourceSpec

        loki: Loki | None = None
        vector: Vector | None = None
        loki_external_url: Annotated[str, Field(alias="lokiExternalUrl")] = ""
        model_config = {"populate_by_name": True}

    class Local(BaseModel):
        class Storage(BaseModel):
            host_path: Annotated[str, Field(alias="hostPath")]
            capacity: str
            model_config = {"populate_by_name": True}

        storage: Storage
        bridge: dict
        host_aliases: Annotated[list[HostAlias], Field(alias="hostAliases")]
        dev_mounts: Annotated[DevMounts | None, Field(alias="devMounts")] = None
        tls: LocalTls
        reload: bool
        model_config = {"populate_by_name": True}

    class Service(BaseModel):
        type: str
        node_port: Annotated[int | None, Field(alias="nodePort")]
        model_config = {"populate_by_name": True}

    class Postgresql(BaseModel):
        class Auth(BaseModel):
            database: str
            username: str
            password: str

        enabled: bool
        auth: Auth

    name_override: Annotated[str, Field(alias="nameOverride")]
    full_name_override: Annotated[str, Field(alias="fullnameOverride")]
    namespace: dict
    image_pull_secrets: Annotated[list[dict], Field(alias="imagePullSecrets")]
    pod_security_context: Annotated[dict, Field(alias="podSecurityContext")]
    container_security_context: Annotated[dict, Field(alias="containerSecurityContext")]
    worker_types: Annotated[list[WorkerType], Field(alias="workerTypes")]
    prodigy_pypi_index_url: Annotated[str, Field(alias="prodigyPypiIndexUrl")]
    image: Image
    broker: Broker
    ingress: Ingress
    cluster: Cluster
    database: Database
    service_account: Annotated[ServiceAccount, Field(alias="serviceAccount")]
    cert_manager: Annotated[CertManager | None, Field(alias="certManager")] = None
    development: Development | None = None
    nfs: NFS
    storage: Storage
    secrets: Secrets
    config_files: Annotated[ConfigFile, Field(alias="configFiles")]
    logging: Logging | None = None
    local: Local | None = None
    network_policy: Annotated[bool, Field(alias="networkPolicy")]
    service: Service
    postgresql: Postgresql
    gcp: GCP | None = None
    automount_service_account_token: Annotated[
        bool | None, Field(alias="automountServiceAccountToken")
    ]
    model_config = {"populate_by_name": True}

    @contextmanager
    def write_to_temp(self) -> Iterator[str]:
        """Write chart values to a temp JSON file, yielding its path."""
        data = self.model_dump(by_alias=True, exclude_none=True)
        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=True) as tmp:
            tmp.write(json.dumps(data, indent=2))
            tmp.flush()
            yield tmp.name


class HelmChartOverrides(BaseModel):
    """Partial overrides for the Helm chart.

    Every field is ``Optional[T] = None``.  When serialized with
    ``exclude_none=True``, unset fields are omitted and Helm falls back
    to chart defaults.  The submodel types mirror HelmChartValues but
    with all-optional fields so partial construction works (e.g. setting
    3 of 7 Secrets fields).
    """

    class Image(BaseModel):
        registry: str | None = None
        broker_version: Annotated[str | None, Field(alias="brokerVersion")] = None
        cpl_version: Annotated[str | None, Field(alias="cplVersion")] = None
        pull_policy: Annotated[str | None, Field(alias="pullPolicy")] = None
        model_config = {"populate_by_name": True}

    class ServiceAccount(BaseModel):
        annotations: dict[str, str] | None = None

    class Broker(BaseModel):
        class BrokerProbe(BaseModel):
            initial_delay_seconds: Annotated[
                int | None, Field(alias="initialDelaySeconds")
            ] = None
            period_seconds: Annotated[int | None, Field(alias="periodSeconds")] = None
            model_config = {"populate_by_name": True}

        replicas: int | None = None
        resources: dict[str, ResourceSpec] | None = None
        readiness_probe: Annotated[
            BrokerProbe | None, Field(alias="readinessProbe")
        ] = None
        liveness_probe: Annotated[BrokerProbe | None, Field(alias="livenessProbe")] = (
            None
        )
        cpl_readiness_probe: Annotated[
            dict | None, Field(alias="cplReadinessProbe")
        ] = None
        cpl_liveness_probe: Annotated[dict | None, Field(alias="cplLivenessProbe")] = (
            None
        )
        node_selector: Annotated[dict[str, str] | None, Field(alias="nodeSelector")] = (
            None
        )
        model_config = {"populate_by_name": True}

    class Ingress(BaseModel):
        class TLS(BaseModel):
            secret_name: Annotated[str | None, Field(alias="secretName")] = None
            model_config = {"populate_by_name": True}

        class_name: Annotated[str | None, Field(alias="className")] = None
        tls: TLS | None = None
        model_config = {"populate_by_name": True}

    class Cluster(BaseModel):
        broker_id: Annotated[str | None, Field(alias="brokerId")] = None
        fqdn: str | None = None
        pam_fqdn: Annotated[str | None, Field(alias="pamFqdn")] = None
        container_registry: Annotated[str | None, Field(alias="containerRegistry")] = (
            None
        )
        stream_tracebacks: Annotated[bool | None, Field(alias="streamTracebacks")] = (
            None
        )
        model_config = {"populate_by_name": True}

    class Database(BaseModel):
        ip: str | None = None
        user: str | None = None
        name: str | None = None

    class CertManager(BaseModel):
        cluster_issuer_name: Annotated[str | None, Field(alias="clusterIssuerName")] = (
            None
        )
        acme_email: Annotated[str | None, Field(alias="acmeEmail")] = None
        model_config = {"populate_by_name": True}

    class Development(BaseModel):
        verbose: bool | None = None
        json_logs: Annotated[bool | None, Field(alias="jsonLogs")] = None
        model_config = {"populate_by_name": True}

    class NFS(BaseModel):
        claim_name: Annotated[str | None, Field(alias="claimName")] = None
        model_config = {"populate_by_name": True}

    class Storage(BaseModel):
        type: str | None = None
        host_path: Annotated[str | None, Field(alias="hostPath")] = None
        capacity: str | None = None
        model_config = {"populate_by_name": True}

    class Secrets(BaseModel):
        infra_secret_name: Annotated[str | None, Field(alias="infraSecretName")] = None
        rollout_checksum: Annotated[str | None, Field(alias="rolloutChecksum")] = None
        container_sa_key: Annotated[str | None, Field(alias="containerSaKey")] = None
        prodigy_license_key: Annotated[str | None, Field(alias="prodigyLicenseKey")] = (
            None
        )
        database_password: Annotated[str | None, Field(alias="databasePassword")] = None
        private_key: Annotated[str | None, Field(alias="privateKey")] = None
        public_key: Annotated[str | None, Field(alias="publicKey")] = None
        model_config = {"populate_by_name": True}

    class GCP(BaseModel):
        project: str | None = None
        zone: str | None = None

    class ConfigFile(BaseModel):
        broker_config: Annotated[str | None, Field(alias="brokerConfig")] = None
        cpl_config: Annotated[str | None, Field(alias="cplConfig")] = None
        model_config = {"populate_by_name": True}

    class Logging(BaseModel):
        class Loki(BaseModel):
            class Storage(BaseModel):
                nfs_sub_path: Annotated[str | None, Field(alias="nfsSubPath")] = None
                retention_period: Annotated[
                    str | None, Field(alias="retentionPeriod")
                ] = None
                model_config = {"populate_by_name": True}

            image: str | None = None
            storage: Storage | None = None
            resources: ResourceSpec | None = None
            node_selector: Annotated[
                dict[str, str] | None, Field(alias="nodeSelector")
            ] = None
            model_config = {"populate_by_name": True}

        class Vector(BaseModel):
            image: str | None = None
            resources: ResourceSpec | None = None

        loki: Loki | None = None
        vector: Vector | None = None
        loki_external_url: Annotated[str | None, Field(alias="lokiExternalUrl")] = None
        model_config = {"populate_by_name": True}

    class Local(BaseModel):
        class Storage(BaseModel):
            host_path: Annotated[str | None, Field(alias="hostPath")] = None
            capacity: str | None = None
            model_config = {"populate_by_name": True}

        storage: Storage | None = None
        bridge: dict | None = None
        host_aliases: Annotated[list[HostAlias] | None, Field(alias="hostAliases")] = (
            None
        )
        dev_mounts: Annotated[DevMounts | None, Field(alias="devMounts")] = None
        tls: LocalTls | None = None
        reload: bool | None = None
        model_config = {"populate_by_name": True}

    class Service(BaseModel):
        type: str | None = None
        node_port: Annotated[int | None, Field(alias="nodePort")] = None
        model_config = {"populate_by_name": True}

    class Postgresql(BaseModel):
        class Auth(BaseModel):
            database: str | None = None
            username: str | None = None
            password: str | None = None

        enabled: bool | None = None
        auth: Auth | None = None

    name_override: Annotated[str | None, Field(alias="nameOverride")] = None
    full_name_override: Annotated[str | None, Field(alias="fullnameOverride")] = None
    namespace: dict | None = None
    image_pull_secrets: Annotated[
        list[dict] | None, Field(alias="imagePullSecrets")
    ] = None
    pod_security_context: Annotated[dict | None, Field(alias="podSecurityContext")] = (
        None
    )
    container_security_context: Annotated[
        dict | None, Field(alias="containerSecurityContext")
    ] = None
    worker_types: Annotated[list[WorkerType] | None, Field(alias="workerTypes")] = None
    prodigy_pypi_index_url: Annotated[
        str | None, Field(alias="prodigyPypiIndexUrl")
    ] = None
    image: Image | None = None
    broker: Broker | None = None
    ingress: Ingress | None = None
    cluster: Cluster | None = None
    database: Database | None = None
    service_account: Annotated[ServiceAccount | None, Field(alias="serviceAccount")] = (
        None
    )
    cert_manager: Annotated[CertManager | None, Field(alias="certManager")] = None
    development: Development | None = None
    nfs: NFS | None = None
    storage: Storage | None = None
    secrets: Secrets | None = None
    config_files: Annotated[ConfigFile | None, Field(alias="configFiles")] = None
    logging: Logging | None = None
    local: Local | None = None
    network_policy: Annotated[bool | None, Field(alias="networkPolicy")] = None
    service: Service | None = None
    postgresql: Postgresql | None = None
    gcp: GCP | None = None
    automount_service_account_token: Annotated[
        bool | None, Field(alias="automountServiceAccountToken")
    ] = None
    model_config = {"populate_by_name": True}

    @contextmanager
    def write_to_temp(self) -> Iterator[str]:
        """Write override values to a temp JSON file, yielding its path."""
        data = self.model_dump(by_alias=True, exclude_none=True)
        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=True) as tmp:
            tmp.write(json.dumps(data, indent=2))
            tmp.flush()
            yield tmp.name


def _check_overrides_match(
    values_cls: type[BaseModel], overrides_cls: type[BaseModel]
) -> None:
    """Assert that overrides has the same fields/aliases as values.

    Call at import time so drift between HelmChartValues and
    HelmChartOverrides is caught immediately, not at deploy time.
    """
    for name, field_info in values_cls.model_fields.items():
        assert name in overrides_cls.model_fields, (
            f"{overrides_cls.__name__} is missing field '{name}' "
            f"from {values_cls.__name__}"
        )
        override_field = overrides_cls.model_fields[name]
        assert field_info.alias == override_field.alias, (
            f"Alias mismatch for '{name}': "
            f"{values_cls.__name__} has '{field_info.alias}', "
            f"{overrides_cls.__name__} has '{override_field.alias}'"
        )
    # Check nested BaseModel subclasses defined as inner classes
    for attr_name, attr in vars(values_cls).items():
        if isinstance(attr, type) and issubclass(attr, BaseModel):
            override_attr = getattr(overrides_cls, attr_name, None)
            assert override_attr is not None and isinstance(override_attr, type), (
                f"{overrides_cls.__name__} is missing nested class '{attr_name}'"
            )
            _check_overrides_match(attr, override_attr)


_check_overrides_match(HelmChartValues, HelmChartOverrides)


def upgrade_from_ref(
    chart_ref: str,
    values: HelmChartOverrides,
    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 values.write_to_temp() as values_path:
        cmd += ["-f", values_path]
        subprocess.run(cmd, check=True, text=True)


def preflight_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 build_values(
    config: "ClusterConfig",
    creds: "BrokerCredentials",
) -> HelmChartOverrides:
    """Build Helm chart overrides from cluster config and PAM credentials."""
    if config.deployment.acme_email and not config.deployment.tls_secret_name:
        raise ValueError(
            "deployment.tls_secret_name is required when deployment.acme_email is set"
        )
    primary_registry = creds.image.registry or config.deployment.container_registry
    return HelmChartOverrides(
        image=HelmChartOverrides.Image(
            broker_version=creds.image.broker_version,
            cpl_version=creds.image.cpl_version,
            registry=primary_registry,
        ),
        cluster=HelmChartOverrides.Cluster(
            broker_id=str(creds.broker_id),
            fqdn=config.deployment.domain,
            pam_fqdn=config.deployment.pam_domain,
            container_registry=config.deployment.container_registry,
            stream_tracebacks=False,
        ),
        database=HelmChartOverrides.Database(**config.deployment.database.model_dump()),
        nfs=HelmChartOverrides.NFS(claim_name=config.deployment.nfs_pvc_name),
        worker_types=[w.to_worker_type() for w in config.node_pools],
        prodigy_pypi_index_url=config.deployment.prodigy_pypi_index_url or None,
        secrets=HelmChartOverrides.Secrets(
            container_sa_key=creds.secrets.container_sa_key,
            prodigy_license_key=creds.secrets.prodigy_license_key,
            infra_secret_name=config.deployment.infra_secret_name,
        ),
        gcp=HelmChartOverrides.GCP(**config.deployment.gcp.model_dump())
        if config.deployment.gcp.project
        else None,
        service_account=HelmChartOverrides.ServiceAccount(
            annotations={
                "iam.gke.io/gcp-service-account": f"cluster-gke-sa@{config.deployment.gcp.project}.iam.gserviceaccount.com",
            },
        )
        if config.deployment.gcp.project
        else None,
        ingress=HelmChartOverrides.Ingress(
            tls=HelmChartOverrides.Ingress.TLS(
                secret_name=config.deployment.tls_secret_name
            ),
        )
        if config.deployment.tls_secret_name
        else None,
        cert_manager=HelmChartOverrides.CertManager(
            acme_email=config.deployment.acme_email,
        )
        if config.deployment.acme_email
        else None,
    )


def upgrade(
    chart_path: Path,
    values: HelmChartOverrides,
    release: str = "prodigy-teams",
    namespace: str = "prodigy-teams",
    wait: bool = False,
    timeout: str = "10m",
) -> None:
    """Run helm upgrade --install for the prodigy-teams chart."""
    # Ensure chart dependency repos are available (idempotent).
    result = subprocess.run(
        ["helm", "repo", "add", "bitnami", "https://charts.bitnami.com/bitnami"],
        text=True,
        capture_output=True,
    )
    if result.returncode != 0 and "already exists" not in result.stderr:
        raise subprocess.CalledProcessError(
            result.returncode, result.args, result.stdout, result.stderr
        )
    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"]
    with values.write_to_temp() as values_path:
        cmd += ["-f", values_path]
        subprocess.run(cmd, check=True, text=True)
