"""Local build pipeline for packaged Firecracker runtime bundles."""

from __future__ import annotations

import argparse
import base64
import hashlib
import io
import json
import os
import shutil
import subprocess
import tarfile
import urllib.error
import urllib.parse
import urllib.request
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, cast

from pyro_mcp.runtime import DEFAULT_PLATFORM
from pyro_mcp.vm_environments import get_environment

DEFAULT_RUNTIME_SOURCE_DIR = Path("runtime_sources")
DEFAULT_RUNTIME_BUILD_DIR = Path("build/runtime_bundle")
DEFAULT_RUNTIME_BUNDLE_DIR = Path("src/pyro_mcp/runtime_bundle")
DEFAULT_RUNTIME_MATERIALIZED_DIR = Path("build/runtime_sources")
DEFAULT_RUNTIME_OCI_LAYOUT_DIR = Path("build/oci_layouts")
DOWNLOAD_CHUNK_SIZE = 1024 * 1024
OCI_IMAGE_MANIFEST_MEDIA_TYPE = "application/vnd.oci.image.manifest.v1+json"
OCI_IMAGE_CONFIG_MEDIA_TYPE = "application/vnd.oci.image.config.v1+json"
OCI_IMAGE_LAYER_MEDIA_TYPE = "application/vnd.oci.image.layer.v1.tar"
OCI_IMAGE_INDEX_MEDIA_TYPE = "application/vnd.oci.image.index.v1+json"
OCI_LAYOUT_VERSION = "1.0.0"
VALIDATION_READ_LIMIT = 1024 * 1024
DEFAULT_OCI_USERNAME_ENV = "OCI_REGISTRY_USERNAME"
DEFAULT_OCI_PASSWORD_ENV = "OCI_REGISTRY_PASSWORD"
DEFAULT_DOCKERHUB_USERNAME_ENV = "DOCKERHUB_USERNAME"
DEFAULT_DOCKERHUB_TOKEN_ENV = "DOCKERHUB_TOKEN"
DEFAULT_OCI_REQUEST_TIMEOUT_ENV = "PYRO_OCI_REQUEST_TIMEOUT_SECONDS"
DEFAULT_OCI_UPLOAD_TIMEOUT_ENV = "PYRO_OCI_UPLOAD_TIMEOUT_SECONDS"
DEFAULT_OCI_UPLOAD_CHUNK_SIZE_ENV = "PYRO_OCI_UPLOAD_CHUNK_SIZE_BYTES"
DEFAULT_OCI_REQUEST_TIMEOUT_SECONDS = 90
DEFAULT_OCI_UPLOAD_TIMEOUT_SECONDS = 900
DEFAULT_OCI_UPLOAD_CHUNK_SIZE_BYTES = 8 * 1024 * 1024
_REGISTRY_BEARER_TOKENS: dict[tuple[str, str, str, str], str] = {}


@dataclass(frozen=True)
class RuntimeBuildLock:
    bundle_version: str
    platform: str
    component_versions: dict[str, str]
    capabilities: dict[str, bool]
    binaries: dict[str, str]
    guest: dict[str, dict[str, str]]
    profiles: dict[str, dict[str, str]]
    upstream: dict[str, Any]


@dataclass(frozen=True)
class RuntimeBuildPaths:
    source_root: Path
    source_platform_root: Path
    build_root: Path
    build_platform_root: Path
    bundle_dir: Path
    materialized_root: Path
    materialized_platform_root: Path
    platform: str


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as fp:
        for block in iter(lambda: fp.read(DOWNLOAD_CHUNK_SIZE), b""):
            digest.update(block)
    return digest.hexdigest()


def _cache_filename(url: str) -> str:
    digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:12]
    return f"{digest}-{Path(url).name}"


