"""Direct Vast worker launch settings, readiness, and watchdog state updates."""

from __future__ import annotations

import json
import logging
import math
import os
import shlex
import time
from dataclasses import dataclass
from pathlib import PurePosixPath
from typing import Any, Callable, Mapping

if __package__:
    from .remote_protocol import REMOTE_PROTOCOL_VERSION
    from .remote.vast_watchdog import VastWatchdogSnapshot
    from .vast_leases import VastLeaseRecord
    from .vast_models import VastInstanceLaunchSpec, VastResourceProfile
    from .vast_ssh import VastSshError, VastSshRunner, VastSshVolumeBackend
else:  # pragma: no cover - direct debugging imports.
    from remote_protocol import REMOTE_PROTOCOL_VERSION
    from remote.vast_watchdog import VastWatchdogSnapshot
    from vast_leases import VastLeaseRecord
    from vast_models import VastInstanceLaunchSpec, VastResourceProfile
    from vast_ssh import VastSshError, VastSshRunner, VastSshVolumeBackend

logger = logging.getLogger(__name__)

VAST_IMAGE_ENV = "COMFY_MODAL_VAST_IMAGE"
DEFAULT_VAST_REMOTE_REPO_ROOT = PurePosixPath("/opt/comfy-remote/repo")
DEFAULT_VAST_REMOTE_COMFYUI_ROOT = PurePosixPath("/opt/comfy-remote/ComfyUI")
DEFAULT_VAST_REMOTE_STORAGE_ROOT = PurePosixPath("/storage")
DEFAULT_VAST_WATCHDOG_STATE_PATH = PurePosixPath(
    "/storage/comfy-vast-watchdog.json"
)


@dataclass(frozen=True)
class VastRuntimeConfiguration:
    """Describe the digest-pinned worker image and direct runtime layout."""

    image: str
    runtime_fingerprint: str
    remote_repo_root: PurePosixPath = DEFAULT_VAST_REMOTE_REPO_ROOT
    remote_comfyui_root: PurePosixPath = DEFAULT_VAST_REMOTE_COMFYUI_ROOT
    remote_storage_root: PurePosixPath = DEFAULT_VAST_REMOTE_STORAGE_ROOT
    watchdog_state_path: PurePosixPath = DEFAULT_VAST_WATCHDOG_STATE_PATH
    startup_timeout_seconds: float = 900.0
    readiness_poll_seconds: float = 2.0

    def __post_init__(self) -> None:
        """Validate launch and fingerprint values."""
        if not self.image.strip() or any(
            character in self.image for character in ("\x00", "\n", "\r")
        ):
            raise ValueError("Vast runtime image must be a non-empty single-line value.")
        if len(self.runtime_fingerprint) != 64:
            raise ValueError("Vast runtime fingerprint must be a SHA-256 hex digest.")
        for field_name, path in (
            ("remote_repo_root", self.remote_repo_root),
            ("remote_comfyui_root", self.remote_comfyui_root),
            ("remote_storage_root", self.remote_storage_root),
            ("watchdog_state_path", self.watchdog_state_path),
        ):
            if not path.is_absolute() or path == PurePosixPath("/"):
                raise ValueError(f"{field_name} must be an absolute non-root path.")
        if self.startup_timeout_seconds <= 0 or self.readiness_poll_seconds <= 0:
            raise ValueError("Vast runtime readiness timing must be positive.")

    @classmethod
    def from_environment(
        cls,
        runtime_fingerprint: str,
        *,
        environment: Mapping[str, str] | None = None,
    ) -> "VastRuntimeConfiguration":
        """Resolve the required worker image without inventing an unpublished tag."""
        source = os.environ if environment is None else environment
        image = str(source.get(VAST_IMAGE_ENV) or "").strip()
        if not image:
            raise RuntimeError(
                f"Set {VAST_IMAGE_ENV} to a published ComfyUI-Modal Vast worker image "
                "before renting an instance."
            )
        return cls(image=image, runtime_fingerprint=runtime_fingerprint)

    def launch_spec(
        self,
        profile: VastResourceProfile,
        label: str,
    ) -> VastInstanceLaunchSpec:
        """Return one direct-SSH launch with worker and watchdog supervision."""
        environment = {
            "COMFY_MODAL_REMOTE_WORKER": "1",
            "COMFY_MODAL_LLM_EXECUTION_TARGET": "vast",
            "COMFYUI_ROOT": str(self.remote_comfyui_root),
            "COMFY_MODAL_COMFYUI_ROOT": str(self.remote_comfyui_root),
            "COMFY_MODAL_LOCAL_STORAGE_ROOT": str(self.remote_storage_root),
            "COMFY_MODAL_REMOTE_STORAGE_ROOT": str(self.remote_storage_root),
            "COMFY_MODAL_EXECUTION_MODE": "local",
            "PYTHONPATH": f"{self.remote_repo_root}:{self.remote_comfyui_root}",
        }
        exported_lines = "\n".join(
            f"{name}={value}" for name, value in sorted(environment.items())
        )
        supervisor_command = shlex.join(
            (
                "python",
                "-m",
                "remote.vast_supervisor",
                "start",
                "--watchdog-state",
                str(self.watchdog_state_path),
            )
        )
        onstart = "\n".join(
            (
                "set -e",
                "mkdir -p /run/comfy-remote /storage/logs",
                "cat >> /etc/environment <<'COMFY_VAST_ENV'",
                exported_lines,
                "COMFY_VAST_ENV",
                supervisor_command,
            )
        )
        return VastInstanceLaunchSpec(
            image=self.image,
            disk_gb=profile.allocated_disk_gb,
            label=label,
            onstart=onstart,
            environment=environment,
            rental_type=profile.rental_type,
        )


