"""Build and publish a current worker image after Vast source drift."""

from __future__ import annotations

import logging
import os
import re
import shlex
import subprocess
import sys
import threading
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import IO

if __package__:
    from .vast_image_registry import (
        VastImageNotFoundError,
        VastImageRegistryError,
        published_runtime_fingerprint,
    )
else:  # pragma: no cover - direct debugging imports.
    from vast_image_registry import (
        VastImageNotFoundError,
        VastImageRegistryError,
        published_runtime_fingerprint,
    )

logger = logging.getLogger(__name__)

VAST_IMAGE_BUILD_COMMAND = (
    "uv",
    "run",
    "python",
    "scripts/build_vast_worker_image.py",
    "--push",
)
_IMAGE_RESULT_PREFIX = "COMFY_MODAL_VAST_IMAGE="
_SOURCE_FINGERPRINT_RESULT_PREFIX = "COMFY_MODAL_VAST_SOURCE_FINGERPRINT="
_DIGEST_REFERENCE_PATTERN = re.compile(r"^\S+@sha256:[0-9a-f]{64}$")
_FINGERPRINT_PATTERN = re.compile(r"^[0-9a-f]{64}$")
_DIGEST_IN_TEXT_PATTERN = re.compile(r"sha256:[0-9a-f]{64}")
_ANSI_ESCAPE_PATTERN = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
_MAXIMUM_STATUS_LINE_CHARACTERS = 300
_BUILD_LOCK = threading.Lock()
_BUILT_IMAGES_BY_FINGERPRINT: dict[str, str] = {}

VastImageBuildStatusCallback = Callable[[str], None]


class VastWorkerImageBuildError(RuntimeError):
    """Report that the automatic Vast worker image build could not finish."""


