scripts/pixelmodel_robustness/cross_codec.py

"""Task 1のcross-codec adapter、条件registry、protected payloadを提供する。"""

from __future__ import annotations

import hashlib
import json
import logging
import re
import subprocess
import tempfile
from collections.abc import Mapping, Sequence
from dataclasses import asdict, dataclass, field
from io import BytesIO
from types import MappingProxyType
from typing import Final

import numpy as np
from PIL import Image, features
from PIL import __version__ as PILLOW_VERSION

from pixelmodel_robustness.jpeg_repair import RepairResult, repair_nonfinite

LOGGER = logging.getLogger(__name__)
MAGICK_PATH: Final[str] = "/etc/profiles/per-user/scratchbrulee/bin/magick"
LOSSY_CODEC_NAMES: Final[tuple[str, ...]] = ("jpeg_q100", "jpeg_q80", "webp_q80", "avif_q70", "jxl_d1")

CONDITION_IDS: Final[tuple[str, ...]] = (
    "png_baseline", "webp_lossless", "jpeg_q100_raw", "jpeg_q100_repair_zero",
    "jpeg_q100_protected_high", "jpeg_q80_raw", "jpeg_q80_repair_zero",
    "jpeg_q80_protected_high", "webp_q80_raw", "webp_q80_repair_zero",
    "webp_q80_protected_high", "avif_q70_raw", "avif_q70_repair_zero",
    "avif_q70_protected_high", "jxl_d1_raw", "jxl_d1_repair_zero",
    "jxl_d1_protected_high",
)


@dataclass(frozen=True)
class CodecSpec:
    """一つのcodecの固定encode/decode契約。"""

    name: str
    format: str
    quality: int | None
    options: Mapping[str, object]
    metadata: Mapping[str, object]
    grayscale: bool = True


@dataclass(frozen=True)
class ConditionSpec:
    """ordered registryのsemantic entry。"""

    condition_id: str
    codec: str
    mode: str
    repair: str | None
    protected: bool


@dataclass(frozen=True)
class CodecPreflight:
    """codec availabilityと実行identity。"""

    name: str
    available: bool
    version: str | None
    reason: str | None
    argv: tuple[str, ...]
    metadata: Mapping[str, object]
    diagnostics: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))


@dataclass(frozen=True)
class ProtectedArtifact:
    """high PNGとlossy low planeのin-memory payload。"""

    high_png_bytes: bytes
    low_codec_bytes: bytes
    low_codec: str

    @property
    def payload(self) -> bytes:
        """high bytesとlow bytesを境界なしで結合したpayloadを返す。"""
        return self.high_png_bytes + self.low_codec_bytes


class JxlCodecError(RuntimeError):
    """JXL subprocessの失敗をdiagnostics付きで表す。"""

    def __init__(self, message: str, *, argv: Sequence[str], stdout: str, stderr: str, returncode: int) -> None:
        super().__init__(f"{message}: returncode={returncode}")
        self.argv = tuple(argv)
        self.stdout = stdout
        self.stderr = stderr
        self.returncode = returncode


def _mapping(values: Mapping[str, object]) -> Mapping[str, object]:
    return MappingProxyType(dict(values))


CODEC_SPECS: Final[Mapping[str, CodecSpec]] = MappingProxyType({
    "png": CodecSpec("png", "PNG", None, _mapping({}), _mapping({})),
    "webp_lossless": CodecSpec(
        "webp_lossless", "WEBP", None,
        _mapping({"lossless": True, "exact": True, "method": 6}), _mapping({}),
    ),
    "jpeg_q100": CodecSpec(
        "jpeg_q100", "JPEG", 100,
        _mapping({"quality": 100, "subsampling": 0, "optimize": False}), _mapping({}),
    ),
    "jpeg_q80": CodecSpec(
        "jpeg_q80", "JPEG", 80,
        _mapping({"quality": 80, "subsampling": 0, "optimize": False}), _mapping({}),
    ),
    "webp_q80": CodecSpec(
        "webp_q80", "WEBP", 80,
        _mapping({"quality": 80, "lossless": False, "method": 6, "exact": True}), _mapping({}),
    ),
    "avif_q70": CodecSpec(
        "avif_q70", "AVIF", 70,
        _mapping({"quality": 70, "subsampling": "4:4:4", "speed": 6, "range": "full", "codec": "aom", "max_threads": 1, "autotiling": False}),
        _mapping({}),
    ),
    "jxl_d1": CodecSpec(
        "jxl_d1", "JXL", 90,
        _mapping({"quality": 90}),
        _mapping({"quality": 90, "libjxl_distance": 1.0, "quality_mapping": "JxlEncoderDistanceFromQuality(90)"}),
    ),
})

