"""Provider-neutral remote execution environment models and scheduling."""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Mapping, Protocol, Sequence

WORKFLOW_REMOTE_CONFIG_KEY = "remote_execution"
WORKFLOW_REMOTE_POLICY_KEY = "policy"
WORKFLOW_REMOTE_AUTO_PLACE_KEY = "auto_place"
WORKFLOW_REMOTE_PREFERRED_ENVIRONMENTS_KEY = "preferred_environment_ids"
WORKFLOW_REMOTE_MINIMUM_VRAM_GB_KEY = "minimum_vram_gb"


class ExecutionProvider(str, Enum):
    """Identify one implementation of remote workflow execution."""

    MODAL = "modal"
    SSH_DOCKER = "ssh_docker"


class ExecutionPolicy(str, Enum):
    """Describe how eligible workflow components select a provider."""

    MODAL = "modal"
    SELF_HOSTED = "self_hosted"
    AUTOMATIC = "automatic"


class EnvironmentHealth(str, Enum):
    """Describe whether a remote environment may receive new work."""

    UNKNOWN = "unknown"
    READY = "ready"
    DEGRADED = "degraded"
    UNAVAILABLE = "unavailable"
    DRAINING = "draining"


@dataclass(frozen=True)
class GpuCapability:
    """Describe one GPU visible to a remote execution environment."""

    uuid: str
    name: str
    total_vram_bytes: int
    free_vram_bytes: int | None = None
    compute_capability: str | None = None
    driver_version: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-compatible representation of this GPU."""
        return {
            "uuid": self.uuid,
            "name": self.name,
            "total_vram_bytes": self.total_vram_bytes,
            "free_vram_bytes": self.free_vram_bytes,
            "compute_capability": self.compute_capability,
            "driver_version": self.driver_version,
        }

    @classmethod
    def from_dict(cls, payload: Mapping[str, Any]) -> "GpuCapability":
        """Build one validated GPU capability from persisted data."""
        return cls(
            uuid=str(payload.get("uuid") or "").strip(),
            name=str(payload.get("name") or "Unknown GPU").strip(),
            total_vram_bytes=_non_negative_int(
                payload.get("total_vram_bytes"), "total_vram_bytes"
            ),
            free_vram_bytes=_optional_non_negative_int(
                payload.get("free_vram_bytes"), "free_vram_bytes"
            ),
            compute_capability=_optional_string(payload.get("compute_capability")),
            driver_version=_optional_string(payload.get("driver_version")),
        )


@dataclass(frozen=True)
class EnvironmentCapabilities:
    """Describe resources and runtime features discovered on one environment."""

    architecture: str
    operating_system: str
    cpu_count: int
    total_ram_bytes: int
    available_ram_bytes: int | None
    available_disk_bytes: int | None
    docker_version: str
    docker_rootless: bool
    nvidia_container_runtime: bool
    gpus: tuple[GpuCapability, ...] = ()
    probed_at_epoch: float | None = None
    reported_cost_usd_per_second: float | None = None

    @property
    def maximum_vram_bytes(self) -> int:
        """Return the largest total VRAM value exposed by one GPU."""
        return max((gpu.total_vram_bytes for gpu in self.gpus), default=0)

    @property
    def maximum_free_vram_bytes(self) -> int:
        """Return the largest currently free VRAM value when reported."""
        return max(
            (
                gpu.free_vram_bytes
                for gpu in self.gpus
                if gpu.free_vram_bytes is not None
            ),
            default=0,
        )

    @property
    def maximum_available_vram_bytes(self) -> int:
        """Return probed free VRAM when known, otherwise nameplate capacity."""
        reported_free = [
            gpu.free_vram_bytes for gpu in self.gpus if gpu.free_vram_bytes is not None
        ]
        if reported_free:
            return max(reported_free)
        return self.maximum_vram_bytes

    @property
    def schedulable_ram_bytes(self) -> int:
        """Return currently available RAM when known, otherwise total RAM."""
        if self.available_ram_bytes is not None:
            return self.available_ram_bytes
        return self.total_ram_bytes

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-compatible representation of these capabilities."""
        return {
            "architecture": self.architecture,
            "operating_system": self.operating_system,
            "cpu_count": self.cpu_count,
            "total_ram_bytes": self.total_ram_bytes,
            "available_ram_bytes": self.available_ram_bytes,
            "available_disk_bytes": self.available_disk_bytes,
            "docker_version": self.docker_version,
            "docker_rootless": self.docker_rootless,
            "nvidia_container_runtime": self.nvidia_container_runtime,
            "gpus": [gpu.to_dict() for gpu in self.gpus],
            "probed_at_epoch": self.probed_at_epoch,
            "reported_cost_usd_per_second": self.reported_cost_usd_per_second,
        }

    @classmethod
    def from_dict(cls, payload: Mapping[str, Any]) -> "EnvironmentCapabilities":
        """Build validated capabilities from persisted data."""
        raw_gpus = payload.get("gpus")
        gpu_payloads = raw_gpus if isinstance(raw_gpus, list) else []
        return cls(
            architecture=str(payload.get("architecture") or "unknown").strip(),
            operating_system=str(payload.get("operating_system") or "unknown").strip(),
            cpu_count=_non_negative_int(payload.get("cpu_count"), "cpu_count"),
            total_ram_bytes=_non_negative_int(
                payload.get("total_ram_bytes"), "total_ram_bytes"
            ),
            available_ram_bytes=_optional_non_negative_int(
                payload.get("available_ram_bytes"), "available_ram_bytes"
            ),
            available_disk_bytes=_optional_non_negative_int(
                payload.get("available_disk_bytes"), "available_disk_bytes"
            ),
            docker_version=str(payload.get("docker_version") or "unknown").strip(),
            docker_rootless=bool(payload.get("docker_rootless", False)),
            nvidia_container_runtime=bool(
                payload.get("nvidia_container_runtime", False)
            ),
            gpus=tuple(
                GpuCapability.from_dict(gpu_payload)
                for gpu_payload in gpu_payloads
                if isinstance(gpu_payload, Mapping)
            ),
            probed_at_epoch=_optional_non_negative_float(
                payload.get("probed_at_epoch"), "probed_at_epoch"
            ),
            reported_cost_usd_per_second=_optional_non_negative_float(
                payload.get("reported_cost_usd_per_second"),
                "reported_cost_usd_per_second",
            ),
        )


