"""Manual Codex built-in image-generation bridge for SMA packages."""
from __future__ import annotations
import hashlib
import os
import shutil
import subprocess
from pathlib import Path
from typing import Any
from PIL import Image, UnidentifiedImageError
CODEX = os.environ.get("SMA_CODEX_BINARY", "/usr/bin/codex")
CODEX_HOME = Path(os.environ.get("CODEX_HOME", "/root/.codex")).resolve()
TRUSTED_WORKSPACE = Path(os.environ.get("SMA_CODEX_WORKSPACE", "/tmp/aa014-codex-lVvqrS")).resolve()
MAX_BYTES = 25 * 1024 * 1024
class CodexImageError(RuntimeError):
"""Safe failure from the Codex image-generation bridge."""
def _png_info(path: Path) -> tuple[int, int, int, str]:
data = path.read_bytes()
if not data or len(data) > MAX_BYTES:
raise CodexImageError("Codex lieferte keine gültigen Bilddaten")
try:
with Image.open(path) as image:
size = image.size
image.verify()
except (UnidentifiedImageError, OSError, SyntaxError) as exc:
raise CodexImageError("Codex lieferte kein lesbares Bild") from exc
return size[0], size[1], len(data), hashlib.sha256(data).hexdigest()
def _image_paths() -> set[Path]:
return {p.resolve() for p in CODEX_HOME.joinpath("generated_images").glob("**/*") if p.is_file() and p.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}}
def generate_with_codex(prompt: str, output_path: Path, *, timeout: int = 600) -> dict[str, Any]:
if not prompt.strip():
raise CodexImageError("Bildprompt fehlt")
if not TRUSTED_WORKSPACE.is_dir():
raise CodexImageError(f"Codex-Arbeitskontext fehlt: {TRUSTED_WORKSPACE}")
before = _image_paths()
child_env = os.environ.copy()
child_env["HOME"] = "/root"
child_env["CODEX_HOME"] = str(CODEX_HOME)
child_env.pop("OPENAI_API_KEY", None)
command = [CODEX, "exec", "--sandbox", "workspace-write", "--skip-git-repo-check", "--ephemeral", "-C", str(TRUSTED_WORKSPACE), "--json", prompt]
try:
completed = subprocess.run(command, env=child_env, text=True, capture_output=True, timeout=timeout, check=False)
except (OSError, subprocess.TimeoutExpired) as exc:
raise CodexImageError(f"Codex-Prozess konnte nicht abgeschlossen werden: {type(exc).__name__}") from exc
if completed.returncode != 0:
raise CodexImageError(f"Codex-Prozess fehlgeschlagen (Exit {completed.returncode})")
new_images = sorted(_image_paths() - before, key=lambda p: p.stat().st_mtime_ns)
if len(new_images) != 1:
raise CodexImageError(f"Codex lieferte nicht genau ein neues Bild: {len(new_images)}")
source = new_images[0]
width, height, byte_count, sha256 = _png_info(source)
output_path = Path(output_path).resolve()
if output_path.exists():
raise CodexImageError(f"Zielbild existiert bereits: {output_path}")
output_path.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(source, output_path)
return {"origin_path": str(source), "path": str(output_path), "file_type": "PNG", "width": width, "height": height, "bytes": byte_count, "sha256": sha256, "codex_output": completed.stdout[-4000:]}