EXPECTED_FORMAT_BY_CODEC: Final[Mapping[str, str]] = MappingProxyType({
    "png": "PNG", "webp_lossless": "WEBP", "jpeg_q100": "JPEG", "jpeg_q80": "JPEG",
    "webp_q80": "WEBP", "avif_q70": "AVIF", "jxl_d1": "JXL",
})


def _build_registry() -> Mapping[str, ConditionSpec]:
    entries: dict[str, ConditionSpec] = {}
    for condition_id in CONDITION_IDS:
        if condition_id == "png_baseline":
            codec, mode, repair, protected = "png", "raw", None, False
        elif condition_id == "webp_lossless":
            codec, mode, repair, protected = "webp_lossless", "raw", None, False
        else:
            codec = _condition_codec(condition_id)
            protected = condition_id.endswith("_protected_high")
            repair = "zero" if condition_id.endswith("_repair_zero") else None
            mode = "protected_high" if protected else "raw" if repair is None else "repair_zero"
        entries[condition_id] = ConditionSpec(condition_id, codec, mode, repair, protected)
    return MappingProxyType(entries)


def _condition_codec(condition_id: str) -> str:
    for codec in ("jpeg_q100", "jpeg_q80", "webp_q80", "avif_q70", "jxl_d1"):
        if condition_id.startswith(codec):
            return codec
    raise ValueError(f"unknown condition: {condition_id}")


CONDITION_REGISTRY: Final[Mapping[str, ConditionSpec]] = _build_registry()


def semantic_hash(registry: Mapping[str, ConditionSpec] | Sequence[ConditionSpec]) -> str:
    """registryの順序と全semantic fieldをcanonical JSONでhashする。"""
    values = registry.values() if isinstance(registry, Mapping) else registry
    entries = [CONDITION_REGISTRY[item] if isinstance(item, str) else item for item in values]
    encoded = json.dumps([asdict(item) for item in entries], sort_keys=True, separators=(",", ":")).encode()
    return hashlib.sha256(encoded).hexdigest()


def _pillow_available(format_name: str) -> tuple[bool, str | None]:
    Image.init()
    if format_name == "AVIF":
        available = bool(features.check("avif"))
    else:
        available = format_name in Image.SAVE
    return available, None if available else f"Pillow codec unavailable: {format_name}"


def _jxl_run(args: Sequence[str]) -> subprocess.CompletedProcess[str]:
    argv = [MAGICK_PATH, *args]
    try:
        result = subprocess.run(argv, capture_output=True, text=True, shell=False, check=False)
    except OSError as error:
        raise JxlCodecError(str(error), argv=argv, stdout="", stderr=str(error), returncode=-1) from error
    if result.returncode != 0:
        raise JxlCodecError("ImageMagick JXL command failed", argv=argv, stdout=result.stdout, stderr=result.stderr, returncode=result.returncode)
    return result


