"""Inspect public OCI worker images without downloading their filesystem layers."""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass
from typing import Any, Mapping
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urlencode
from urllib.request import Request, urlopen

logger = logging.getLogger(__name__)

RUNTIME_FINGERPRINT_LABEL = "comfy.remote.runtime-fingerprint"
_MANIFEST_ACCEPT = ", ".join(
    (
        "application/vnd.oci.image.index.v1+json",
        "application/vnd.oci.image.manifest.v1+json",
        "application/vnd.docker.distribution.manifest.list.v2+json",
        "application/vnd.docker.distribution.manifest.v2+json",
    )
)
_BEARER_PARAMETER_PATTERN = re.compile(r'(\w+)="([^"]*)"')


class VastImageRegistryError(RuntimeError):
    """Report that published worker image metadata could not be inspected."""


class VastImageNotFoundError(VastImageRegistryError):
    """Report that a configured worker image is absent from its registry."""


@dataclass(frozen=True)
class OciImageReference:
    """Identify one registry repository and tag or digest reference."""

    registry: str
    repository: str
    reference: str


def parse_oci_image_reference(image: str) -> OciImageReference:
    """Parse a Docker-compatible image reference for registry HTTP requests."""
    normalized = image.strip()
    if not normalized:
        raise ValueError("Worker image reference must not be empty.")
    name, separator, digest = normalized.partition("@")
    if separator:
        reference = digest
    else:
        final_slash = name.rfind("/")
        final_colon = name.rfind(":")
        if final_colon > final_slash:
            reference = name[final_colon + 1 :]
            name = name[:final_colon]
        else:
            reference = "latest"
    components = name.split("/")
    first = components[0]
    explicit_registry = "." in first or ":" in first or first == "localhost"
    if explicit_registry:
        if len(components) < 2:
            raise ValueError(f"Worker image {image!r} is missing a repository name.")
        registry = first
        repository = "/".join(components[1:])
    else:
        registry = "registry-1.docker.io"
        repository = name if "/" in name else f"library/{name}"
    if registry in {"docker.io", "index.docker.io"}:
        registry = "registry-1.docker.io"
    if not repository or not reference:
        raise ValueError(f"Worker image reference {image!r} is malformed.")
    return OciImageReference(
        registry=registry,
        repository=repository,
        reference=reference,
    )


def published_runtime_fingerprint(
    image: str,
    *,
    timeout_seconds: float = 30.0,
) -> str | None:
    """Return the runtime label from a public image's linux/amd64 config."""
    reference = parse_oci_image_reference(image)
    token: str | None = None
    manifest, token = _registry_json(
        reference,
        resource=f"manifests/{quote(reference.reference, safe=':')}",
        accept=_MANIFEST_ACCEPT,
        token=token,
        timeout_seconds=timeout_seconds,
    )
    selected_digest = _linux_amd64_digest(manifest)
    if selected_digest is not None:
        manifest, token = _registry_json(
            reference,
            resource=f"manifests/{quote(selected_digest, safe=':')}",
            accept=_MANIFEST_ACCEPT,
            token=token,
            timeout_seconds=timeout_seconds,
        )
    config_digest = _config_digest(manifest)
    config, _ = _registry_json(
        reference,
        resource=f"blobs/{quote(config_digest, safe=':')}",
        accept="application/octet-stream",
        token=token,
        timeout_seconds=timeout_seconds,
    )
    labels = _image_labels(config)
    fingerprint = labels.get(RUNTIME_FINGERPRINT_LABEL)
    logger.info(
        "Inspected published Vast worker image image=%s fingerprint=%s.",
        image,
        fingerprint[:12] if isinstance(fingerprint, str) else "missing",
    )
    return fingerprint if isinstance(fingerprint, str) else None