@dataclass
class VastRuntimeManager:
    """Ensure the direct worker process matches the selected runtime image."""

    runner: VastSshRunner
    configuration: VastRuntimeConfiguration
    clock: Callable[[], float] = time.time
    monotonic: Callable[[], float] = time.monotonic
    sleep: Callable[[float], None] = time.sleep

    def ensure_worker(self) -> dict[str, Any]:
        """Wait for the supervised worker socket and exact runtime fingerprint."""
        started_at = self.monotonic()
        deadline = started_at + self.configuration.startup_timeout_seconds
        last_error: str | None = None
        last_logged_error: str | None = None
        attempt = 0
        while self.monotonic() < deadline:
            attempt += 1
            try:
                info = self.runtime_info()
            except (VastSshError, ValueError) as exc:
                last_error = str(exc)
            else:
                actual_fingerprint = str(info.get("runtime_fingerprint") or "")
                if actual_fingerprint != self.configuration.runtime_fingerprint:
                    actual_protocol = int(info.get("protocol_version") or 0)
                    if actual_protocol != REMOTE_PROTOCOL_VERSION:
                        raise VastSshError(
                            "Vast worker image fingerprint mismatch: expected "
                            f"{self.configuration.runtime_fingerprint[:12]}, found "
                            f"{actual_fingerprint[:12] or 'empty'}, with incompatible "
                            f"protocol {actual_protocol}."
                        )
                    logger.warning(
                        "Reusing Vast worker with compatible protocol despite source "
                        "fingerprint drift expected=%s actual=%s protocol=%d.",
                        self.configuration.runtime_fingerprint[:12],
                        actual_fingerprint[:12] or "empty",
                        actual_protocol,
                    )
                if bool(info.get("worker_socket_ready")):
                    logger.info(
                        "Vast worker became ready after %d probe attempts (%.1fs).",
                        attempt,
                        self.monotonic() - started_at,
                    )
                    return info
                self._start_supervisor()
                last_error = "worker socket is not ready"
            if last_error != last_logged_error:
                logger.warning(
                    "Vast worker readiness probe attempt=%d failed: %s",
                    attempt,
                    last_error,
                )
                last_logged_error = last_error
            self.sleep(self.configuration.readiness_poll_seconds)
        raise TimeoutError(
            "Vast worker did not become ready within "
            f"{self.configuration.startup_timeout_seconds:.0f}s: "
            f"{last_error or 'no diagnostics'}."
        )

    def runtime_info(self) -> dict[str, Any]:
        """Return validated JSON from the direct worker image."""
        result = self.runner.run(
            ("python", "-m", "remote.ssh_worker", "runtime-info"),
            check=True,
        )
        try:
            payload = json.loads(result.stdout_text)
        except json.JSONDecodeError as exc:
            raise ValueError("Vast worker returned invalid runtime-info JSON.") from exc
        if not isinstance(payload, dict):
            raise ValueError("Vast worker runtime-info must be a JSON object.")
        return payload

    def restart_worker(self) -> None:
        """Restart only the managed worker while preserving watchdog and storage."""
        self.runner.run(
            (
                "python",
                "-m",
                "remote.vast_supervisor",
                "restart-worker",
                "--watchdog-state",
                str(self.configuration.watchdog_state_path),
            ),
            timeout_seconds=60.0,
        )
        self.ensure_worker()

    def update_watchdog(self, lease: VastLeaseRecord) -> None:
        """Atomically publish controller activity for the in-instance fail-safe."""
        if not math.isfinite(lease.idle_deadline_epoch):
            raise ValueError("Vast lease idle deadline must be finite.")
        snapshot = VastWatchdogSnapshot(
            instance_id=lease.instance_id,
            owner_label=lease.label,
            idle_deadline_epoch=lease.idle_deadline_epoch,
            active_invocations=lease.active_invocations,
            updated_at_epoch=self.clock(),
        )
        backend = VastSshVolumeBackend(
            self.runner,
            storage_root=self.configuration.remote_storage_root,
        )
        relative_path = self.configuration.watchdog_state_path.relative_to(
            self.configuration.remote_storage_root
        )
        backend.put_bytes(snapshot.to_json_bytes(), relative_path.as_posix())

    def _start_supervisor(self) -> None:
        """Idempotently start both direct managed processes."""
        self.runner.run(
            (
                "python",
                "-m",
                "remote.vast_supervisor",
                "start",
                "--watchdog-state",
                str(self.configuration.watchdog_state_path),
            ),
            timeout_seconds=60.0,
        )


__all__ = [
    "DEFAULT_VAST_REMOTE_COMFYUI_ROOT",
    "DEFAULT_VAST_REMOTE_REPO_ROOT",
    "DEFAULT_VAST_REMOTE_STORAGE_ROOT",
    "DEFAULT_VAST_WATCHDOG_STATE_PATH",
    "VAST_IMAGE_ENV",
    "VastRuntimeConfiguration",
    "VastRuntimeManager",
]