def _jxl_version() -> CodecPreflight:
    argv = (MAGICK_PATH, "-version")
    try:
        result = _jxl_run(["-version"])
    except JxlCodecError as error:
        return CodecPreflight(
            "jxl_d1", False, None, str(error), tuple(error.argv), _mapping({"quality": 90, "libjxl_distance": 1.0}),
            _mapping({"argv": error.argv, "stdout": error.stdout, "stderr": error.stderr, "returncode": error.returncode}),
        )
    if "jxl" not in result.stdout.lower():
        diagnostics = _mapping({"argv": argv, "stdout": result.stdout, "stderr": result.stderr, "returncode": result.returncode})
        return CodecPreflight("jxl_d1", False, result.stdout, "ImageMagick version has no JXL delegate", argv, _mapping({}), diagnostics)
    try:
        formats = _jxl_run(["-list", "format"])
        delegate_version = _parse_libjxl_delegate_version(formats.stdout)
    except JxlCodecError as error:
        return CodecPreflight(
            "jxl_d1", False, result.stdout, str(error), argv, _mapping({}),
            _mapping({"argv": error.argv, "stdout": error.stdout, "stderr": error.stderr, "returncode": error.returncode}),
        )
    if delegate_version is None:
        diagnostics = _mapping({"argv": (MAGICK_PATH, "-list", "format"), "stdout": formats.stdout, "stderr": formats.stderr, "returncode": formats.returncode})
        return CodecPreflight("jxl_d1", False, result.stdout, "libjxl delegate version could not be parsed", argv, _mapping({}), diagnostics)
    probe_image = Image.fromarray(np.arange(8 * 8 * 3, dtype=np.uint8).reshape(8, 8, 3), "RGB")
    try:
        q90 = _jxl_encode(probe_image, quality=90)
        q50 = _jxl_encode(probe_image, quality=50)
        decoded_q90 = _jxl_decode(q90)
        decoded_q50 = _jxl_decode(q50)
        if not q90 or not q50 or q90 == q50 or decoded_q90.size != probe_image.size or decoded_q50.size != probe_image.size:
            raise ValueError("JXL probe payloads are empty, identical, or have wrong dimensions")
    except JxlCodecError as error:
        stages = _mapping({
            "version": _mapping({"argv": argv, "stdout": result.stdout, "stderr": result.stderr, "returncode": result.returncode}),
            "format": _mapping({"argv": (MAGICK_PATH, "-list", "format"), "stdout": formats.stdout, "stderr": formats.stderr, "returncode": formats.returncode}),
            "probe": _mapping({"argv": error.argv, "stdout": error.stdout, "stderr": error.stderr, "returncode": error.returncode}),
        })
        return CodecPreflight("jxl_d1", False, result.stdout, str(error), tuple(error.argv), _mapping({}), stages)
    except (OSError, ValueError) as error:
        return CodecPreflight(
            "jxl_d1", False, result.stdout, str(error), argv,
            _mapping({"quality": 90, "libjxl_distance": 1.0, "stderr": result.stderr}),
        )
    metadata = {
        "quality": 90,
        "libjxl_distance": 1.0,
        "quality_mapping": "JxlEncoderDistanceFromQuality(quality)",
        "imagemagick_full_version": result.stdout,
        "libjxl_delegate_version": delegate_version,
        "stderr": result.stderr,
        "probe_q90_bytes": len(q90),
        "probe_q50_bytes": len(q50),
        "probe_q90_sha256": hashlib.sha256(q90).hexdigest(),
        "probe_q50_sha256": hashlib.sha256(q50).hexdigest(),
        "probe_dimensions": (probe_image.height, probe_image.width),
        "probe_q50_dimensions": (decoded_q50.height, decoded_q50.width),
        "argv_template": (MAGICK_PATH, "-quality", "90", "<source>", "<target>"),
    }
    diagnostics = _mapping({"argv": argv, "stdout": result.stdout, "stderr": result.stderr, "returncode": result.returncode})
    return CodecPreflight("jxl_d1", True, result.stdout, None, argv, _mapping(metadata), diagnostics)


def _parse_libjxl_delegate_version(output: str) -> str | None:
    match = re.search(r"libjxl\s+([0-9]+\.[0-9]+\.[0-9]+)", output, flags=re.IGNORECASE)
    return match.group(1) if match else None


def preflight_codecs() -> tuple[CodecPreflight, ...]:
    """全codecのavailabilityをfail-closedで返す。"""
    results: list[CodecPreflight] = []
    for name, spec in CODEC_SPECS.items():
        if name == "jxl_d1":
            results.append(_jxl_version())
            continue
        available, reason = _pillow_available(spec.format)
        results.append(CodecPreflight(name, available, PILLOW_VERSION if available else None, reason, ("Pillow", spec.format), spec.metadata))
    return tuple(results)


def _verified_status(spec: CodecSpec, verified_preflight: Mapping[str, CodecPreflight] | str | None) -> CodecPreflight | None:
    """verified preflight tokenを検証し、standalone時はNoneを返す。"""
    if verified_preflight is None:
        return None
    if isinstance(verified_preflight, str):
        if verified_preflight != spec.name:
            raise ValueError(f"verified preflight token mismatch: {verified_preflight} != {spec.name}")
        return CodecPreflight(spec.name, True, "verified-token", None, ("verified", spec.name), spec.metadata)
    status = verified_preflight.get(spec.name)
    if status is None or status.name != spec.name or not status.available:
        raise RuntimeError(status.reason if status is not None and status.reason else f"codec unavailable: {spec.name}")
    return status


def _require_available(spec: CodecSpec, verified_preflight: Mapping[str, CodecPreflight] | str | None = None) -> None:
    status = _verified_status(spec, verified_preflight)
    if status is None and spec.name == "jxl_d1":
        status = _jxl_version()
    elif status is None:
        available, reason = _pillow_available(spec.format)
        status = CodecPreflight(spec.name, available, PILLOW_VERSION if available else None, reason, ("Pillow", spec.format), spec.metadata)
    if not status.available:
        raise RuntimeError(status.reason or f"codec unavailable: {spec.name}")