@dataclass(frozen=True)
class ComponentResourceRequirements:
    """Describe hard requirements and cost inputs for one remote component."""

    minimum_vram_bytes: int = 0
    minimum_ram_bytes: int = 0
    gpu_required: bool = True
    architecture: str | None = None
    estimated_execution_seconds: float = 0.0
    estimated_execution_seconds_by_environment: Mapping[str, float] = field(
        default_factory=dict
    )
    estimated_transfer_seconds: float = 0.0
    required_tags: frozenset[str] = frozenset()
    preferred_environment_ids: tuple[str, ...] = ()
    required_provider: ExecutionProvider | None = None


@dataclass(frozen=True)
class EnvironmentSchedulingState:
    """Combine configured and transient state used for scheduling."""

    environment_id: str
    provider: ExecutionProvider
    enabled: bool
    health: EnvironmentHealth
    cost_usd_per_second: float | None
    capabilities: EnvironmentCapabilities | None
    tags: frozenset[str] = frozenset()
    active_workers: int = 0
    maximum_workers: int = 1
    queue_delay_seconds: float = 0.0
    cold_start_seconds: float = 0.0
    cached_artifact_keys: frozenset[str] = frozenset()


@dataclass(frozen=True)
class ExecutionAssignment:
    """Record the selected environment and its predicted scheduling cost."""

    environment_id: str
    provider: ExecutionProvider
    predicted_cost_usd: float | None
    predicted_completion_seconds: float
    reasons: tuple[str, ...] = ()


