import asyncio
import json
import logging
import os
from enum import Enum
from pathlib import Path
from typing import Any, ClassVar, Dict, List, Optional, Set

from pydantic import (
    BaseModel,
    Field,
    PrivateAttr,
    field_serializer,
    field_validator,
    model_validator,
)
from runpod.endpoint.runner import Job

from ..api.runpod import RunpodGraphQLClient
from ..exceptions import RunpodAPIKeyError
from ..utils.backoff import get_backoff_delay
from .base import DeployableResource
from .cloud import runpod
from .constants import (
    CONSOLE_URL,
    DEFAULT_WORKERS_MAX,
    DEFAULT_WORKERS_MIN,
    validate_python_version as _validate_python_version,
)
from .cpu import CpuInstanceType
from .gpu import GpuGroup, GpuType
from .network_volume import NetworkVolume, DataCenter, CPU_DATACENTERS
from .template import KeyValuePair, PodTemplate
from .resource_manager import ResourceManager


# Prefix applied to endpoint names during live provisioning
LIVE_PREFIX = "live-"


log = logging.getLogger(__name__)


class ServerlessScalerType(Enum):
    QUEUE_DELAY = "QUEUE_DELAY"
    REQUEST_COUNT = "REQUEST_COUNT"


class ServerlessType(Enum):
    """
    Serverless endpoint execution model.

    QB (Queue-based): Traditional queue processing with automatic retries.
                      Requests are placed in queue and processed sequentially.
                      JSON input/output only. Higher latency but built-in error recovery.

    LB (Load-balancer): Direct HTTP routing to healthy workers.
                        Supports custom HTTP endpoints and any data format.
                        Lower latency but no automatic retries.
    """

    QB = "QB"
    LB = "LB"


class CudaVersion(Enum):
    V11_8 = "11.8"
    V12_0 = "12.0"
    V12_1 = "12.1"
    V12_2 = "12.2"
    V12_3 = "12.3"
    V12_4 = "12.4"
    V12_5 = "12.5"
    V12_6 = "12.6"
    V12_7 = "12.7"
    V12_8 = "12.8"