def encode_rgb(image: Image.Image, spec: CodecSpec, *, verified_preflight: Mapping[str, CodecPreflight] | str | None = None) -> bytes:
    """RGB imageを固定optionでmemory encodeする。"""
    _require_available(spec, verified_preflight)
    rgb = image.convert("RGB")
    if spec.name == "jxl_d1":
        return _jxl_encode(rgb)
    return _encode_pillow(rgb, spec)


def _encode_pillow(image: Image.Image, spec: CodecSpec) -> bytes:
    """Pillowへ入力modeを変更せず固定optionで渡す。"""
    stream = BytesIO()
    image.save(stream, format=spec.format, **dict(spec.options))
    return stream.getvalue()


def decode_rgb(payload: bytes, spec: CodecSpec, *, verified_preflight: Mapping[str, CodecPreflight] | str | None = None) -> Image.Image:
    """bytesをRGB imageへdecodeし、file handleを閉じる。"""
    _verified_status(spec, verified_preflight)
    return _decode_l_image(payload, spec).convert("RGB")


def _decode_image(payload: bytes, spec: CodecSpec) -> Image.Image:
    """payloadをdecoderの実modeのままmemory imageへdecodeする。"""
    if spec.name == "jxl_d1":
        return _jxl_decode(payload)
    with Image.open(BytesIO(payload)) as image:
        return image.copy()


def encode_l_plane(plane: np.ndarray, spec: CodecSpec, *, verified_preflight: Mapping[str, CodecPreflight] | str | None = None) -> bytes:
    """grayscale planeを固定codec optionでmemory encodeする。"""
    array = np.asarray(plane, dtype=np.uint8)
    if array.ndim != 2:
        raise ValueError("L plane must be two-dimensional")
    _require_available(spec, verified_preflight)
    image = Image.fromarray(array, "L")
    if spec.name == "jxl_d1":
        return _jxl_encode(image)
    return _encode_pillow(image, spec)


def decode_l_plane(payload: bytes, spec: CodecSpec, *, verified_preflight: Mapping[str, CodecPreflight] | str | None = None) -> np.ndarray:
    """bytesをgrayscale uint8 planeへdecodeする。"""
    _verified_status(spec, verified_preflight)
    image = _decode_l_image(payload, spec)
    return np.asarray(image.convert("L"), dtype=np.uint8).copy()


def _decode_l_image(payload: bytes, spec: CodecSpec) -> Image.Image:
    """L plane payloadをcodec format gate付きでdecodeする。"""
    expected_format = EXPECTED_FORMAT_BY_CODEC.get(spec.name)
    if expected_format is None:
        raise ValueError(f"unknown codec name: {spec.name}")
    if spec.format != expected_format:
        raise ValueError(f"codec format is not canonical for {spec.name}: {spec.format}")
    if spec.name == "jxl_d1":
        return _jxl_decode(payload)
    with Image.open(BytesIO(payload)) as image:
        if image.format != expected_format:
            raise ValueError(f"payload format {image.format} does not match codec format {expected_format}")
        return image.copy()


def _jxl_encode(image: Image.Image, *, quality: int = 90) -> bytes:
    with tempfile.TemporaryDirectory(prefix="pixelmodel-jxl-") as directory:
        source, target = f"{directory}/source.png", f"{directory}/payload.jxl"
        image.save(source, format="PNG")
        _jxl_run(["-quality", str(quality), source, target])
        with open(target, "rb") as stream:
            return stream.read()


def _jxl_decode(payload: bytes) -> Image.Image:
    with tempfile.TemporaryDirectory(prefix="pixelmodel-jxl-") as directory:
        source, target = f"{directory}/payload.jxl", f"{directory}/decoded.png"
        with open(source, "wb") as stream:
            stream.write(payload)
        identify_argv = [MAGICK_PATH, "identify", "-format", "%m", source]
        identified = _jxl_run(["identify", "-format", "%m", source])
        if identified.stdout.strip() != "JXL":
            raise JxlCodecError(
                "JXL identify format gate failed", argv=identify_argv,
                stdout=identified.stdout, stderr=identified.stderr, returncode=identified.returncode,
            )
        _jxl_run([source, target])
        with Image.open(target) as image:
            return image.copy()