@dataclass(frozen=True)
class WorkflowExecutionPreferences:
    """Describe provider and placement preferences saved with one workflow."""

    policy: ExecutionPolicy = ExecutionPolicy.MODAL
    auto_place: bool = False
    preferred_environment_ids: tuple[str, ...] = ()
    minimum_vram_bytes: int = 0

    @classmethod
    def from_workflow(cls, workflow: object) -> "WorkflowExecutionPreferences":
        """Read generic preferences while preserving legacy Modal defaults."""
        if not isinstance(workflow, Mapping):
            return cls()
        extra = workflow.get("extra")
        if not isinstance(extra, Mapping):
            return cls()
        payload = extra.get(WORKFLOW_REMOTE_CONFIG_KEY)
        if not isinstance(payload, Mapping):
            return cls()
        raw_preferred = payload.get(WORKFLOW_REMOTE_PREFERRED_ENVIRONMENTS_KEY)
        preferred_values = raw_preferred if isinstance(raw_preferred, list) else []
        raw_minimum_vram_gb = payload.get(WORKFLOW_REMOTE_MINIMUM_VRAM_GB_KEY, 0)
        minimum_vram_gb = _optional_non_negative_float(
            raw_minimum_vram_gb,
            WORKFLOW_REMOTE_MINIMUM_VRAM_GB_KEY,
        )
        return cls(
            policy=ExecutionPolicy(
                str(payload.get(WORKFLOW_REMOTE_POLICY_KEY) or "modal").strip().lower()
            ),
            auto_place=bool(payload.get(WORKFLOW_REMOTE_AUTO_PLACE_KEY, False)),
            preferred_environment_ids=tuple(
                str(value).strip() for value in preferred_values if str(value).strip()
            ),
            minimum_vram_bytes=int((minimum_vram_gb or 0.0) * 1024**3),
        )


class ExecutionEnvironmentClient(Protocol):
    """Provider-neutral client used to prepare and invoke remote environments."""

    def execute_payload(
        self,
        payload: Mapping[str, Any],
        kwargs: Mapping[str, Any],
    ) -> Sequence[Any]:
        """Execute one component payload and return deserialized outputs."""

    async def execute_payload_async(
        self,
        payload: Mapping[str, Any],
        kwargs: Mapping[str, Any],
    ) -> Sequence[Any]:
        """Execute one component payload asynchronously."""


class NoCompatibleExecutionEnvironmentError(RuntimeError):
    """Raised when no configured environment satisfies component requirements."""