class ServerlessResource(DeployableResource):
    """
    Base class for GPU serverless resource
    """

    _input_only = {
        "id",
        "cudaVersions",
        "datacenter",
        "env",
        "gpus",
        "flashboot",
        "flashEnvironmentId",
        "imageName",
        "networkVolume",
        "networkVolumes",
        "python_version",
    }

    _hashed_fields = {
        "datacenter",
        "env",
        "gpuIds",
        "executionTimeoutMs",
        "gpuCount",
        "locations",
        "name",
        "networkVolumeId",
        "networkVolumes",
        "python_version",
        "scalerType",
        "scalerValue",
        "workersMax",
        "workersMin",
        "workersPFBTarget",
        "allowedCudaVersions",
        "type",
    }

    # Fields assigned by API that shouldn't affect drift detection
    # When adding new fields to ServerlessResource, evaluate if they are:
    # 1. User-specified (include in hash)
    # 2. API-assigned/runtime (add to RUNTIME_FIELDS)
    # 3. Dynamic identifiers (already excluded via "id")
    RUNTIME_FIELDS: ClassVar[Set[str]] = {
        "template",
        "templateId",
        "aiKey",
        "userId",
        "createdAt",
        "activeBuildid",
        "computeType",
        "hubRelease",
        "repo",
    }

    EXCLUDED_HASH_FIELDS: ClassVar[Set[str]] = {"id"}

    # === Input-only Fields ===
    cudaVersions: Optional[List[CudaVersion]] = []  # for allowedCudaVersions
    env: Optional[Dict[str, str]] = Field(default=None)
    flashboot: Optional[bool] = True
    gpus: Optional[List[GpuGroup | GpuType]] = [GpuGroup.ANY]  # for gpuIds
    imageName: Optional[str] = ""  # for template.imageName
    networkVolume: Optional[NetworkVolume] = None
    networkVolumes: Optional[List[NetworkVolume]] = None
    # accepts a single DataCenter or a list for multi-dc deployments
    datacenter: Optional[List[DataCenter] | DataCenter] = Field(default=None)
    python_version: Optional[str] = Field(
        default=None,
        description="Python version for runtime image selection. Defaults to the local interpreter version at build time.",
    )

    # === Input Fields ===
    executionTimeoutMs: Optional[int] = 0
    gpuCount: Optional[int] = 1
    idleTimeout: Optional[int] = 60
    instanceIds: Optional[List[CpuInstanceType]] = None
    locations: Optional[str] = None
    name: str
    networkVolumeId: Optional[str] = None
    flashEnvironmentId: Optional[str] = None
    scalerType: Optional[ServerlessScalerType] = ServerlessScalerType.QUEUE_DELAY
    scalerValue: Optional[int] = 4
    templateId: Optional[str] = None
    type: Optional[ServerlessType] = ServerlessType.QB
    workersMax: Optional[int] = DEFAULT_WORKERS_MAX
    workersMin: Optional[int] = DEFAULT_WORKERS_MIN
    workersPFBTarget: Optional[int] = 0

    # === Private Attributes ===
    _deployed_volume_ids: list[str] = PrivateAttr(default_factory=list)

    # === Runtime Fields ===
    activeBuildid: Optional[str] = None
    aiKey: Optional[str] = None
    allowedCudaVersions: Optional[str] = ""
    computeType: Optional[str] = None
    createdAt: Optional[str] = None  # TODO: use datetime
    gpuIds: Optional[str] = ""
    hubRelease: Optional[str] = None
    repo: Optional[str] = None
    template: Optional[PodTemplate] = None
    userId: Optional[str] = None

    def __str__(self) -> str:
        return f"{self.__class__.__name__}:{self.id}"

    @property
    def url(self) -> str:
        if not self.id:
            raise ValueError("Missing self.id")
        return CONSOLE_URL % self.id

    @property
    def endpoint(self) -> runpod.Endpoint:
        """
        Returns the Runpod endpoint object for this serverless resource.
        """
        if not self.id:
            raise ValueError("Missing self.id")
        return runpod.Endpoint(self.id)

    @property
    def endpoint_url(self) -> str:
        base_url = self.endpoint.rp_client.endpoint_url_base
        return f"{base_url}/{self.id}"

    @field_serializer("scalerType")
    def serialize_scaler_type(
        self, value: Optional[ServerlessScalerType]
    ) -> Optional[str]:
        """Convert ServerlessScalerType enum to string.

        Handles both enum instances and pre-stringified values that may occur
        during nested model serialization or when values are already deserialized.
        """
        if value is None:
            return None
        return value.value if isinstance(value, ServerlessScalerType) else value

    @field_serializer("type")
    def serialize_type(self, value: Optional[ServerlessType]) -> Optional[str]:
        """Convert ServerlessType enum to string.

        Handles both enum instances and pre-stringified values that may occur
        during nested model serialization or when values are already deserialized.
        """
        if value is None:
            return None
        return value.value if isinstance(value, ServerlessType) else value

    @field_serializer("datacenter")
    def serialize_datacenter(
        self,
        value: Optional[List[DataCenter]],
    ) -> Optional[List[str]]:
        """Convert DataCenter enum list to strings."""
        if value is None:
            return None
        return [dc.value if isinstance(dc, DataCenter) else str(dc) for dc in value]

    @model_validator(mode="after")
    def normalize_network_volumes(self):
        """Merge networkVolume (singular) into networkVolumes list.

        Validates that no two volumes share the same datacenter.
        """
        volumes: list[NetworkVolume] = []
        if self.networkVolumes:
            volumes.extend(self.networkVolumes)
        if self.networkVolume and self.networkVolume not in volumes:
            volumes.append(self.networkVolume)

        if not volumes:
            return self

        # validate one volume per datacenter
        seen_dcs: dict[str, str] = {}
        for v in volumes:
            dc_id = v.dataCenterId.value
            label = v.name or v.id or "unknown"
            if dc_id in seen_dcs:
                raise ValueError(
                    f"Multiple volumes in datacenter {dc_id} "
                    f"('{seen_dcs[dc_id]}' and '{label}'). "
                    f"Only one network volume is allowed per datacenter."
                )
            seen_dcs[dc_id] = label

        self.networkVolumes = volumes
        # keep networkVolume pointing at the first for backward compat
        self.networkVolume = volumes[0] if volumes else None

        return self

    @field_validator("datacenter", mode="before")
    @classmethod
    def normalize_datacenter(cls, value):
        """Normalize datacenter to a list of DataCenter enums.

        Accepts a single DataCenter, a string, a list of either, or None.
        """
        if value is None:
            return None
        if isinstance(value, DataCenter):
            return [value]
        if isinstance(value, str):
            return [DataCenter.from_string(value)]
        if isinstance(value, list):
            result = []
            for item in value:
                if isinstance(item, DataCenter):
                    result.append(item)
                elif isinstance(item, str):
                    result.append(DataCenter.from_string(item))
                else:
                    raise ValueError(f"Invalid datacenter value: {item!r}")
            return result
        raise ValueError(f"Invalid datacenter value: {value!r}")

    @field_validator("gpus")
    @classmethod
    def validate_gpus(cls, value: List[GpuGroup | GpuType]) -> List[GpuGroup | GpuType]:
        """Expand ANY to all GPU groups"""
        if not value:
            return value
        if GpuGroup.ANY in value or GpuType.ANY in value:
            return GpuGroup.all()
        return value

    @field_validator("python_version")
    @classmethod
    def validate_python_version(cls, v: Optional[str]) -> Optional[str]:
        if v is not None:
            _validate_python_version(v)
        return v

    @property
    def config_hash(self) -> str:
        """Get config hash excluding runtime-assigned fields.

        Fields that are None are excluded via exclude_none. When env is None
        (default), it is not included in the hash. When env is explicitly set,
        it IS included and triggers drift detection.

        Hashes user-specified configuration including env vars.
        """
        import hashlib
        import json

        resource_type = self.__class__.__name__

        # Exclude runtime fields and id from hash
        exclude_fields = (
            self.__class__.RUNTIME_FIELDS | self.__class__.EXCLUDED_HASH_FIELDS
        )
        config_dict = self.model_dump(
            exclude_none=True, exclude=exclude_fields, mode="json"
        )

        # strip runtime-assigned IDs from nested network volumes so that
        # pre-deploy and post-deploy hashes match
        for key in ("networkVolume", "networkVolumes"):
            if key not in config_dict:
                continue
            val = config_dict[key]
            if isinstance(val, dict):
                val.pop("id", None)
            elif isinstance(val, list):
                for entry in val:
                    if isinstance(entry, dict):
                        entry.pop("id", None)

        # networkVolumeId is derived from volume deployment, not user config
        config_dict.pop("networkVolumeId", None)

        # Convert to JSON string for hashing
        config_str = json.dumps(config_dict, sort_keys=True)
        hash_obj = hashlib.md5(f"{resource_type}:{config_str}".encode())
        hash_value = hash_obj.hexdigest()

        return hash_value

    @model_validator(mode="after")
    def sync_input_fields(self):
        """Sync between temporary inputs and exported fields.

        Idempotent: Can be called multiple times safely without changing the result.
        """
        # Prepend live- prefix for live provisioning context
        # Must happen BEFORE flashboot suffix to get: live-my-endpoint-fb
        is_live_provisioning = (
            os.getenv("FLASH_IS_LIVE_PROVISIONING", "").lower() == "true"
        )

        if is_live_provisioning:
            # Remove existing live- prefixes for idempotency
            while self.name.startswith(LIVE_PREFIX):
                self.name = self.name[len(LIVE_PREFIX) :]
            # Add prefix once
            self.name = f"{LIVE_PREFIX}{self.name}"

        if self.flashboot and not self.name.endswith("-fb"):
            # Remove all trailing '-fb' suffixes, then add one
            while self.name.endswith("-fb"):
                self.name = self.name[:-3]
            self.name += "-fb"

        # sync datacenter list to locations field for API
        env_locations = os.getenv("RUNPOD_DEFAULT_LOCATIONS")
        env_datacenter = os.getenv("RUNPOD_DEFAULT_DATACENTER")
        if env_locations:
            self.locations = env_locations
        elif not self.locations:
            if env_datacenter:
                try:
                    self.locations = DataCenter(env_datacenter).value
                except ValueError:
                    self.locations = env_datacenter
            elif self.datacenter:
                self.locations = ",".join(dc.value for dc in self.datacenter)

        # validate that all network volume DCs are within the endpoint's datacenter list
        all_volumes = self.networkVolumes or (
            [self.networkVolume] if self.networkVolume else []
        )
        if all_volumes and self.datacenter:
            for vol in all_volumes:
                if vol.dataCenterId not in self.datacenter:
                    dc_values = ", ".join(dc.value for dc in self.datacenter)
                    raise ValueError(
                        f"Network volume datacenter ({vol.dataCenterId.value}) "
                        f"is not in the endpoint's datacenter list ({dc_values})"
                    )

        # backward compat: sync single volume ID for legacy code paths
        if self.networkVolume and self.networkVolume.is_created:
            self.networkVolumeId = self.networkVolume.id

        self._sync_input_fields_gpu()

        return self

    @model_validator(mode="after")
    def validate_cpu_datacenters(self):
        """Ensure CPU endpoints only target data centers that support CPU."""
        if not self._has_cpu_instances():
            return self
        if not self.datacenter:
            return self

        dc_list = self.datacenter
        unsupported = [dc for dc in dc_list if dc not in CPU_DATACENTERS]
        if unsupported:
            unsupported_str = ", ".join(dc.value for dc in unsupported)
            supported_str = ", ".join(
                dc.value for dc in sorted(CPU_DATACENTERS, key=lambda d: d.value)
            )
            raise ValueError(
                f"CPU endpoints are not available in: {unsupported_str}. "
                f"Supported CPU data centers: {supported_str}"
            )
        return self

    @model_validator(mode="after")
    def validate_worker_range(self):
        """Ensure worker scaling bounds are valid."""
        if (
            self.workersMin is not None
            and self.workersMax is not None
            and self.workersMin > self.workersMax
        ):
            raise ValueError(
                f"workersMin ({self.workersMin}) cannot be greater than "
                f"workersMax ({self.workersMax})"
            )
        return self

    def _has_cpu_instances(self) -> bool:
        """Check if endpoint has CPU instances configured.

        Returns:
            True if instanceIds field is present and non-empty, False otherwise.
        """
        return (
            hasattr(self, "instanceIds")
            and self.instanceIds is not None
            and len(self.instanceIds) > 0
        )

    def _get_cpu_disk_limit(self) -> Optional[int]:
        """Calculate max disk size for CPU instances.

        Returns:
            Maximum allowed disk size in GB, or None if no CPU instances.
        """
        if not self._has_cpu_instances():
            return None

        from .cpu import get_max_disk_size_for_instances

        return get_max_disk_size_for_instances(self.instanceIds)

    def _apply_smart_disk_sizing(self, template: PodTemplate) -> None:
        """Apply smart disk sizing based on instance type detection.

        If CPU instances are detected and using the default disk size,
        auto-sizes the disk to the CPU instance limit.

        Args:
            template: PodTemplate to configure.
        """
        cpu_limit = self._get_cpu_disk_limit()

        if cpu_limit is None:
            return  # No CPU instances, keep default

        # Auto-size if using default value
        default_disk_size = PodTemplate.model_fields["containerDiskInGb"].default
        if template.containerDiskInGb == default_disk_size:
            log.debug(
                f"Auto-sizing containerDiskInGb from {default_disk_size}GB "
                f"to {cpu_limit}GB (CPU instance limit)"
            )
            template.containerDiskInGb = cpu_limit

    def _validate_cpu_disk_size(self) -> None:
        """Validate disk size doesn't exceed CPU instance limits.

        Raises:
            ValueError: If disk size exceeds CPU instance limits.
        """
        cpu_limit = self._get_cpu_disk_limit()

        if cpu_limit is None:
            return  # No CPU instances, no validation needed

        if not self.template or not self.template.containerDiskInGb:
            return

        if self.template.containerDiskInGb > cpu_limit:
            from .cpu import CPU_INSTANCE_DISK_LIMITS

            instance_limits = [
                f"{inst.value}: max {CPU_INSTANCE_DISK_LIMITS[inst]}GB"
                for inst in self.instanceIds
            ]

            raise ValueError(
                f"Container disk size {self.template.containerDiskInGb}GB exceeds "
                f"the maximum allowed for CPU instances. "
                f"Instance limits: {', '.join(instance_limits)}. "
                f"Maximum allowed: {cpu_limit}GB. "
                f"Consider using CpuServerlessEndpoint or CpuLiveServerless classes "
                f"for CPU-only deployments."
            )

    def _create_new_template(self) -> PodTemplate:
        """Create a new PodTemplate with standard configuration."""
        return PodTemplate(
            name=self.resource_id,
            imageName=self.imageName,
            env=KeyValuePair.from_dict(self.env or {}),
        )

    def _configure_existing_template(self) -> None:
        """Configure an existing template with necessary overrides."""
        if self.template is None:
            return

        self.template.name = f"{self.resource_id}__{self.template.resource_id}"

        if self.imageName:
            self.template.imageName = self.imageName
        if self.env is not None:
            has_explicit_template_env = "env" in getattr(
                self.template, "model_fields_set", set()
            )
            if self.env or not has_explicit_template_env:
                self.template.env = KeyValuePair.from_dict(self.env)

    async def _sync_graphql_object_with_inputs(
        self, returned_endpoint: "ServerlessResource"
    ):
        for _input_field in self._input_only or set():
            if getattr(self, _input_field) is not None:
                # sync input only fields stripped from gql request back to endpoint
                setattr(returned_endpoint, _input_field, getattr(self, _input_field))

        return returned_endpoint

    def _sync_input_fields_gpu(self):
        # GPU-specific fields (idempotent - only set if not already set)
        if self.gpus and not self.gpuIds:
            # Convert gpus list to gpuIds string
            self.gpuIds = GpuGroup.to_gpu_ids_str(self.gpus)
        elif self.gpuIds and not self.gpus:
            # Convert gpuIds string to gpus list (from backend responses)
            self.gpus = GpuGroup.from_gpu_ids_str(self.gpuIds)

        if self.cudaVersions and not self.allowedCudaVersions:
            # Convert cudaVersions list to allowedCudaVersions string
            self.allowedCudaVersions = ",".join(v.value for v in self.cudaVersions)
        elif self.allowedCudaVersions and not self.cudaVersions:
            # Convert allowedCudaVersions string to cudaVersions list (from backend responses)
            version_values = [
                v.strip() for v in self.allowedCudaVersions.split(",") if v.strip()
            ]
            self.cudaVersions = [CudaVersion(value) for value in version_values]

        return self

    async def _ensure_network_volume_deployed(self) -> None:
        """Ensures all network volumes are deployed.

        Deploys each volume in networkVolumes and collects their IDs.
        Sets networkVolumeId (singular) for backward compat with the first volume.
        Populates _deployed_volume_ids for multi-volume API payloads.
        """
        self._deployed_volume_ids = []

        if self.networkVolumeId:
            self._deployed_volume_ids.append(self.networkVolumeId)

        volumes = self.networkVolumes or (
            [self.networkVolume] if self.networkVolume else []
        )
        for vol in volumes:
            if vol.is_created and vol.id:
                if vol.id not in self._deployed_volume_ids:
                    self._deployed_volume_ids.append(vol.id)
            else:
                deployed = await vol.deploy()
                if deployed.id and deployed.id not in self._deployed_volume_ids:
                    self._deployed_volume_ids.append(deployed.id)

        # backward compat: set singular field from first volume
        if self._deployed_volume_ids and not self.networkVolumeId:
            self.networkVolumeId = self._deployed_volume_ids[0]

    async def is_deployed(self) -> bool:
        """
        Checks if the serverless resource is deployed and available.
        """
        try:
            if not self.id:
                return False

            # During flash run, skip the health check. Newly-created endpoints
            # can fail health checks due to RunPod propagation delay — the
            # endpoint exists but the health API hasn't registered it yet.
            # Trusting the cached ID is correct here; actual failures surface
            # on the first real run/runsync call.
            if os.getenv("FLASH_IS_LIVE_PROVISIONING", "").lower() == "true":
                return True

            response = self.endpoint.health()
            return response is not None
        except Exception as e:
            log.debug(f"Error checking {self}: {e}")
            return False

    def _inject_multi_volume_payload(self, payload: dict) -> None:
        """Replace singular networkVolumeId with networkVolumeIds when multiple volumes are deployed."""
        if len(self._deployed_volume_ids) > 1:
            payload["networkVolumeIds"] = [
                {"networkVolumeId": vid} for vid in self._deployed_volume_ids
            ]
            payload.pop("networkVolumeId", None)

    def _payload_exclude(self) -> Set[str]:
        # flashEnvironmentId is input-only but must be sent when provided
        exclude_fields = set(self._input_only or set())
        exclude_fields.discard("flashEnvironmentId")
        # When templateId is already set, exclude template from the payload.
        # RunPod rejects requests that contain both fields simultaneously.
        # Both can coexist after deploy mutates config (sets templateId while
        # template remains from initialization) — templateId takes precedence.
        if self.templateId:
            exclude_fields.add("template")
        return exclude_fields

    @staticmethod
    def _build_template_update_payload(
        template: PodTemplate,
        template_id: str,
        *,
        skip_env: bool = False,
    ) -> Dict[str, Any]:
        """Build saveTemplate payload from template model.

        Keep this to fields supported by saveTemplate to avoid passing endpoint-only
        fields to the template mutation.

        Args:
            template: Template model with desired configuration.
            template_id: ID of the template to update.
            skip_env: When True, omit ``env`` from the payload so
                saveTemplate preserves the existing template env vars.
                This prevents removing platform-injected vars (e.g.
                PORT, PORT_HEALTH on LB endpoints) when the user's
                env hasn't actually changed.
        """
        template_data = template.model_dump(exclude_none=True, mode="json")
        allowed_fields = {
            "name",
            "imageName",
            "containerDiskInGb",
            "dockerArgs",
            "env",
            "readme",
        }
        if skip_env:
            allowed_fields.discard("env")
        payload = {
            key: value for key, value in template_data.items() if key in allowed_fields
        }
        # savetemplate mutation requires volumeInGb, but for sls this is always 0
        payload["volumeInGb"] = 0
        payload["id"] = template_id
        return payload

    def _check_makes_remote_calls(self) -> bool:
        """Check if resource makes remote calls from build manifest.

        Reads flash_manifest.json to determine if this resource config
        has makes_remote_calls=True.

        Returns:
            True if makes remote calls, False if local-only,
            True (safe default) if manifest not found.
        """
        try:
            manifest_path = Path.cwd() / "flash_manifest.json"
            if not manifest_path.exists():
                # Try alternative locations
                manifest_path = Path("/flash_manifest.json")  # Container path

            if not manifest_path.exists():
                log.debug("Manifest not found, assuming makes_remote_calls=True")
                return True  # Safe default

            with open(manifest_path) as f:
                manifest_data = json.load(f)

            resources = manifest_data.get("resources", {})

            # Strip -fb suffix and live- prefix to match manifest name
            lookup_name = self.name
            if lookup_name.endswith("-fb"):
                lookup_name = lookup_name[:-3]
            if lookup_name.startswith(LIVE_PREFIX):
                lookup_name = lookup_name[len(LIVE_PREFIX) :]

            resource_config = resources.get(lookup_name)

            if not resource_config:
                log.debug(
                    f"Resource '{lookup_name}' (from '{self.name}') not in manifest, assuming makes_remote_calls=True"
                )
                return True  # Safe default

            makes_remote_calls = resource_config.get("makes_remote_calls", True)
            log.debug(
                f"Resource '{lookup_name}' (from '{self.name}') makes_remote_calls={makes_remote_calls}"
            )
            return makes_remote_calls

        except Exception as e:
            log.warning(
                f"Failed to read manifest: {e}, assuming makes_remote_calls=True"
            )
            return True  # Safe default on error

    def _get_module_path(self) -> Optional[str]:
        """Get module_path from build manifest for this resource.

        Returns:
            Dotted module path (e.g., 'preprocess.first_pass'), or None if not found.
        """
        try:
            manifest_path = Path.cwd() / "flash_manifest.json"
            if not manifest_path.exists():
                manifest_path = Path("/flash_manifest.json")
            if not manifest_path.exists():
                return None

            with open(manifest_path) as f:
                manifest_data = json.load(f)

            resources = manifest_data.get("resources", {})

            lookup_name = self.name
            if lookup_name.endswith("-fb"):
                lookup_name = lookup_name[:-3]
            if lookup_name.startswith(LIVE_PREFIX):
                lookup_name = lookup_name[len(LIVE_PREFIX) :]

            resource_config = resources.get(lookup_name)
            if not resource_config:
                return None

            return resource_config.get("module_path")

        except Exception:
            return None

    def _inject_template_env(self, key: str, value: str) -> None:
        """Append a KeyValuePair to self.template.env if the key isn't already present.

        This injects runtime env vars directly into the template without
        mutating self.env, which would cause false config drift on subsequent
        deploys.
        """
        if self.template is None:
            return
        if self.template.env is None:
            self.template.env = []
        existing_keys = {kv.key for kv in self.template.env}
        if key not in existing_keys:
            self.template.env.append(KeyValuePair(key=key, value=value))

    def _inject_runtime_template_vars(self) -> None:
        """Inject runtime env vars into template.env without mutating self.env.

        For QB endpoints making remote calls: injects RUNPOD_API_KEY.
        For LB endpoints: injects FLASH_MODULE_PATH.

        Called by both _do_deploy (initial) and update (env changes) so
        runtime vars survive template updates.
        """
        env_dict = self.env or {}

        if self.type == ServerlessType.QB:
            if self._check_makes_remote_calls():
                if "RUNPOD_API_KEY" not in env_dict:
                    from runpod_flash.core.credentials import get_api_key

                    api_key = get_api_key()
                    if api_key:
                        self._inject_template_env("RUNPOD_API_KEY", api_key)
                        log.debug(
                            f"{self.name}: Injected RUNPOD_API_KEY for remote calls "
                            f"(makes_remote_calls=True)"
                        )
                    else:
                        log.warning(
                            f"{self.name}: makes_remote_calls=True but RUNPOD_API_KEY not set. "
                            f"Remote calls to other endpoints will fail."
                        )

        elif self.type == ServerlessType.LB:
            module_path = self._get_module_path()
            if module_path and "FLASH_MODULE_PATH" not in env_dict:
                self._inject_template_env("FLASH_MODULE_PATH", module_path)
                log.debug(f"{self.name}: Injected FLASH_MODULE_PATH={module_path}")

    async def _preserve_platform_env(
        self,
        client: "RunpodGraphQLClient",
        template_id: str,
        old_env: Optional[Dict[str, str]] = None,
    ) -> None:
        """Read current template env and re-add platform-injected vars.

        The platform injects env vars (e.g. PORT, PORT_HEALTH) once at
        initial deploy and does not re-inject them on saveTemplate.  When
        we send a full env replacement we must carry those vars forward.

        A live template key is considered platform-injected only if it
        was NOT in the previous user config (old_env).  Keys that were
        in old_env but are absent from the new config were intentionally
        removed by the user and must not be resurrected.
        """
        if self.template is None:
            return

        try:
            live = await client.get_template(template_id)
        except Exception:
            log.debug(
                f"{self.name}: Could not fetch template '{template_id}', "
                f"skipping platform env preservation"
            )
            return

        live_env = live.get("env") or []
        if not live_env:
            return

        new_keys = {kv.key for kv in (self.template.env or [])}
        old_keys = set(old_env or {})
        for entry in live_env:
            key = entry.get("key", "")
            if not key or key in new_keys:
                continue
            # Key was in old user config — user intentionally removed it
            if key in old_keys:
                log.debug(f"{self.name}: User removed env var '{key}', not preserving")
                continue
            self.template.env = self.template.env or []
            self.template.env.append(
                KeyValuePair(key=key, value=entry.get("value", ""))
            )
            log.debug(f"{self.name}: Preserved platform env var '{key}'")

    async def _do_deploy(self) -> "DeployableResource":
        """
        Deploys the serverless resource using the provided configuration.

        For queue-based endpoints that make remote calls, injects RUNPOD_API_KEY.
        For load-balanced endpoints, injects FLASH_MODULE_PATH.

        Returns a DeployableResource object.
        """
        try:
            # If the resource is already deployed, return it
            if await self.is_deployed():
                log.debug(f"{self} exists")
                return self

            self._inject_runtime_template_vars()

            # Ensure network volumes are deployed first
            await self._ensure_network_volume_deployed()

            async with RunpodGraphQLClient() as client:
                payload = self.model_dump(
                    exclude=self._payload_exclude(), exclude_none=True, mode="json"
                )

                # inject multi-volume IDs if available
                self._inject_multi_volume_payload(payload)

                result = await client.save_endpoint(payload)

            if endpoint := self.__class__(**result):
                endpoint = await self._sync_graphql_object_with_inputs(endpoint)
                self.id = endpoint.id
                self.templateId = endpoint.templateId
                self.template = (
                    None  # templateId takes precedence; clear to avoid conflict
                )
                return endpoint

            raise ValueError("Deployment failed, no endpoint was returned.")

        except RunpodAPIKeyError:
            raise
        except Exception as e:
            log.error(f"{self} failed to deploy: {e}")
            raise

    async def update(self, new_config: "ServerlessResource") -> "ServerlessResource":
        """Update existing endpoint with new configuration.

        Uses saveEndpoint mutation which handles both version-triggering and
        rolling changes. Version-triggering changes (GPU, template, volumes)
        automatically increment version and trigger worker recreation server-side.

        Args:
            new_config: New configuration to apply

        Returns:
            Updated ServerlessResource instance

        Raises:
            ValueError: If endpoint not deployed or update fails
        """
        if not self.id:
            raise ValueError("Cannot update: endpoint not deployed")

        try:
            resolved_template_id = self.templateId or new_config.templateId
            # Check for version-triggering changes
            if not self._has_structural_changes(new_config):
                log.info(f"Updating endpoint '{self.name}' (ID: {self.id})")

            # Ensure network volumes are deployed if specified
            await new_config._ensure_network_volume_deployed()

            async with RunpodGraphQLClient() as client:
                # Include the endpoint ID to trigger update
                payload = new_config.model_dump(
                    exclude=new_config._payload_exclude(),
                    exclude_none=True,
                    mode="json",
                )
                payload["id"] = self.id  # Critical: include ID for update

                # inject multi-volume IDs if available
                new_config._inject_multi_volume_payload(payload)

                result = await client.save_endpoint(payload)
                resolved_template_id = (
                    result.get("templateId") or self.templateId or new_config.templateId
                )

                if new_config.template:
                    if resolved_template_id:
                        # Skip env in the template payload when the user's env
                        # hasn't changed.  This lets the platform keep vars it
                        # injected (e.g. PORT, PORT_HEALTH on LB endpoints)
                        # and avoids a spurious rolling release.
                        #
                        # Also check template.env: if env is empty but the
                        # caller provided explicit template env entries, those
                        # must not be silently dropped.
                        env_unchanged = self.env == new_config.env
                        template_fields_set = getattr(
                            new_config.template, "model_fields_set", set()
                        )
                        has_explicit_template_env = (
                            not new_config.env and "env" in template_fields_set
                        )
                        skip_env = env_unchanged and not has_explicit_template_env

                        if not skip_env:
                            # Inject runtime vars (RUNPOD_API_KEY, FLASH_MODULE_PATH)
                            # so they survive the template env overwrite.
                            new_config._inject_runtime_template_vars()

                            # Preserve platform-injected env vars (e.g. PORT,
                            # PORT_HEALTH) that the platform sets once at initial
                            # deploy and does not re-inject on template updates.
                            await new_config._preserve_platform_env(
                                client, resolved_template_id, self.env
                            )

                        template_payload = self._build_template_update_payload(
                            new_config.template,
                            resolved_template_id,
                            skip_env=skip_env,
                        )
                        await client.update_template(template_payload)
                        log.debug(
                            f"Updated template '{resolved_template_id}' for endpoint '{self.name}'"
                        )
                    else:
                        log.warning(
                            "Template provided during endpoint update but no templateId could "
                            "be resolved; skipping separate saveTemplate call"
                        )

            if updated := self.__class__(**result):
                if not updated.templateId:
                    updated.templateId = (
                        resolved_template_id or self.templateId or new_config.templateId
                    )
                # Keep local input-only state on the hydrated model. The GraphQL
                # response does not include many user-provided fields (for example
                # env, networkVolume, datacenter), and dropping them causes
                # repeated false drift on subsequent deploys.
                updated = await new_config._sync_graphql_object_with_inputs(updated)
                log.debug(
                    f"Successfully updated endpoint '{self.name}' (ID: {self.id})"
                )
                return updated

            raise ValueError("Update failed, no endpoint was returned.")

        except Exception as e:
            log.error(f"Failed to update {self.name}: {e}")
            raise

    def _has_structural_changes(self, new_config: "ServerlessResource") -> bool:
        """Check if config changes are version-triggering.

        Version-triggering changes cause server-side version increment and
        worker recreation:
        - Image changes (imageName via templateId)
        - GPU configuration (gpus, gpuIds, allowedCudaVersions, gpuCount)
        - Hardware allocation (instanceIds, locations)
        - Storage changes (networkVolumeId)
        - Flashboot toggle

        Rolling changes (no version increment):
        - Worker scaling (workersMin, workersMax)
        - Scaler configuration (scalerType, scalerValue)
        - Timeout values (idleTimeout, executionTimeoutMs)
        - Environment variables (env)

        Note: This method is now informational for logging. The actual
        version-triggering logic runs server-side when saveEndpoint is called.

        Runtime fields (template, templateId, aiKey, userId) are excluded
        to prevent false positives when comparing deployed vs new config.

        Args:
            new_config: New configuration to compare against

        Returns:
            True if version-triggering changes detected (workers will be recreated)
        """
        structural_fields = [
            "gpus",
            "gpuIds",
            "imageName",
            "flashboot",
            "allowedCudaVersions",
            "cudaVersions",
            "instanceIds",
        ]

        for field in structural_fields:
            old_val = getattr(self, field, None)
            new_val = getattr(new_config, field, None)

            # Handle list comparison
            if isinstance(old_val, list) and isinstance(new_val, list):
                if sorted(str(v) for v in old_val) != sorted(str(v) for v in new_val):
                    return True
            # Handle other types
            elif old_val != new_val:
                return True

        return False

    async def deploy(self) -> "DeployableResource":
        resource_manager = ResourceManager()
        resource = await resource_manager.get_or_deploy_resource(self)
        # hydrate the id onto the resource so it's usable when this is called directly
        # on a config
        self.id = resource.id
        self.templateId = getattr(resource, "templateId", None)
        return self

    async def _do_undeploy(self) -> bool:
        """
        Undeploys (deletes) the serverless endpoint.

        If deletion fails, verifies the endpoint still exists. If not, treats it as
        successfully undeployed (handles cases where endpoint was deleted externally).

        Returns:
            True if successfully undeployed or endpoint doesn't exist, False otherwise
        """
        if not self.id:
            log.warning(f"{self} has no endpoint ID, cannot undeploy")
            return False

        try:
            async with RunpodGraphQLClient() as client:
                result = await client.delete_endpoint(self.id)
                success = result.get("success", False)

                if success:
                    log.debug(f"{self} successfully undeployed")
                    return True
                else:
                    log.debug(f"{self} failed to undeploy")
                    return False

        except Exception as e:
            log.debug(f"{self} failed to undeploy: {e}")

            # Deletion failed. Check if endpoint still exists.
            # If it doesn't exist, treat as successful cleanup (orphaned endpoint).
            try:
                async with RunpodGraphQLClient() as client:
                    if not await client.endpoint_exists(self.id):
                        log.debug(
                            f"{self} no longer exists on RunPod, removing from cache"
                        )
                        return True
            except Exception as check_error:
                log.warning(f"Could not verify endpoint existence: {check_error}")

            return False

    async def undeploy(self) -> Dict[str, Any]:
        resource_manager = ResourceManager()
        result = await resource_manager.undeploy_resource(self.resource_id)
        log.debug(f"undeployment result: {result}")
        return result

    async def runsync(self, payload: Dict[str, Any]) -> "JobOutput":
        """
        Executes a serverless endpoint request with the payload.
        Returns a JobOutput object.
        """
        if not self.id:
            raise ValueError("Serverless is not deployed")

        def _fetch_job():
            log.info(f"{self} | API /runsync")
            return self.endpoint.rp_client.post(
                f"{self.id}/runsync", payload, timeout=60
            )

        try:
            # log.debug(f"[{self}] Payload: {payload}")

            response = await asyncio.to_thread(_fetch_job)
            return JobOutput(**response)

        except Exception as e:
            health = await asyncio.to_thread(self.endpoint.health)
            health = ServerlessHealth(**health)
            log.debug(f"{self} | Health {health.workers.status}")
            log.error(f"{self} | Exception: {e}")
            raise

    async def run(self, payload: Dict[str, Any]) -> "JobOutput":
        """
        Executes a serverless endpoint async request with the payload.
        Returns a JobOutput object.
        """
        if not self.id:
            raise ValueError("Serverless is not deployed")

        job: Optional[Job] = None

        try:
            # log.debug(f"[{self}] Payload: {payload}")

            # Create a job using the endpoint
            log.info(f"{self} | API /run")
            job = await asyncio.to_thread(self.endpoint.run, request_input=payload)

            log_subgroup = f"Job:{job.job_id}"

            log.info(f"{self} | Started {log_subgroup}")

            current_pace = 0
            attempt = 0
            job_status = Status.UNKNOWN
            last_status = job_status

            # Poll for job status
            while True:
                await asyncio.sleep(current_pace)

                # Check job status
                job_status = await asyncio.to_thread(job.status)

                if last_status == job_status:
                    # nothing changed, increase the gap
                    attempt += 1
                    indicator = "." * (attempt // 2) if attempt % 2 == 0 else ""
                    if indicator:
                        log.info(f"{log_subgroup} | {indicator}")
                else:
                    # status changed, reset the gap
                    log.info(f"{log_subgroup} | Status: {job_status}")
                    attempt = 0

                last_status = job_status

                # Adjust polling pace appropriately
                current_pace = get_backoff_delay(attempt)

                if job_status in ("COMPLETED", "FAILED", "CANCELLED"):
                    response = await asyncio.to_thread(job._fetch_job)
                    return JobOutput(**response)

        except Exception as e:
            if job and job.job_id:
                log.debug(f"{self} | Cancelling job {job.job_id}")
                await asyncio.to_thread(job.cancel)

            log.error(f"{self} | Exception: {e}")
            raise


class ServerlessEndpoint(ServerlessResource):
    """
    Represents a serverless endpoint distinct from a live serverless.
    Inherits from ServerlessResource.
    """

    @model_validator(mode="after")
    def validate_instance_mutual_exclusivity(self):
        """Ensure gpuIds and instanceIds are mutually exclusive.

        When instanceIds is specified, clears GPU configuration since CPU and GPU
        are mutually exclusive resources. Prevents mixing GPU and CPU configurations.
        """
        has_cpu = (
            hasattr(self, "instanceIds")
            and self.instanceIds is not None
            and len(self.instanceIds) > 0
        )

        if has_cpu:
            # Clear GPU configuration if CPU instances are specified
            # This makes CPU intent explicit
            self.gpus = []
            self.gpuIds = ""
            self.gpuCount = 0

        return self

    @model_validator(mode="after")
    def set_serverless_template(self):
        """Create template from imageName if not provided.

        Must run after sync_input_fields to ensure all input fields are synced.
        Applies smart disk sizing and validates configuration.
        """
        if not any([self.imageName, self.template, self.templateId]):
            raise ValueError(
                "Either imageName, template, or templateId must be provided"
            )

        if not self.templateId and not self.template:
            self.template = self._create_new_template()
            # Apply smart disk sizing to new template
            self._apply_smart_disk_sizing(self.template)
        elif self.template:
            self._configure_existing_template()
            # Apply smart disk sizing to existing template
            self._apply_smart_disk_sizing(self.template)

        # Validate CPU disk size if applicable
        self._validate_cpu_disk_size()

        return self


class JobOutput(BaseModel):
    id: str
    workerId: str
    status: str
    delayTime: int
    executionTime: int
    output: Optional[Any] = None
    error: Optional[str] = ""

    def model_post_init(self, _: Any) -> None:
        log_group = f"Worker:{self.workerId}"
        log.info(f"{log_group} | Delay Time: {self.delayTime} ms")
        log.info(f"{log_group} | Execution Time: {self.executionTime} ms")


class Status(str, Enum):
    READY = "READY"
    INITIALIZING = "INITIALIZING"
    THROTTLED = "THROTTLED"
    UNHEALTHY = "UNHEALTHY"
    UNKNOWN = "UNKNOWN"


class WorkersHealth(BaseModel):
    idle: int
    initializing: int
    ready: int
    running: int
    throttled: int
    unhealthy: int

    @property
    def status(self) -> Status:
        if self.ready or self.idle or self.running:
            return Status.READY

        if self.initializing:
            return Status.INITIALIZING

        if self.throttled:
            return Status.THROTTLED

        if self.unhealthy:
            return Status.UNHEALTHY

        return Status.UNKNOWN


class JobsHealth(BaseModel):
    completed: int
    failed: int
    inProgress: int
    inQueue: int
    retried: int


class ServerlessHealth(BaseModel):
    workers: WorkersHealth
    jobs: JobsHealth

    @property
    def is_ready(self) -> bool:
        return self.workers.status == Status.READY