def encode_protected_high(values: np.ndarray, shape: tuple[int, int], low_spec: CodecSpec, *, verified_preflight: Mapping[str, CodecPreflight] | str | None = None) -> ProtectedArtifact:
    """high byteをPNG、low byteを指定lossy codecで保存する。"""
    if low_spec.name not in LOSSY_CODEC_NAMES:
        raise ValueError(f"protected low codec must be lossy: {low_spec.name}")
    canonical = CODEC_SPECS[low_spec.name]
    if (
        low_spec.format != canonical.format
        or low_spec.quality != canonical.quality
        or low_spec.options != canonical.options
        or low_spec.metadata != canonical.metadata
        or low_spec.grayscale != canonical.grayscale
    ):
        raise ValueError(f"protected low codec spec is not semantically canonical: {low_spec.name}")
    flat = np.asarray(values, dtype=np.float16).reshape(-1)
    capacity = int(np.prod(shape))
    if flat.size > capacity:
        raise ValueError("values exceed shape capacity")
    bits = np.zeros(capacity, dtype=np.uint16)
    bits[: flat.size] = flat.view(np.uint16)
    high = (bits >> 8).astype(np.uint8).reshape(shape)
    low = bits.astype(np.uint8).reshape(shape)
    high_stream = BytesIO()
    Image.fromarray(high, "L").save(high_stream, format="PNG")
    return ProtectedArtifact(high_stream.getvalue(), encode_l_plane(low, low_spec, verified_preflight=verified_preflight), low_spec.name)


def decode_protected_high(artifact: ProtectedArtifact, shape: tuple[int, int], source_values: np.ndarray, *, verified_preflight: Mapping[str, CodecPreflight] | str | None = None) -> tuple[np.ndarray, dict[str, object]]:
    """protected payloadを復元し、high exactnessとbytes/hashを検証する。"""
    with Image.open(BytesIO(artifact.high_png_bytes)) as high_image:
        if high_image.format != "PNG":
            raise ValueError(f"protected high plane format must be PNG, got {high_image.format}")
        if high_image.mode != "L":
            raise ValueError(f"protected high plane mode must be L, got {high_image.mode}")
        high = np.asarray(high_image.convert("L"), dtype=np.uint8).copy()
        high_mode, high_dimensions = high_image.mode, [high_image.height, high_image.width]
    if artifact.low_codec not in LOSSY_CODEC_NAMES:
        raise ValueError(f"protected low codec must be lossy: {artifact.low_codec}")
    _verified_status(CODEC_SPECS[artifact.low_codec], verified_preflight)
    low_image = _decode_l_image(artifact.low_codec_bytes, CODEC_SPECS[artifact.low_codec])
    low_mode = low_image.mode
    low = np.asarray(low_image.convert("L"), dtype=np.uint8).copy()
    if high.shape != shape or low.shape != shape:
        raise ValueError("protected plane dimensions mismatch")
    source_flat = np.asarray(source_values, dtype=np.float16).reshape(-1)
    if source_flat.size > int(np.prod(shape)):
        raise ValueError("source values exceed shape capacity")
    expected_bits = np.zeros(int(np.prod(shape)), dtype=np.uint16)
    expected_bits[: source_flat.size] = source_flat.view(np.uint16)
    source_high = (expected_bits >> 8).astype(np.uint8).reshape(shape)
    if not np.array_equal(high, source_high):
        raise ValueError("decoded protected high plane is not bit-exact")
    bits = ((high.reshape(-1).astype(np.uint16) << 8) | low.reshape(-1).astype(np.uint16)).view(np.float16)
    metadata: dict[str, object] = {
        "dimensions": list(shape), "high_mode": high_mode, "low_mode": low_mode, "high_dimensions": high_dimensions,
        "low_dimensions": [low_image.height, low_image.width],
        "decoded_high_exact": True, "decoded_high_sha256": hashlib.sha256(high.tobytes()).hexdigest(),
        "combined_payload_bytes": len(artifact.payload), "combined_payload_sha256": hashlib.sha256(artifact.payload).hexdigest(),
        "low_codec": artifact.low_codec, "source_value_count": int(source_flat.size),
        "low_changed_count": int(np.count_nonzero(low != expected_bits.astype(np.uint8).reshape(shape))),
        "high_encoded_bytes": len(artifact.high_png_bytes), "high_encoded_sha256": hashlib.sha256(artifact.high_png_bytes).hexdigest(),
        "low_encoded_bytes": len(artifact.low_codec_bytes), "low_encoded_sha256": hashlib.sha256(artifact.low_codec_bytes).hexdigest(),
        "decoded_low_sha256": hashlib.sha256(low.tobytes()).hexdigest(),
        "pillow_version": PILLOW_VERSION,
    }
    return bits, metadata


def repair_nonfinite_zero(values: np.ndarray, layers: Mapping[str, tuple[int, int]]) -> RepairResult:
    """NaN/+Inf/-Infだけをzeroへ置換し、finite fp16 bit patternを保持する。"""
    return repair_nonfinite(values, "zero", layers)