class CostAwareEnvironmentScheduler:
    """Select the least-cost compatible environment with deterministic tie-breaking."""

    def choose(
        self,
        environments: Sequence[EnvironmentSchedulingState],
        requirements: ComponentResourceRequirements,
    ) -> ExecutionAssignment:
        """Return the best currently compatible execution environment."""
        candidates: list[tuple[tuple[Any, ...], ExecutionAssignment]] = []
        rejection_reasons: list[str] = []
        for environment in environments:
            rejection = self._incompatibility_reason(environment, requirements)
            if rejection is not None:
                rejection_reasons.append(f"{environment.environment_id}: {rejection}")
                continue

            execution_seconds = (
                requirements.estimated_execution_seconds_by_environment.get(
                    environment.environment_id,
                    requirements.estimated_execution_seconds,
                )
            )
            completion_seconds = max(
                0.0,
                environment.queue_delay_seconds
                + environment.cold_start_seconds
                + requirements.estimated_transfer_seconds
                + execution_seconds,
            )
            predicted_cost = (
                completion_seconds * environment.cost_usd_per_second
                if environment.cost_usd_per_second is not None
                else None
            )
            preferred_rank = self._preferred_rank(
                environment.environment_id,
                requirements.preferred_environment_ids,
            )
            unknown_cost_rank = 1 if predicted_cost is None else 0
            cost_rank = predicted_cost if predicted_cost is not None else math.inf
            assignment = ExecutionAssignment(
                environment_id=environment.environment_id,
                provider=environment.provider,
                predicted_cost_usd=predicted_cost,
                predicted_completion_seconds=completion_seconds,
                reasons=(
                    f"compatible with {len(environment.capabilities.gpus) if environment.capabilities else 0} GPU(s)",
                    (
                        "requires at least "
                        f"{requirements.minimum_vram_bytes / 1024**3:.2f} GiB "
                        "GPU VRAM"
                        if requirements.minimum_vram_bytes
                        else "no explicit GPU VRAM floor"
                    ),
                    (
                        f"estimated cost ${predicted_cost:.6f}"
                        if predicted_cost is not None
                        else "cost unknown"
                    ),
                    f"estimated execution {execution_seconds:.3f}s",
                ),
            )
            candidates.append(
                (
                    (
                        preferred_rank,
                        unknown_cost_rank,
                        cost_rank,
                        completion_seconds,
                        environment.environment_id,
                    ),
                    assignment,
                )
            )

        if not candidates:
            details = "; ".join(rejection_reasons) or "no environments are configured"
            raise NoCompatibleExecutionEnvironmentError(
                f"No compatible remote execution environment is available: {details}."
            )
        return min(candidates, key=lambda candidate: candidate[0])[1]

    def _incompatibility_reason(
        self,
        environment: EnvironmentSchedulingState,
        requirements: ComponentResourceRequirements,
    ) -> str | None:
        """Return a human-readable hard-constraint failure when incompatible."""
        if not environment.enabled:
            return "disabled"
        if (
            requirements.required_provider is not None
            and environment.provider is not requirements.required_provider
        ):
            return f"provider is not {requirements.required_provider.value!r}"
        if environment.health not in {
            EnvironmentHealth.READY,
            EnvironmentHealth.DEGRADED,
        }:
            return f"health is {environment.health.value}"
        if environment.maximum_workers <= 0:
            return "worker limit is zero"
        capabilities = environment.capabilities
        if capabilities is None:
            return "capabilities have not been probed"
        if (
            requirements.architecture
            and capabilities.architecture != requirements.architecture
        ):
            return (
                f"architecture {capabilities.architecture!r} does not satisfy "
                f"{requirements.architecture!r}"
            )
        if capabilities.schedulable_ram_bytes < requirements.minimum_ram_bytes:
            return (
                "insufficient RAM "
                f"({capabilities.schedulable_ram_bytes / 1024**3:.2f} GiB available, "
                f"{requirements.minimum_ram_bytes / 1024**3:.2f} GiB required)"
            )
        if not requirements.required_tags.issubset(environment.tags):
            return "required tags are missing"
        if requirements.gpu_required:
            if not capabilities.nvidia_container_runtime:
                return "NVIDIA container runtime is unavailable"
            if (
                capabilities.maximum_available_vram_bytes
                < requirements.minimum_vram_bytes
            ):
                return (
                    "insufficient GPU VRAM "
                    "("
                    f"{capabilities.maximum_available_vram_bytes / 1024**3:.2f} "
                    "GiB available, "
                    f"{requirements.minimum_vram_bytes / 1024**3:.2f} GiB required)"
                )
        return None

    def _preferred_rank(
        self,
        environment_id: str,
        preferred_environment_ids: tuple[str, ...],
    ) -> int:
        """Return a stable rank for explicit environment preferences."""
        if not preferred_environment_ids:
            return 0
        try:
            return preferred_environment_ids.index(environment_id)
        except ValueError:
            return len(preferred_environment_ids) + 1


def _non_negative_int(value: Any, field_name: str) -> int:
    """Return one non-negative integer or raise a validation error."""
    if isinstance(value, bool):
        raise ValueError(f"{field_name} must be a non-negative integer.")
    try:
        normalized = int(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field_name} must be a non-negative integer.") from exc
    if normalized < 0:
        raise ValueError(f"{field_name} must be a non-negative integer.")
    return normalized


def _optional_non_negative_int(value: Any, field_name: str) -> int | None:
    """Return an optional non-negative integer."""
    if value is None:
        return None
    return _non_negative_int(value, field_name)


def _optional_non_negative_float(value: Any, field_name: str) -> float | None:
    """Return an optional finite non-negative floating-point number."""
    if value is None:
        return None
    if isinstance(value, bool):
        raise ValueError(f"{field_name} must be a non-negative number.")
    try:
        normalized = float(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field_name} must be a non-negative number.") from exc
    if not math.isfinite(normalized) or normalized < 0:
        raise ValueError(f"{field_name} must be a non-negative number.")
    return normalized


def _optional_string(value: Any) -> str | None:
    """Return stripped optional text."""
    if value is None:
        return None
    normalized = str(value).strip()
    return normalized or None