def _load_lock(paths: RuntimeBuildPaths) -> RuntimeBuildLock:
    lock_path = paths.source_platform_root / "runtime.lock.json"
    payload = json.loads(lock_path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise RuntimeError(f"invalid runtime lock file: {lock_path}")
    return RuntimeBuildLock(
        bundle_version=str(payload["bundle_version"]),
        platform=str(payload["platform"]),
        component_versions={
            str(key): str(value) for key, value in dict(payload["component_versions"]).items()
        },
        capabilities={
            str(key): bool(value) for key, value in dict(payload["capabilities"]).items()
        },
        binaries={str(key): str(value) for key, value in dict(payload["binaries"]).items()},
        guest={
            str(key): {str(k): str(v) for k, v in dict(value).items()}
            for key, value in dict(payload["guest"]).items()
        },
        profiles={
            str(key): {str(k): str(v) for k, v in dict(value).items()}
            for key, value in dict(payload["profiles"]).items()
        },
        upstream={str(key): value for key, value in dict(payload.get("upstream", {})).items()},
    )


def _copy_file(source: Path, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, dest)


def _copy_notice(paths: RuntimeBuildPaths) -> None:
    _copy_file(paths.source_root / "NOTICE", paths.build_root / "NOTICE")


def _resolved_source_path(paths: RuntimeBuildPaths, relative_path: str) -> Path:
    materialized = paths.materialized_platform_root / relative_path
    if materialized.exists():
        return materialized
    return paths.source_platform_root / relative_path


def _download(url: str, dest: Path) -> None:  # pragma: no cover - integration helper
    dest.parent.mkdir(parents=True, exist_ok=True)
    with urllib.request.urlopen(url) as response, dest.open("wb") as fp:  # noqa: S310
        while True:
            chunk = response.read(DOWNLOAD_CHUNK_SIZE)
            if chunk == b"":
                break
            fp.write(chunk)


def _run(command: list[str]) -> None:  # pragma: no cover - integration helper
    completed = subprocess.run(command, text=True, capture_output=True, check=False)
    if completed.returncode != 0:
        stderr = completed.stderr.strip() or completed.stdout.strip()
        raise RuntimeError(f"command {' '.join(command)!r} failed: {stderr}")


def _sha256_bytes(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def _path_contains_marker(
    path: Path,
    marker: bytes,
    *,
    read_limit: int = VALIDATION_READ_LIMIT,
) -> bool:
    bytes_remaining = read_limit
    overlap = max(len(marker) - 1, 0)
    previous_tail = b""
    with path.open("rb") as fp:
        while bytes_remaining > 0:
            chunk = fp.read(min(DOWNLOAD_CHUNK_SIZE, bytes_remaining))
            if chunk == b"":
                break
            payload = previous_tail + chunk
            if marker in payload:
                return True
            previous_tail = payload[-overlap:] if overlap > 0 else b""
            bytes_remaining -= len(chunk)
    return False


def _environment_slug(environment: str) -> str:
    return environment.replace(":", "_").replace("/", "_")


def _platform_to_oci_platform(platform: str) -> tuple[str, str]:
    os_name, separator, architecture = platform.partition("-")
    if separator == "" or os_name == "" or architecture == "":
        raise RuntimeError(f"unsupported runtime platform format: {platform}")
    architecture_aliases = {"x86_64": "amd64", "aarch64": "arm64"}
    return os_name, architecture_aliases.get(architecture, architecture)


def _blob_path(blobs_dir: Path, digest: str) -> Path:
    algorithm, separator, value = digest.partition(":")
    if separator == "" or value == "":
        raise RuntimeError(f"invalid OCI digest: {digest}")
    return blobs_dir / algorithm / value


def _write_blob_bytes(
    blobs_dir: Path,
    payload: bytes,
) -> tuple[str, int, Path]:
    digest = f"sha256:{_sha256_bytes(payload)}"
    blob_path = _blob_path(blobs_dir, digest)
    blob_path.parent.mkdir(parents=True, exist_ok=True)
    blob_path.write_bytes(payload)
    return digest, len(payload), blob_path


def _normalized_tar_info(name: str, size: int, *, mode: int = 0o644) -> tarfile.TarInfo:
    info = tarfile.TarInfo(name=name)
    info.size = size
    info.mode = mode
    info.mtime = 0
    info.uid = 0
    info.gid = 0
    info.uname = ""
    info.gname = ""
    return info


def _write_tar_blob_from_path(
    blobs_dir: Path,
    *,
    source_path: Path,
    arcname: str,
) -> tuple[str, int, Path]:
    temp_dir = blobs_dir / "sha256"
    temp_dir.mkdir(parents=True, exist_ok=True)
    temp_path = temp_dir / f".tmp-{uuid.uuid4().hex}.tar"
    try:
        with temp_path.open("wb") as fp:
            with tarfile.open(fileobj=fp, mode="w") as archive:
                file_size = source_path.stat().st_size
                info = _normalized_tar_info(arcname, file_size)
                with source_path.open("rb") as source_fp:
                    archive.addfile(info, source_fp)
        digest = f"sha256:{_sha256(temp_path)}"
        blob_path = _blob_path(blobs_dir, digest)
        blob_path.parent.mkdir(parents=True, exist_ok=True)
        size = temp_path.stat().st_size
        temp_path.replace(blob_path)
        return digest, size, blob_path
    except Exception:
        temp_path.unlink(missing_ok=True)
        raise


def _write_tar_blob_from_bytes(
    blobs_dir: Path,
    *,
    payload: bytes,
    arcname: str,
) -> tuple[str, int, Path]:
    temp_dir = blobs_dir / "sha256"
    temp_dir.mkdir(parents=True, exist_ok=True)
    temp_path = temp_dir / f".tmp-{uuid.uuid4().hex}.tar"
    try:
        with temp_path.open("wb") as fp:
            with tarfile.open(fileobj=fp, mode="w") as archive:
                info = _normalized_tar_info(arcname, len(payload))
                archive.addfile(info, io.BytesIO(payload))
        digest = f"sha256:{_sha256(temp_path)}"
        blob_path = _blob_path(blobs_dir, digest)
        blob_path.parent.mkdir(parents=True, exist_ok=True)
        size = temp_path.stat().st_size
        temp_path.replace(blob_path)
        return digest, size, blob_path
    except Exception:
        temp_path.unlink(missing_ok=True)
        raise


def _normalize_headers(headers: Any) -> dict[str, str]:
    return {str(key).lower(): str(value) for key, value in headers.items()}


def _basic_auth_header(username: str | None, password: str | None) -> str | None:
    if username is None or password is None:
        return None
    token = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode("ascii")
    return f"Basic {token}"


def _parse_authenticate_parameters(raw: str) -> dict[str, str]:
    params: dict[str, str] = {}
    segments: list[str] = []
    current: list[str] = []
    in_quotes = False
    for char in raw:
        if char == '"':
            in_quotes = not in_quotes
        if char == "," and not in_quotes:
            segment = "".join(current).strip()
            if segment:
                segments.append(segment)
            current = []
            continue
        current.append(char)
    tail = "".join(current).strip()
    if tail:
        segments.append(tail)

    for segment in segments:
        if "=" not in segment:
            continue
        key, value = segment.split("=", 1)
        params[key.strip()] = value.strip().strip('"')
    return params


def _registry_url(registry: str, repository: str, suffix: str) -> str:
    return f"https://{registry}/v2/{repository}/{suffix}"


def _registry_credentials(
    *,
    username_env: str = DEFAULT_OCI_USERNAME_ENV,
    password_env: str = DEFAULT_OCI_PASSWORD_ENV,
) -> tuple[str | None, str | None]:
    username = os.environ.get(username_env) or os.environ.get(DEFAULT_DOCKERHUB_USERNAME_ENV)
    password = os.environ.get(password_env) or os.environ.get(DEFAULT_DOCKERHUB_TOKEN_ENV)
    return (
        username if isinstance(username, str) and username != "" else None,
        password if isinstance(password, str) and password != "" else None,
    )


def _rewind_body(data: object) -> None:
    if hasattr(data, "seek"):
        data.seek(0)


def _read_response_body(response: Any) -> bytes:
    if hasattr(response, "read"):
        return bytes(response.read())
    return b""


def _positive_int_from_env(name: str, default: int) -> int:
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    try:
        value = int(raw)
    except ValueError as exc:
        raise RuntimeError(f"{name} must be an integer; got {raw!r}") from exc
    if value <= 0:
        raise RuntimeError(f"{name} must be greater than zero; got {value}")
    return value


def _oci_request_timeout_seconds() -> int:
    return _positive_int_from_env(
        DEFAULT_OCI_REQUEST_TIMEOUT_ENV,
        DEFAULT_OCI_REQUEST_TIMEOUT_SECONDS,
    )


def _oci_upload_timeout_seconds() -> int:
    return _positive_int_from_env(
        DEFAULT_OCI_UPLOAD_TIMEOUT_ENV,
        DEFAULT_OCI_UPLOAD_TIMEOUT_SECONDS,
    )


def _oci_upload_chunk_size_bytes() -> int:
    return _positive_int_from_env(
        DEFAULT_OCI_UPLOAD_CHUNK_SIZE_ENV,
        DEFAULT_OCI_UPLOAD_CHUNK_SIZE_BYTES,
    )


def _registry_token_cache_key(
    *,
    url: str,
    repository: str,
    scope_actions: str,
    username: str | None,
) -> tuple[str, str, str, str]:
    return (
        urllib.parse.urlsplit(url).netloc,
        repository,
        scope_actions,
        username or "",
    )


def _fetch_registry_token(
    authenticate: str,
    *,
    repository: str,
    scope_actions: str,
    username: str | None,
    password: str | None,
) -> str:
    if not authenticate.startswith("Bearer "):
        raise RuntimeError("unsupported OCI registry authentication scheme")
    params = _parse_authenticate_parameters(authenticate[len("Bearer ") :])
    realm = params.get("realm")
    if realm is None:
        raise RuntimeError("OCI auth challenge did not include a token realm")
    query = {
        "service": params.get("service", ""),
        "scope": params.get("scope", f"repository:{repository}:{scope_actions}"),
    }
    headers: dict[str, str] = {}
    basic_auth = _basic_auth_header(username, password)
    if basic_auth is not None:
        headers["Authorization"] = basic_auth
    request = urllib.request.Request(
        f"{realm}?{urllib.parse.urlencode(query)}",
        headers=headers,
        method="GET",
    )
    with urllib.request.urlopen(request, timeout=_oci_request_timeout_seconds()) as response:  # noqa: S310
        payload = json.loads(_read_response_body(response).decode("utf-8"))
    if not isinstance(payload, dict):
        raise RuntimeError("OCI auth token response was not a JSON object")
    raw_token = payload.get("token") or payload.get("access_token")
    if not isinstance(raw_token, str) or raw_token == "":
        raise RuntimeError("OCI auth token response did not include a bearer token")
    return raw_token


def _request_registry(
    url: str,
    *,
    method: str,
    repository: str,
    scope_actions: str,
    headers: dict[str, str] | None = None,
    data: object = None,
    allow_statuses: tuple[int, ...] = (),
    username: str | None = None,
    password: str | None = None,
    timeout_seconds: int | None = None,
) -> tuple[int, bytes, dict[str, str]]:
    request_headers = dict(headers or {})
    request_timeout = timeout_seconds or _oci_request_timeout_seconds()
    cache_key = _registry_token_cache_key(
        url=url,
        repository=repository,
        scope_actions=scope_actions,
        username=username,
    )
    cached_token = _REGISTRY_BEARER_TOKENS.get(cache_key)
    if cached_token is not None and "Authorization" not in request_headers:
        request_headers["Authorization"] = f"Bearer {cached_token}"
    request = urllib.request.Request(
        url,
        data=cast(Any, data),
        headers=request_headers,
        method=method,
    )
    try:
        with urllib.request.urlopen(request, timeout=request_timeout) as response:  # noqa: S310
            return (
                response.status,
                _read_response_body(response),
                _normalize_headers(response.headers),
            )
    except urllib.error.HTTPError as exc:
        if exc.code == 401:
            authenticate = exc.headers.get("WWW-Authenticate")
            if authenticate is None:
                raise RuntimeError("OCI registry denied access without an auth challenge") from exc
            token = _fetch_registry_token(
                authenticate,
                repository=repository,
                scope_actions=scope_actions,
                username=username,
                password=password,
            )
            _REGISTRY_BEARER_TOKENS[cache_key] = token
            authenticated_headers = {**request_headers, "Authorization": f"Bearer {token}"}
            _rewind_body(data)
            retry = urllib.request.Request(
                url,
                data=cast(Any, data),
                headers=authenticated_headers,
                method=method,
            )
            try:
                with urllib.request.urlopen(retry, timeout=request_timeout) as response:  # noqa: S310
                    return (
                        response.status,
                        _read_response_body(response),
                        _normalize_headers(response.headers),
                    )
            except urllib.error.HTTPError as retry_exc:
                if retry_exc.code in allow_statuses:
                    return (
                        retry_exc.code,
                        _read_response_body(retry_exc),
                        _normalize_headers(retry_exc.headers),
                    )
                raise RuntimeError(f"registry request failed for {url}: {retry_exc}") from retry_exc
            except urllib.error.URLError as retry_exc:
                raise RuntimeError(f"registry request failed for {url}: {retry_exc}") from retry_exc
        if exc.code in allow_statuses:
            return exc.code, _read_response_body(exc), _normalize_headers(exc.headers)
        raise RuntimeError(f"registry request failed for {url}: {exc}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"registry request failed for {url}: {exc}") from exc


def _load_oci_layout_manifest(
    layout_dir: Path,
) -> tuple[dict[str, Any], bytes, dict[str, Any], str]:
    index_path = layout_dir / "index.json"
    if not index_path.exists():
        raise RuntimeError(f"OCI layout index not found: {index_path}")
    index_payload = json.loads(index_path.read_text(encoding="utf-8"))
    manifests = index_payload.get("manifests")
    if not isinstance(manifests, list) or len(manifests) != 1:
        raise RuntimeError("OCI layout must contain exactly one manifest descriptor")
    descriptor = manifests[0]
    if not isinstance(descriptor, dict):
        raise RuntimeError("OCI layout manifest descriptor is malformed")
    raw_digest = descriptor.get("digest")
    if not isinstance(raw_digest, str):
        raise RuntimeError("OCI layout manifest descriptor is missing a digest")
    manifest_path = _blob_path(layout_dir / "blobs", raw_digest)
    manifest_bytes = manifest_path.read_bytes()
    manifest = json.loads(manifest_bytes.decode("utf-8"))
    if not isinstance(manifest, dict):
        raise RuntimeError("OCI layout manifest payload is malformed")
    media_type = descriptor.get("mediaType")
    if not isinstance(media_type, str) or media_type == "":
        media_type = str(manifest.get("mediaType") or OCI_IMAGE_MANIFEST_MEDIA_TYPE)
    return descriptor, manifest_bytes, manifest, media_type


def _blob_exists(
    *,
    registry: str,
    repository: str,
    digest: str,
    username: str | None,
    password: str | None,
) -> bool:
    status, _, _ = _request_registry(
        _registry_url(registry, repository, f"blobs/{digest}"),
        method="HEAD",
        repository=repository,
        scope_actions="pull,push",
        allow_statuses=(404,),
        username=username,
        password=password,
    )
    return status == 200


def _upload_blob(
    *,
    registry: str,
    repository: str,
    digest: str,
    blob_path: Path,
    username: str | None,
    password: str | None,
) -> dict[str, Any]:
    if _blob_exists(
        registry=registry,
        repository=repository,
        digest=digest,
        username=username,
        password=password,
    ):
        return {"digest": digest, "size": blob_path.stat().st_size, "uploaded": False}

    status, _, headers = _request_registry(
        _registry_url(registry, repository, "blobs/uploads/"),
        method="POST",
        repository=repository,
        scope_actions="pull,push",
        headers={"Content-Length": "0"},
        allow_statuses=(202,),
        username=username,
        password=password,
    )
    if status != 202:
        raise RuntimeError(f"unexpected registry status when starting blob upload: {status}")
    location = headers.get("location")
    if location is None:
        raise RuntimeError("registry did not return a blob upload location")
    upload_url = urllib.parse.urljoin(f"https://{registry}", location)
    size = blob_path.stat().st_size
    upload_timeout_seconds = _oci_upload_timeout_seconds()
    upload_chunk_size = _oci_upload_chunk_size_bytes()
    with blob_path.open("rb") as blob_fp:
        uploaded_bytes = 0
        chunk_index = 0
        while True:
            chunk = blob_fp.read(upload_chunk_size)
            if chunk == b"":
                break
            try:
                status, _, headers = _request_registry(
                    upload_url,
                    method="PATCH",
                    repository=repository,
                    scope_actions="pull,push",
                    headers={
                        "Content-Length": str(len(chunk)),
                        "Content-Type": "application/octet-stream",
                    },
                    data=chunk,
                    allow_statuses=(202,),
                    username=username,
                    password=password,
                    timeout_seconds=upload_timeout_seconds,
                )
            except RuntimeError as exc:
                raise RuntimeError(
                    "registry blob upload patch failed for "
                    f"{digest} at byte offset {uploaded_bytes} (chunk {chunk_index}): {exc}"
                ) from exc
            if status != 202:
                raise RuntimeError(
                    f"unexpected registry status when uploading blob chunk {chunk_index}: {status}"
                )
            location = headers.get("location")
            if location is None:
                raise RuntimeError(
                    "registry did not return a blob upload location after chunk "
                    f"{chunk_index}"
                )
            upload_url = urllib.parse.urljoin(f"https://{registry}", location)
            uploaded_bytes += len(chunk)
            chunk_index += 1

    separator = "&" if "?" in upload_url else "?"
    finalize_url = f"{upload_url}{separator}{urllib.parse.urlencode({'digest': digest})}"
    try:
        status, _, _ = _request_registry(
            finalize_url,
            method="PUT",
            repository=repository,
            scope_actions="pull,push",
            headers={
                "Content-Length": "0",
                "Content-Type": "application/octet-stream",
            },
            data=b"",
            allow_statuses=(201,),
            username=username,
            password=password,
            timeout_seconds=upload_timeout_seconds,
        )
    except RuntimeError as exc:
        raise RuntimeError(f"registry blob upload finalize failed for {digest}: {exc}") from exc
    if status != 201:
        raise RuntimeError(f"unexpected registry status when uploading blob: {status}")
    return {"digest": digest, "size": size, "uploaded": True}


def publish_environment_oci_layout(
    *,
    environment: str,
    layout_root: Path,
    registry: str | None = None,
    repository: str | None = None,
    reference: str | None = None,
    username_env: str = DEFAULT_OCI_USERNAME_ENV,
    password_env: str = DEFAULT_OCI_PASSWORD_ENV,
) -> dict[str, Any]:
    spec = get_environment(environment)
    resolved_registry = registry or spec.oci_registry
    resolved_repository = repository or spec.oci_repository
    resolved_reference = reference or spec.oci_reference or spec.version
    if resolved_registry is None or resolved_repository is None:
        raise RuntimeError(f"environment {environment!r} does not define an OCI registry target")

    layout_dir = layout_root / _environment_slug(environment)
    descriptor, manifest_bytes, manifest, media_type = _load_oci_layout_manifest(layout_dir)
    username, password = _registry_credentials(
        username_env=username_env,
        password_env=password_env,
    )
    if username is None or password is None:
        raise RuntimeError(
            "OCI registry credentials are not configured; set "
            f"{username_env}/{password_env} or "
            f"{DEFAULT_DOCKERHUB_USERNAME_ENV}/{DEFAULT_DOCKERHUB_TOKEN_ENV}"
        )

    config = manifest.get("config")
    layers = manifest.get("layers")
    if not isinstance(config, dict):
        raise RuntimeError("OCI layout manifest is missing a config descriptor")
    if not isinstance(layers, list):
        raise RuntimeError("OCI layout manifest is missing layers")

    uploaded_blobs: list[dict[str, Any]] = []
    descriptors = [config, *layers]
    for entry in descriptors:
        if not isinstance(entry, dict):
            raise RuntimeError("OCI layout descriptor is malformed")
        raw_digest = entry.get("digest")
        if not isinstance(raw_digest, str):
            raise RuntimeError("OCI layout descriptor is missing a digest")
        blob_path = _blob_path(layout_dir / "blobs", raw_digest)
        uploaded_blobs.append(
            _upload_blob(
                registry=resolved_registry,
                repository=resolved_repository,
                digest=raw_digest,
                blob_path=blob_path,
                username=username,
                password=password,
            )
        )

    status, _, headers = _request_registry(
        _registry_url(resolved_registry, resolved_repository, f"manifests/{resolved_reference}"),
        method="PUT",
        repository=resolved_repository,
        scope_actions="pull,push",
        headers={
            "Content-Length": str(len(manifest_bytes)),
            "Content-Type": media_type,
        },
        data=manifest_bytes,
        allow_statuses=(201,),
        username=username,
        password=password,
    )
    if status != 201:
        raise RuntimeError(f"unexpected registry status when publishing manifest: {status}")
    return {
        "environment": spec.name,
        "layout_dir": str(layout_dir),
        "registry": resolved_registry,
        "repository": resolved_repository,
        "reference": resolved_reference,
        "manifest_digest": headers.get("docker-content-digest", str(descriptor["digest"])),
        "uploaded_blobs": uploaded_blobs,
    }


def validate_sources(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> None:
    firecracker_source = _resolved_source_path(paths, lock.binaries["firecracker"])
    jailer_source = _resolved_source_path(paths, lock.binaries["jailer"])
    has_shim_binaries = (
        _path_contains_marker(firecracker_source, b"bundled firecracker shim")
        or _path_contains_marker(jailer_source, b"bundled jailer shim")
    )

    has_placeholder_profiles = False
    for profile in lock.profiles.values():
        for kind in ("kernel", "rootfs"):
            source = _resolved_source_path(paths, profile[kind])
            if _path_contains_marker(source, b"placeholder-"):
                has_placeholder_profiles = True
                break
        if has_placeholder_profiles:
            break

    if any(lock.capabilities.values()) and (has_shim_binaries or has_placeholder_profiles):
        raise RuntimeError(
            "runtime lock advertises guest-capable features while source artifacts are still "
            "shim/placeholder inputs"
        )


def materialize_binaries(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> None:
    release = lock.upstream.get("firecracker_release")
    if not isinstance(release, dict):
        raise RuntimeError("runtime lock is missing upstream.firecracker_release configuration")
    archive_url = release.get("archive_url")
    archive_sha256 = release.get("archive_sha256")
    firecracker_member = release.get("firecracker_member")
    jailer_member = release.get("jailer_member")
    if not all(
        isinstance(value, str)
        for value in (
            archive_url,
            archive_sha256,
            firecracker_member,
            jailer_member,
        )
    ):
        raise RuntimeError("upstream.firecracker_release is incomplete")
    archive_url = str(archive_url)
    archive_sha256 = str(archive_sha256)
    firecracker_member = str(firecracker_member)
    jailer_member = str(jailer_member)

    cache_dir = paths.materialized_root / "_downloads"
    archive_path = cache_dir / _cache_filename(archive_url)
    if not archive_path.exists():
        _download(archive_url, archive_path)
    actual_archive_sha256 = _sha256(archive_path)
    if actual_archive_sha256 != archive_sha256:
        raise RuntimeError(
            "firecracker release archive checksum mismatch: expected "
            f"{archive_sha256}, got {actual_archive_sha256}"
        )

    targets = {
        firecracker_member: paths.materialized_platform_root / lock.binaries["firecracker"],
        jailer_member: paths.materialized_platform_root / lock.binaries["jailer"],
    }
    with tarfile.open(archive_path, "r:gz") as archive:
        members = {member.name: member for member in archive.getmembers()}
        for member_name, dest in targets.items():
            member = members.get(member_name)
            if member is None:
                raise RuntimeError(f"release archive is missing {member_name}")
            extracted = archive.extractfile(member)
            if extracted is None:
                raise RuntimeError(f"unable to extract {member_name}")
            dest.parent.mkdir(parents=True, exist_ok=True)
            dest.write_bytes(extracted.read())
            dest.chmod(dest.stat().st_mode | 0o111)


def materialize_kernel(
    paths: RuntimeBuildPaths, lock: RuntimeBuildLock
) -> None:  # pragma: no cover - integration helper
    kernel_build = lock.upstream.get("kernel_build")
    if not isinstance(kernel_build, dict):
        raise RuntimeError("runtime lock is missing upstream.kernel_build configuration")
    script = kernel_build.get("script")
    linux_version = kernel_build.get("linux_version")
    source_url = kernel_build.get("source_url")
    config_url = kernel_build.get("config_url")
    builder_image = kernel_build.get("builder_image")
    if not all(
        isinstance(value, str)
        for value in (
            script,
            linux_version,
            source_url,
            config_url,
            builder_image,
        )
    ):
        raise RuntimeError("upstream.kernel_build is incomplete")
    script = str(script)
    linux_version = str(linux_version)
    source_url = str(source_url)
    config_url = str(config_url)
    builder_image = str(builder_image)

    script_path = paths.source_platform_root / script
    if not script_path.exists():
        raise RuntimeError(f"kernel build script not found: {script_path}")

    shared_output = paths.materialized_platform_root / "profiles/_shared/vmlinux"
    shared_output.parent.mkdir(parents=True, exist_ok=True)
    workdir = paths.materialized_root / "_kernel_work"
    workdir.mkdir(parents=True, exist_ok=True)
    _run(
        [
            str(script_path),
            "--builder-image",
            builder_image,
            "--linux-version",
            linux_version,
            "--source-url",
            source_url,
            "--config-url",
            config_url,
            "--workdir",
            str(workdir),
            "--output",
            str(shared_output),
        ]
    )

    for profile in lock.profiles.values():
        dest = paths.materialized_platform_root / profile["kernel"]
        _copy_file(shared_output, dest)


def materialize_rootfs(
    paths: RuntimeBuildPaths, lock: RuntimeBuildLock
) -> None:  # pragma: no cover - integration helper
    rootfs_build = lock.upstream.get("rootfs_build")
    if not isinstance(rootfs_build, dict):
        raise RuntimeError("runtime lock is missing upstream.rootfs_build configuration")
    script = rootfs_build.get("script")
    builder_image = rootfs_build.get("builder_image")
    debian_release = rootfs_build.get("debian_release")
    debian_snapshot = rootfs_build.get("debian_snapshot")
    package_files = rootfs_build.get("package_files")
    guest_init = rootfs_build.get("guest_init")
    agent_service = rootfs_build.get("agent_service")
    if not isinstance(package_files, dict):
        raise RuntimeError("upstream.rootfs_build.package_files must be a mapping")
    if not all(
        isinstance(value, str)
        for value in (
            script,
            builder_image,
            debian_release,
            debian_snapshot,
            guest_init,
            agent_service,
        )
    ):
        raise RuntimeError("upstream.rootfs_build is incomplete")
    script = str(script)
    builder_image = str(builder_image)
    debian_release = str(debian_release)
    debian_snapshot = str(debian_snapshot)
    guest_init = str(guest_init)
    agent_service = str(agent_service)

    script_path = paths.source_platform_root / script
    guest_agent_path = paths.source_platform_root / lock.guest["agent"]["path"]
    guest_init_path = paths.source_platform_root / guest_init
    service_path = paths.source_platform_root / agent_service
    if not script_path.exists():
        raise RuntimeError(f"rootfs build script not found: {script_path}")

    workdir = paths.materialized_root / "_rootfs_work"
    workdir.mkdir(parents=True, exist_ok=True)
    for profile_name, profile in lock.profiles.items():
        raw_packages_path = package_files.get(profile_name)
        if not isinstance(raw_packages_path, str):
            raise RuntimeError(f"missing package file for profile {profile_name!r}")
        packages_path = paths.source_platform_root / raw_packages_path
        output_path = paths.materialized_platform_root / profile["rootfs"]
        output_path.parent.mkdir(parents=True, exist_ok=True)
        profile_workdir = workdir / f"{profile_name}-{uuid.uuid4().hex[:8]}"
        _run(
            [
                str(script_path),
                "--builder-image",
                builder_image,
                "--debian-release",
                debian_release,
                "--debian-snapshot",
                debian_snapshot,
                "--packages-file",
                str(packages_path),
                "--guest-agent",
                str(guest_agent_path),
                "--guest-init",
                str(guest_init_path),
                "--agent-service",
                str(service_path),
                "--workdir",
                str(profile_workdir),
                "--output",
                str(output_path),
            ]
        )


def stage_binaries(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> None:
    for name, relative_path in lock.binaries.items():
        del name
        source = _resolved_source_path(paths, relative_path)
        dest = paths.build_platform_root / relative_path
        _copy_file(source, dest)
        dest.chmod(dest.stat().st_mode | 0o111)


def stage_kernel(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> None:
    for profile in lock.profiles.values():
        source = _resolved_source_path(paths, profile["kernel"])
        dest = paths.build_platform_root / profile["kernel"]
        _copy_file(source, dest)


def stage_rootfs(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> None:
    for profile in lock.profiles.values():
        source = _resolved_source_path(paths, profile["rootfs"])
        dest = paths.build_platform_root / profile["rootfs"]
        _copy_file(source, dest)


def stage_agent(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> None:
    for artifact in lock.guest.values():
        source = _resolved_source_path(paths, artifact["path"])
        dest = paths.build_platform_root / artifact["path"]
        _copy_file(source, dest)
        dest.chmod(dest.stat().st_mode | 0o111)


def export_environment_oci_layout(
    paths: RuntimeBuildPaths,
    *,
    environment: str,
    output_dir: Path,
    reference: str | None = None,
) -> dict[str, Any]:
    lock = _load_lock(paths)
    validate_sources(paths, lock)

    spec = get_environment(environment)
    try:
        profile = lock.profiles[spec.source_profile]
    except KeyError as exc:
        raise RuntimeError(
            f"runtime lock does not define source profile {spec.source_profile!r} "
            f"for environment {environment!r}"
        ) from exc

    kernel_path = _resolved_source_path(paths, profile["kernel"])
    rootfs_path = _resolved_source_path(paths, profile["rootfs"])
    if not kernel_path.exists() or not rootfs_path.exists():
        raise RuntimeError(
            f"missing artifacts for environment {environment!r}; expected "
            f"{kernel_path} and {rootfs_path}"
        )

    layout_dir = output_dir / _environment_slug(environment)
    if layout_dir.exists():
        shutil.rmtree(layout_dir)
    blobs_dir = layout_dir / "blobs"
    blob_sha_dir = blobs_dir / "sha256"
    blob_sha_dir.mkdir(parents=True, exist_ok=True)

    metadata_payload = {
        "environment": spec.name,
        "version": spec.version,
        "platform": spec.platform,
        "distribution": spec.distribution,
        "distribution_version": spec.distribution_version,
        "description": spec.description,
        "default_packages": list(spec.default_packages),
        "source_profile": spec.source_profile,
        "bundle_version": lock.bundle_version,
        "component_versions": lock.component_versions,
        "capabilities": lock.capabilities,
    }
    metadata_bytes = json.dumps(metadata_payload, indent=2, sort_keys=True).encode("utf-8") + b"\n"

    kernel_digest, kernel_size, _ = _write_tar_blob_from_path(
        blobs_dir,
        source_path=kernel_path,
        arcname="vmlinux",
    )
    rootfs_digest, rootfs_size, _ = _write_tar_blob_from_path(
        blobs_dir,
        source_path=rootfs_path,
        arcname="rootfs.ext4",
    )
    metadata_digest, metadata_size, _ = _write_tar_blob_from_bytes(
        blobs_dir,
        payload=metadata_bytes,
        arcname="environment.json",
    )

    created_at = datetime.now(UTC).isoformat().replace("+00:00", "Z")
    ref_name = reference or spec.version
    os_name, architecture = _platform_to_oci_platform(spec.platform)
    labels = {
        "io.pyro.environment": spec.name,
        "io.pyro.environment.version": spec.version,
        "io.pyro.source_profile": spec.source_profile,
        "org.opencontainers.image.title": spec.name,
        "org.opencontainers.image.version": spec.version,
    }
    config_payload = {
        "created": created_at,
        "architecture": architecture,
        "os": os_name,
        "config": {"Labels": labels},
        "rootfs": {
            "type": "layers",
            "diff_ids": [kernel_digest, rootfs_digest, metadata_digest],
        },
        "history": [
            {"created": created_at, "created_by": "pyro runtime_build export-environment-oci"}
        ],
    }
    config_bytes = json.dumps(config_payload, indent=2, sort_keys=True).encode("utf-8") + b"\n"
    config_digest, config_size, _ = _write_blob_bytes(blobs_dir, config_bytes)

    manifest_payload = {
        "schemaVersion": 2,
        "mediaType": OCI_IMAGE_MANIFEST_MEDIA_TYPE,
        "config": {
            "mediaType": OCI_IMAGE_CONFIG_MEDIA_TYPE,
            "digest": config_digest,
            "size": config_size,
        },
        "layers": [
            {
                "mediaType": OCI_IMAGE_LAYER_MEDIA_TYPE,
                "digest": kernel_digest,
                "size": kernel_size,
                "annotations": {"org.opencontainers.image.title": "vmlinux"},
            },
            {
                "mediaType": OCI_IMAGE_LAYER_MEDIA_TYPE,
                "digest": rootfs_digest,
                "size": rootfs_size,
                "annotations": {"org.opencontainers.image.title": "rootfs.ext4"},
            },
            {
                "mediaType": OCI_IMAGE_LAYER_MEDIA_TYPE,
                "digest": metadata_digest,
                "size": metadata_size,
                "annotations": {"org.opencontainers.image.title": "environment.json"},
            },
        ],
        "annotations": labels,
    }
    manifest_bytes = json.dumps(manifest_payload, indent=2, sort_keys=True).encode("utf-8") + b"\n"
    manifest_digest, manifest_size, _ = _write_blob_bytes(blobs_dir, manifest_bytes)

    index_payload = {
        "schemaVersion": 2,
        "mediaType": OCI_IMAGE_INDEX_MEDIA_TYPE,
        "manifests": [
            {
                "mediaType": OCI_IMAGE_MANIFEST_MEDIA_TYPE,
                "digest": manifest_digest,
                "size": manifest_size,
                "annotations": {
                    "org.opencontainers.image.ref.name": ref_name,
                    "org.opencontainers.image.title": spec.name,
                },
                "platform": {"os": os_name, "architecture": architecture},
            }
        ],
    }
    (layout_dir / "oci-layout").write_text(
        json.dumps({"imageLayoutVersion": OCI_LAYOUT_VERSION}, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    (layout_dir / "index.json").write_text(
        json.dumps(index_payload, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )

    return {
        "environment": spec.name,
        "version": spec.version,
        "reference": ref_name,
        "layout_dir": str(layout_dir),
        "manifest_digest": manifest_digest,
        "config_digest": config_digest,
        "layers": [
            {"title": "vmlinux", "digest": kernel_digest, "size": kernel_size},
            {"title": "rootfs.ext4", "digest": rootfs_digest, "size": rootfs_size},
            {"title": "environment.json", "digest": metadata_digest, "size": metadata_size},
        ],
    }


def generate_manifest(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> dict[str, Any]:
    manifest: dict[str, Any] = {
        "bundle_version": lock.bundle_version,
        "platform": lock.platform,
        "component_versions": lock.component_versions,
        "capabilities": lock.capabilities,
        "binaries": {},
        "guest": {},
        "profiles": {},
    }
    for name, relative_path in lock.binaries.items():
        full_path = paths.build_platform_root / relative_path
        manifest["binaries"][name] = {"path": relative_path, "sha256": _sha256(full_path)}
    for name, artifact in lock.guest.items():
        full_path = paths.build_platform_root / artifact["path"]
        manifest["guest"][name] = {
            "path": artifact["path"],
            "sha256": _sha256(full_path),
        }
    for name, profile in lock.profiles.items():
        kernel_path = paths.build_platform_root / profile["kernel"]
        rootfs_path = paths.build_platform_root / profile["rootfs"]
        manifest["profiles"][name] = {
            "description": profile["description"],
            "kernel": {"path": profile["kernel"], "sha256": _sha256(kernel_path)},
            "rootfs": {"path": profile["rootfs"], "sha256": _sha256(rootfs_path)},
        }
    manifest_path = paths.build_platform_root / "manifest.json"
    manifest_path.write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return manifest


def sync_bundle(paths: RuntimeBuildPaths) -> None:
    bundle_platform_dir = paths.bundle_dir / paths.platform
    bundle_notice_path = paths.bundle_dir / "NOTICE"
    if bundle_platform_dir.exists():
        shutil.rmtree(bundle_platform_dir)
    bundle_platform_dir.parent.mkdir(parents=True, exist_ok=True)
    shutil.copytree(paths.build_platform_root, bundle_platform_dir)
    _copy_file(paths.build_root / "NOTICE", bundle_notice_path)


def build_bundle(paths: RuntimeBuildPaths, *, sync: bool) -> dict[str, Any]:
    lock = _load_lock(paths)
    if lock.platform != paths.platform:
        raise RuntimeError(
            "runtime lock platform "
            f"{lock.platform!r} does not match requested platform {paths.platform!r}"
        )
    validate_sources(paths, lock)
    if paths.build_root.exists():
        shutil.rmtree(paths.build_root)
    paths.build_platform_root.mkdir(parents=True, exist_ok=True)
    _copy_notice(paths)
    stage_binaries(paths, lock)
    stage_kernel(paths, lock)
    stage_rootfs(paths, lock)
    stage_agent(paths, lock)
    manifest = generate_manifest(paths, lock)
    if sync:
        sync_bundle(paths)
    return manifest


def materialize_sources(paths: RuntimeBuildPaths) -> None:
    lock = _load_lock(paths)
    materialize_binaries(paths, lock)
    materialize_kernel(paths, lock)
    materialize_rootfs(paths, lock)


def _build_paths(
    *,
    source_dir: Path,
    build_dir: Path,
    bundle_dir: Path,
    platform: str,
    materialized_dir: Path = DEFAULT_RUNTIME_MATERIALIZED_DIR,
) -> RuntimeBuildPaths:
    return RuntimeBuildPaths(
        source_root=source_dir,
        source_platform_root=source_dir / platform,
        build_root=build_dir,
        build_platform_root=build_dir / platform,
        bundle_dir=bundle_dir,
        materialized_root=materialized_dir,
        materialized_platform_root=materialized_dir / platform,
        platform=platform,
    )


def _build_parser() -> argparse.ArgumentParser:  # pragma: no cover - CLI wiring
    parser = argparse.ArgumentParser(description="Build packaged runtime bundles for pyro-mcp.")
    parser.add_argument(
        "command",
        choices=[
            "fetch-binaries",
            "build-kernel",
            "build-rootfs",
            "materialize",
            "export-environment-oci",
            "publish-environment-oci",
            "stage-binaries",
            "stage-kernel",
            "stage-rootfs",
            "stage-agent",
            "validate",
            "manifest",
            "sync",
            "bundle",
        ],
    )
    parser.add_argument("--platform", default=DEFAULT_PLATFORM)
    parser.add_argument("--source-dir", default=str(DEFAULT_RUNTIME_SOURCE_DIR))
    parser.add_argument("--build-dir", default=str(DEFAULT_RUNTIME_BUILD_DIR))
    parser.add_argument("--bundle-dir", default=str(DEFAULT_RUNTIME_BUNDLE_DIR))
    parser.add_argument("--materialized-dir", default=str(DEFAULT_RUNTIME_MATERIALIZED_DIR))
    parser.add_argument("--environment")
    parser.add_argument("--output-dir", default=str(DEFAULT_RUNTIME_OCI_LAYOUT_DIR))
    parser.add_argument("--layout-root", default=str(DEFAULT_RUNTIME_OCI_LAYOUT_DIR))
    parser.add_argument("--registry")
    parser.add_argument("--repository")
    parser.add_argument("--reference")
    parser.add_argument("--username-env", default=DEFAULT_OCI_USERNAME_ENV)
    parser.add_argument("--password-env", default=DEFAULT_OCI_PASSWORD_ENV)
    return parser


def main() -> None:  # pragma: no cover - CLI wiring
    args = _build_parser().parse_args()
    paths = _build_paths(
        source_dir=Path(args.source_dir),
        build_dir=Path(args.build_dir),
        bundle_dir=Path(args.bundle_dir),
        materialized_dir=Path(args.materialized_dir),
        platform=args.platform,
    )
    if args.command == "fetch-binaries":
        lock = _load_lock(paths)
        materialize_binaries(paths, lock)
        return
    if args.command == "build-kernel":
        lock = _load_lock(paths)
        materialize_kernel(paths, lock)
        return
    if args.command == "build-rootfs":
        lock = _load_lock(paths)
        materialize_rootfs(paths, lock)
        return
    if args.command == "materialize":
        materialize_sources(paths)
        return
    if args.command == "export-environment-oci":
        if not isinstance(args.environment, str) or args.environment == "":
            raise RuntimeError("--environment is required for export-environment-oci")
        result = export_environment_oci_layout(
            paths,
            environment=args.environment,
            output_dir=Path(args.output_dir),
            reference=args.reference,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
        return
    if args.command == "publish-environment-oci":
        if not isinstance(args.environment, str) or args.environment == "":
            raise RuntimeError("--environment is required for publish-environment-oci")
        layout_root = Path(args.layout_root)
        layout_index = layout_root / _environment_slug(args.environment) / "index.json"
        if not layout_index.exists():
            export_environment_oci_layout(
                paths,
                environment=args.environment,
                output_dir=layout_root,
                reference=args.reference,
            )
        result = publish_environment_oci_layout(
            environment=args.environment,
            layout_root=layout_root,
            registry=args.registry,
            repository=args.repository,
            reference=args.reference,
            username_env=args.username_env,
            password_env=args.password_env,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
        return
    if args.command == "bundle":
        build_bundle(paths, sync=True)
        return
    if args.command == "stage-binaries":
        lock = _load_lock(paths)
        paths.build_platform_root.mkdir(parents=True, exist_ok=True)
        _copy_notice(paths)
        stage_binaries(paths, lock)
        return
    if args.command == "stage-kernel":
        lock = _load_lock(paths)
        paths.build_platform_root.mkdir(parents=True, exist_ok=True)
        stage_kernel(paths, lock)
        return
    if args.command == "stage-rootfs":
        lock = _load_lock(paths)
        paths.build_platform_root.mkdir(parents=True, exist_ok=True)
        stage_rootfs(paths, lock)
        return
    if args.command == "stage-agent":
        lock = _load_lock(paths)
        paths.build_platform_root.mkdir(parents=True, exist_ok=True)
        stage_agent(paths, lock)
        return
    if args.command == "validate":
        lock = _load_lock(paths)
        validate_sources(paths, lock)
        return
    if args.command == "manifest":
        lock = _load_lock(paths)
        generate_manifest(paths, lock)
        return
    if args.command == "sync":
        sync_bundle(paths)
        return
    raise RuntimeError(f"unknown command: {args.command}")


if __name__ == "__main__":  # pragma: no cover - CLI entrypoint
    main()