def _registry_json(
    reference: OciImageReference,
    *,
    resource: str,
    accept: str,
    token: str | None,
    timeout_seconds: float,
) -> tuple[Mapping[str, Any], str | None]:
    """Fetch one JSON registry resource, obtaining an anonymous token if needed."""
    url = f"https://{reference.registry}/v2/{reference.repository}/{resource}"
    try:
        return _request_json(url, accept, token, timeout_seconds), token
    except HTTPError as error:
        if error.code != 401:
            if error.code == 404:
                raise VastImageNotFoundError(
                    f"Worker image was not found in {reference.registry}."
                ) from error
            raise VastImageRegistryError(
                f"Registry returned HTTP {error.code} for {reference.registry}."
            ) from error
        challenge = (
            error.headers.get("WWW-Authenticate", "")
            if error.headers is not None
            else ""
        )
    parameters = _bearer_parameters(challenge)
    realm = parameters.get("realm")
    if realm is None:
        raise VastImageRegistryError(
            f"Registry {reference.registry} requires unsupported authentication."
        )
    query = {
        key: value
        for key, value in (
            ("service", parameters.get("service")),
            (
                "scope",
                parameters.get("scope") or f"repository:{reference.repository}:pull",
            ),
        )
        if value is not None
    }
    query_separator = "&" if "?" in realm else "?"
    token_url = f"{realm}{query_separator}{urlencode(query)}" if query else realm
    try:
        token_payload = _request_json(
            token_url,
            "application/json",
            None,
            timeout_seconds,
        )
    except HTTPError as error:
        raise VastImageRegistryError(
            f"Registry token service returned HTTP {error.code}."
        ) from error
    resolved_token = token_payload.get("token") or token_payload.get("access_token")
    if not isinstance(resolved_token, str) or not resolved_token:
        raise VastImageRegistryError(
            f"Registry {reference.registry} did not return an anonymous pull token."
        )
    try:
        document = _request_json(url, accept, resolved_token, timeout_seconds)
    except HTTPError as error:
        if error.code == 404:
            raise VastImageNotFoundError(
                f"Worker image was not found in {reference.registry}."
            ) from error
        raise VastImageRegistryError(
            f"Registry rejected its pull token with HTTP {error.code}."
        ) from error
    return document, resolved_token


def _request_json(
    url: str,
    accept: str,
    token: str | None,
    timeout_seconds: float,
) -> Mapping[str, Any]:
    """Perform one bounded registry request and decode its JSON object."""
    headers = {"Accept": accept, "User-Agent": "ComfyUI-Modal/worker-preflight"}
    if token is not None:
        headers["Authorization"] = f"Bearer {token}"
    request = Request(url, headers=headers)
    try:
        with urlopen(request, timeout=timeout_seconds) as response:
            payload = response.read()
    except HTTPError:
        raise
    except (OSError, URLError) as error:
        raise VastImageRegistryError(
            f"Could not reach worker image registry: {error}"
        ) from error
    try:
        document = json.loads(payload)
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise VastImageRegistryError(
            "Worker image registry returned invalid JSON."
        ) from error
    if not isinstance(document, Mapping):
        raise VastImageRegistryError(
            "Worker image registry returned a non-object response."
        )
    return document


def _bearer_parameters(challenge: str) -> dict[str, str]:
    """Return quoted parameters from one standard Bearer challenge."""
    if not challenge.casefold().startswith("bearer "):
        return {}
    return dict(_BEARER_PARAMETER_PATTERN.findall(challenge))


def _linux_amd64_digest(manifest: Mapping[str, Any]) -> str | None:
    """Select the linux/amd64 manifest digest from an OCI image index."""
    entries = manifest.get("manifests")
    if entries is None:
        return None
    if not isinstance(entries, list):
        raise VastImageRegistryError(
            "Worker image index has invalid manifests metadata."
        )
    for entry in entries:
        if not isinstance(entry, Mapping):
            continue
        platform = entry.get("platform")
        if not isinstance(platform, Mapping):
            continue
        if platform.get("os") == "linux" and platform.get("architecture") == "amd64":
            digest = entry.get("digest")
            if isinstance(digest, str) and digest.startswith("sha256:"):
                return digest
    raise VastImageRegistryError("Worker image has no linux/amd64 manifest.")


def _config_digest(manifest: Mapping[str, Any]) -> str:
    """Return the image config digest from one platform manifest."""
    config = manifest.get("config")
    digest = config.get("digest") if isinstance(config, Mapping) else None
    if not isinstance(digest, str) or not digest.startswith("sha256:"):
        raise VastImageRegistryError(
            "Worker image manifest has no valid config digest."
        )
    return digest


def _image_labels(config: Mapping[str, Any]) -> Mapping[str, Any]:
    """Return image config labels or an empty mapping when none are present."""
    configuration = config.get("config")
    labels = configuration.get("Labels") if isinstance(configuration, Mapping) else None
    return labels if isinstance(labels, Mapping) else {}


__all__ = [
    "RUNTIME_FINGERPRINT_LABEL",
    "VastImageNotFoundError",
    "VastImageRegistryError",
    "parse_oci_image_reference",
    "published_runtime_fingerprint",
]