@dataclass(frozen=True)
class VastWorkerImageBuilder:
    """Run the repository image builder and stream bounded progress messages."""

    repo_root: Path
    comfyui_root: Path | None
    modal_gpu: str | None = None
    environment: Mapping[str, str] | None = None

    def ensure_published_image(
        self,
        image: str,
        expected_fingerprint: str,
        *,
        status_callback: VastImageBuildStatusCallback | None = None,
    ) -> str:
        """Return a current published image, rebuilding before rental if stale."""
        self._emit(status_callback, "Checking the published Vast worker image")
        try:
            actual_fingerprint = published_runtime_fingerprint(image)
        except VastImageNotFoundError:
            self._emit(
                status_callback,
                "Published Vast worker image is missing; building before requesting "
                "capacity",
            )
            return self.build_and_push(
                expected_fingerprint,
                status_callback=status_callback,
            )
        except (VastImageRegistryError, ValueError) as exc:
            raise VastWorkerImageBuildError(
                self._manual_build_message(
                    "Unable to inspect the configured Vast worker image before "
                    f"renting capacity: {exc} Ensure the image is publicly readable."
                )
            ) from exc
        if actual_fingerprint == expected_fingerprint:
            self._emit(status_callback, "Published Vast worker image is current")
            return image
        actual_summary = (
            actual_fingerprint[:12]
            if actual_fingerprint is not None
            else "missing label"
        )
        logger.warning(
            "Published Vast worker image is stale image=%s expected=%s actual=%s.",
            image,
            expected_fingerprint[:12],
            actual_summary,
        )
        self._emit(
            status_callback,
            "Published Vast worker image is stale; rebuilding before requesting "
            "capacity",
        )
        return self.build_and_push(
            expected_fingerprint,
            status_callback=status_callback,
        )

    def build_and_push(
        self,
        expected_fingerprint: str,
        *,
        status_callback: VastImageBuildStatusCallback | None = None,
    ) -> str:
        """Build once per fingerprint and return the published digest reference."""
        with _BUILD_LOCK:
            cached = _BUILT_IMAGES_BY_FINGERPRINT.get(expected_fingerprint)
            if cached is not None:
                self._emit(
                    status_callback,
                    "Using the worker image published by another workflow",
                )
                return cached
            image = self._run_build(
                expected_fingerprint=expected_fingerprint,
                status_callback=status_callback,
            )
            _BUILT_IMAGES_BY_FINGERPRINT[expected_fingerprint] = image
            return image

    def _run_build(
        self,
        *,
        expected_fingerprint: str,
        status_callback: VastImageBuildStatusCallback | None,
    ) -> str:
        """Execute the image builder and extract its immutable result reference."""
        command = self._automatic_build_command(expected_fingerprint)
        command_text = shlex.join(command)
        self._emit(status_callback, "Building the current Vast worker image")
        environment = dict(os.environ if self.environment is None else self.environment)
        environment.setdefault("BUILDKIT_PROGRESS", "plain")
        if self.comfyui_root is not None:
            environment["COMFYUI_ROOT"] = str(self.comfyui_root)
        if self.modal_gpu is not None:
            environment["COMFY_MODAL_GPU"] = self.modal_gpu
        try:
            process = subprocess.Popen(
                command,
                cwd=self.repo_root,
                env=environment,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                encoding="utf-8",
                errors="replace",
                bufsize=1,
            )
        except OSError as exc:
            raise VastWorkerImageBuildError(
                self._manual_build_message(
                    f"Unable to start {command[0]!r}: {exc}"
                )
            ) from exc
        image_reference, source_fingerprint, recent_output = self._consume_output(
            process.stdout,
            status_callback=status_callback,
        )
        return_code = process.wait()
        if (
            source_fingerprint is not None
            and source_fingerprint != expected_fingerprint
        ):
            raise VastWorkerImageBuildError(
                "Automatic Vast worker image publication stopped because local "
                "runtime source changed after ComfyUI started: expected "
                f"{expected_fingerprint[:12]}, found "
                f"{source_fingerprint[:12]}. Restart ComfyUI and retry the "
                "workflow; no manual image build is needed."
            )
        if return_code != 0:
            diagnostic = self._failure_diagnostic(recent_output)
            raise VastWorkerImageBuildError(
                self._manual_build_message(
                    f"{command_text} exited with status {return_code}: {diagnostic}"
                )
            )
        if image_reference is None:
            raise VastWorkerImageBuildError(
                self._manual_build_message(
                    "The build completed without returning an immutable image digest."
                )
            )
        if source_fingerprint is None:
            raise VastWorkerImageBuildError(
                self._manual_build_message(
                    "The build did not report its baked runtime fingerprint."
                )
            )
        self._emit(status_callback, "Published the current Vast worker image")
        return image_reference

    @staticmethod
    def _automatic_build_command(expected_fingerprint: str) -> tuple[str, ...]:
        """Run with the exact interpreter hosting the active ComfyUI process."""
        return (
            sys.executable,
            "scripts/build_vast_worker_image.py",
            "--push",
            "--expected-fingerprint",
            expected_fingerprint,
        )

    @staticmethod
    def _failure_diagnostic(recent_output: tuple[str, ...]) -> str:
        """Return the last substantive build line instead of traceback footer noise."""
        ignored_lines = {"<no Python frame>", "[stderr]", "]"}
        for line in reversed(recent_output):
            if line not in ignored_lines and not line.startswith("Current thread "):
                return line
        return recent_output[-1] if recent_output else "no build output"

    def _consume_output(
        self,
        output: IO[str] | None,
        *,
        status_callback: VastImageBuildStatusCallback | None,
    ) -> tuple[str | None, str | None, tuple[str, ...]]:
        """Forward output and return its digest, source identity, and diagnostics."""
        if output is None:
            return None, None, ()
        image_reference: str | None = None
        source_fingerprint: str | None = None
        recent_output: list[str] = []
        for raw_line in output:
            line = raw_line.strip()
            if not line:
                continue
            if line.startswith(_IMAGE_RESULT_PREFIX):
                candidate = line.removeprefix(_IMAGE_RESULT_PREFIX).strip()
                if _DIGEST_REFERENCE_PATTERN.fullmatch(candidate):
                    image_reference = candidate
                continue
            if line.startswith(_SOURCE_FINGERPRINT_RESULT_PREFIX):
                candidate = line.removeprefix(
                    _SOURCE_FINGERPRINT_RESULT_PREFIX
                ).strip()
                if _FINGERPRINT_PATTERN.fullmatch(candidate):
                    source_fingerprint = candidate
                    continue
            bounded = self._safe_status_line(line)
            recent_output.append(bounded)
            del recent_output[:-10]
            logger.info("Vast worker image build: %s", bounded)
            self._emit(status_callback, f"Vast image build: {bounded}")
        return image_reference, source_fingerprint, tuple(recent_output)

    @staticmethod
    def _safe_status_line(line: str) -> str:
        """Bound one progress line and remove terminal codes and exact digests."""
        without_ansi = _ANSI_ESCAPE_PATTERN.sub("", line)
        redacted = _DIGEST_IN_TEXT_PATTERN.sub("sha256:[redacted]", without_ansi)
        printable = "".join(
            character for character in redacted if character.isprintable()
        )
        return printable[:_MAXIMUM_STATUS_LINE_CHARACTERS]

    def _manual_build_message(self, cause: str) -> str:
        """Return an actionable error without leaking subprocess environment data."""
        recovery_command = shlex.join(VAST_IMAGE_BUILD_COMMAND)
        return (
            f"Automatic Vast worker image publication failed. {cause} Run `"
            + recovery_command
            + "` from "
            + str(self.repo_root)
            + ", set COMFY_MODAL_VAST_IMAGE to the printed digest, restart ComfyUI, "
            "and retry the workflow."
        )

    @staticmethod
    def _emit(
        status_callback: VastImageBuildStatusCallback | None,
        message: str,
    ) -> None:
        """Send one build update when workflow progress is available."""
        if status_callback is not None:
            status_callback(message)


__all__ = [
    "VAST_IMAGE_BUILD_COMMAND",
    "VastWorkerImageBuildError",
    "VastWorkerImageBuilder",
]