workbuddy_skills/draw-ui/scripts/generate_image.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import base64
import binascii
import json
import mimetypes
import os
import re
import sys
import tempfile
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterable

DEFAULT_MODEL = os.getenv("DRAW_MODEL", "openai/gpt-image-2")
DEFAULT_PROVIDER = os.getenv("DRAW_PROVIDER", "zenmux")
DEFAULT_BASE_URL = os.getenv("ZENMUX_VERTEX_BASE_URL", "https://zenmux.ai/api/vertex-ai")
DEFAULT_OUTPUT_ROOT = Path.home() / ".local" / "share" / "draw" / "outputs"
DEFAULT_MIME = "image/png"
DEFAULT_CODEX_MODEL = os.getenv("DRAW_CODEX_MODEL", "gpt-5.6")




def _read_env_value(path: Path, key: str) -> str:
    try:
        text = path.read_text(encoding="utf-8")
    except FileNotFoundError:
        return ""
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or "=" not in stripped:
            continue
        name, value = stripped.split("=", 1)
        if name.strip() != key:
            continue
        value = value.strip().strip('"').strip("'")
        return value
    return ""


def resolve_api_key() -> str:
    env = os.getenv("ZENMUX_API_KEY", "").strip()
    if env:
        return env
    cwd = Path.cwd().resolve()
    for directory in [cwd, *cwd.parents]:
        found = _read_env_value(directory / ".env.local", "ZENMUX_API_KEY")
        if found:
            return found
    config_path = Path.home() / ".config" / "see" / "api_key"
    try:
        return config_path.read_text(encoding="utf-8").strip()
    except FileNotFoundError:
        return ""


def sanitize_name(value: str, fallback: str = "image") -> str:
    value = value.strip()
    value = re.sub(r"[\\/:*?\"<>|]+", "-", value)
    value = re.sub(r"\s+", "-", value)
    value = re.sub(r"-+", "-", value).strip("-_.")
    if not value:
        return fallback
    return value[:80]


def build_output_path(*, output_arg: str, image_type: str, topic: str, explicit_name: str, ext: str) -> Path:
    if output_arg:
        out = Path(output_arg).expanduser().resolve()
        if out.suffix:
            return out
        return out.with_suffix(ext)
    now = datetime.now()
    day_dir = DEFAULT_OUTPUT_ROOT / now.strftime("%Y-%m-%d")
    day_dir.mkdir(parents=True, exist_ok=True)
    base_name = sanitize_name(explicit_name or topic, fallback=image_type)
    return day_dir / f"{now.strftime('%Y%m%d-%H%M%S')}__{image_type}__{base_name}{ext}"


def metadata_path_for(image_path: Path) -> Path:
    return image_path.with_suffix(image_path.suffix + ".json")


def _output_paths_for_check(output_path: Path) -> list[Path]:
    """Return every path that a writer may choose before its MIME type is known."""
    output_path = output_path.resolve()
    if output_path.suffix:
        return [output_path]

    # The current writers default to PNG, while render_response may infer another
    # image extension from the model response. Include existing siblings so an
    # unknown future MIME type cannot silently replace one of them.
    paths = [output_path.with_suffix(".png")]
    pattern = f"{output_path.name}.*"
    paths.extend(path for path in output_path.parent.glob(pattern) if path.is_file())
    return list(dict.fromkeys(paths))


def ensure_output_available(output_path: Path) -> None:
    """Reject an output or its metadata before any remote generation starts."""
    conflicts: list[Path] = []
    for candidate in _output_paths_for_check(output_path):
        if candidate.exists():
            conflicts.append(candidate)
        metadata_path = metadata_path_for(candidate)
        if metadata_path.exists():
            conflicts.append(metadata_path)
    if conflicts:
        paths = ", ".join(str(path) for path in dict.fromkeys(conflicts))
        raise FileExistsError(f"输出或 metadata 已存在，拒绝覆盖：{paths}")


def _write_new_bytes(path: Path, data: bytes) -> None:
    """Create a file without ever replacing an existing output."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("xb") as handle:
        handle.write(data)


def guess_extension(mime_type: str | None) -> str:
    if not mime_type:
        return ".png"
    guessed = mimetypes.guess_extension(mime_type)
    if guessed == ".jpe":
        return ".jpg"
    return guessed or ".png"


# Type only controls aspect ratio, prompt is fully controlled by caller
ASPECT_RATIOS = {
    "ultrawide": "21:9",
    "wide": "16:9",
    "square": "1:1",
    "portrait": "3:4",
    "classic": "4:3",
}

CODEX_SIZE_PRESETS = {
    "ultrawide": "1536x640",
    "wide": "1536x864",
    "classic": "1024x768",
    "square": "1024x1024",
    "portrait": "768x1024",
}

MODE_PROMPTS = {
    "normal": "",
    "replicate": (
        "Use the reference image as the primary visual source. Recreate the UI screen as closely as possible. "
        "Preserve layout, spacing, typography hierarchy, colors, shadows, border radius, icon style, density, "
        "and the relative position of every major element. Do not redesign unless the prompt explicitly asks for changes. "
        "If text is unreadable, preserve its visual length, alignment, and hierarchy. Output only the clean UI mockup, "
        "with no browser chrome, watermark, annotations, or surrounding device frame."
    ),
    "frame-lock": (
        "Use the first reference image as a locked application frame. Preserve the sidebar, top navigation, brand area, "
        "and persistent chrome as closely as possible. Redesign or generate only the content area requested by the prompt. "
        "Keep the result as a clean full-screen UI mockup with no browser chrome or watermark."
    ),
    "asset-redraw": (
        "Use the reference image to recreate only the requested visual asset as a clean standalone asset. Remove surrounding "
        "UI, labels, browser chrome, mockup frames, and unrelated elements unless explicitly requested. Preserve the source "
        "asset's proportions, material, color, and brand feel with high clarity and generous padding."
    ),
}


def effective_prompt(prompt: str, mode: str) -> str:
    mode_prompt = MODE_PROMPTS.get(mode, "")
    if not mode_prompt:
        return prompt
    return f"{mode_prompt}\n\nUser request:\n{prompt}"


def download_file(url: str, dest: Path, timeout: int = 120) -> None:
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req, timeout=timeout) as response:
        dest.write_bytes(response.read())


def resolve_ref(raw: str, tmp_dir: Path) -> Path:
    parsed = urllib.parse.urlparse(raw)
    if parsed.scheme in {"http", "https"}:
        suffix = Path(parsed.path).suffix or ".png"
        dest = tmp_dir / f"ref-{len(list(tmp_dir.iterdir())) + 1}{suffix}"
        download_file(raw, dest)
        return dest
    path = Path(raw).expanduser().resolve()
    if not path.exists():
        raise FileNotFoundError(f"Reference not found: {raw}")
    return path


def load_genai():
    try:
        from google import genai
        from google.genai import types
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "[ERROR] Missing dependency `google-genai`. Re-run via scripts/ask_draw.sh so it can auto-install it."
        ) from exc
    return genai, types


def build_contents(*, prompt: str, types, refs: Iterable[Path]):
    parts = [types.Part.from_text(text=prompt)]
    for ref in refs:
        mime = mimetypes.guess_type(ref.name)[0] or DEFAULT_MIME
        parts.append(types.Part.from_bytes(data=ref.read_bytes(), mime_type=mime))
    return parts


def extract_parts(response) -> list:
    if getattr(response, "parts", None):
        return list(response.parts)
    candidates = getattr(response, "candidates", None) or []
    parts = []
    for candidate in candidates:
        content = getattr(candidate, "content", None)
        if content and getattr(content, "parts", None):
            parts.extend(content.parts)
    return parts


def render_response(*, response, output_path: Path) -> tuple[str, str]:
    text_parts: list[str] = []
    image_written = False
    image_mime = DEFAULT_MIME
    pending_bytes: bytes | None = None

    for part in extract_parts(response):
        text = getattr(part, "text", None)
        if text:
            text_parts.append(text.strip())
        inline_data = getattr(part, "inline_data", None)
        if inline_data:
            data = inline_data.data
            if isinstance(data, str):
                pending_bytes = base64.b64decode(data)
            else:
                pending_bytes = data
            image_mime = getattr(inline_data, "mime_type", None) or DEFAULT_MIME

    if pending_bytes:
        final_path = output_path
        if not output_path.suffix:
            final_path = output_path.with_suffix(guess_extension(image_mime))
        _write_new_bytes(final_path, pending_bytes)
        image_written = True
    else:
        final_path = output_path

    if not image_written:
        raise RuntimeError("Model returned no image data.")

    return final_path.as_posix(), "\n".join([t for t in text_parts if t]).strip()


def resolve_codex_api_key() -> str:
    key = os.getenv("OPENAI_IMAGE_API_KEY") or os.getenv("OPENAI_API_KEY")
    return key.strip() if key else ""


def resolve_codex_base_url() -> str:
    base_url = (os.getenv("OPENAI_IMAGE_BASE_URL") or "https://api.openai.com/v1").rstrip("/")
    parsed = urllib.parse.urlparse(base_url)
    if parsed.scheme and parsed.netloc and parsed.path in ("", "/"):
        return f"{base_url}/v1"
    return base_url


def resolve_codex_model(override: str = "") -> str:
    return override.strip() or DEFAULT_CODEX_MODEL


def join_endpoint(base_url: str, endpoint: str) -> str:
    base = base_url.rstrip("/")
    endpoint = endpoint.lstrip("/")
    if base.endswith("/v1") and endpoint.startswith("v1/"):
        endpoint = endpoint[3:]
    return f"{base}/{endpoint}"


def ref_to_input_image(path: Path) -> dict:
    mime = mimetypes.guess_type(path.name)[0] or DEFAULT_MIME
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return {"type": "input_image", "image_url": f"data:{mime};base64,{encoded}"}


def looks_like_base64_image(value: str) -> bool:
    if len(value) < 200:
        return False
    compact = value.strip()
    if compact.startswith("data:image/"):
        compact = compact.split(",", 1)[-1]
    try:
        head = base64.b64decode(compact[:256] + "==", validate=False)
    except Exception:
        return False
    return head.startswith(b"\x89PNG") or head.startswith(b"\xff\xd8\xff") or head.startswith(b"RIFF")


def find_image_result_recursive(value: object) -> str | None:
    if isinstance(value, dict):
        value_type = value.get("type")
        for key in ("result", "b64_json", "image_base64"):
            item = value.get(key)
            if isinstance(item, str) and (value_type == "image_generation_call" or looks_like_base64_image(item)):
                return item.split(",", 1)[-1] if item.startswith("data:image/") else item
        for item in value.values():
            found = find_image_result_recursive(item)
            if found:
                return found
    elif isinstance(value, list):
        for item in value:
            found = find_image_result_recursive(item)
            if found:
                return found
    elif isinstance(value, str) and looks_like_base64_image(value):
        return value.split(",", 1)[-1] if value.startswith("data:image/") else value
    return None


def request_codex_image(
    *, prompt: str, refs: list[Path], image_type: str, model: str, output_path: Path
) -> Path:
    api_key = resolve_codex_api_key()
    if not api_key:
        raise RuntimeError("No OPENAI_IMAGE_API_KEY or OPENAI_API_KEY found.")

    base_url = resolve_codex_base_url()
    endpoint = join_endpoint(base_url, "responses")
    content: list[dict] = [{"type": "input_text", "text": prompt}]
    content.extend(ref_to_input_image(ref) for ref in refs)
    payload = {
        "model": model,
        "instructions": "Use the image_generation tool to create exactly the requested image. Do not add extra text.",
        "stream": False,
        "store": False,
        "input": [{"role": "user", "content": content}],
        "tools": [{"type": "image_generation", "size": CODEX_SIZE_PRESETS.get(image_type, "1024x1024")}],
        "tool_choice": "required",
    }
    request = urllib.request.Request(
        endpoint,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=600) as response:
            raw = response.read().decode("utf-8")
    except urllib.error.HTTPError as exc:
        details = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} from Codex image API:\n{details}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"Could not reach Codex image API: {exc.reason}") from exc

    try:
        response_payload = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Codex image API returned non-JSON response:\n{raw[:1500]}") from exc

    image_b64 = find_image_result_recursive(response_payload)
    if not image_b64:
        raise RuntimeError("No image result found in Codex Responses API output.")

    final_path = output_path if output_path.suffix else output_path.with_suffix(".png")
    final_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        image_bytes = base64.b64decode(image_b64, validate=True)
    except (binascii.Error, ValueError, TypeError) as exc:
        raise RuntimeError("Codex image API returned invalid base64 image data.") from exc
    _write_new_bytes(final_path, image_bytes)
    return final_path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate UI images via ZenMux or a Codex OpenAI-compatible provider.")
    parser.add_argument("--type", choices=sorted(ASPECT_RATIOS.keys()), default="wide", help="Aspect ratio preset.")
    parser.add_argument("--prompt", required=True, help="Full prompt for image generation.")
    parser.add_argument("--ref", action="append", default=[], help="Reference image path or URL (repeatable).")
    parser.add_argument("--name", default="", help="Optional short output name.")
    parser.add_argument("-o", "--output", default="", help="Output image path.")
    parser.add_argument("--provider", choices=["zenmux", "codex"], default=DEFAULT_PROVIDER, help="Image backend.")
    parser.add_argument("--mode", choices=sorted(MODE_PROMPTS.keys()), default="normal", help="UI prompt wrapper.")
    parser.add_argument("--model", default="", help="Model override (defaults per provider).")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL, help=argparse.SUPPRESS)
    return parser.parse_args()


def _uses_generate_images_api(model: str) -> bool:
    """Models that require the generate_images / edit_image API instead of generate_content."""
    return model.startswith("openai/")


def _run_generate_images(*, client, model: str, prompt: str, refs: list[Path], types) -> tuple[bytes, str]:
    """Call generate_images (or edit_image when refs are provided) and return (image_bytes, response_text)."""
    if refs:
        # Use edit_image with reference images
        # First ref becomes the base image
        base_image_path = refs[0]
        base_mime = mimetypes.guess_type(base_image_path.name)[0] or DEFAULT_MIME
        base_image = types.Image(image_bytes=base_image_path.read_bytes(), mime_type=base_mime)
        reference_images = [
            types.RawReferenceImage(reference_id=1, reference_image=base_image)
        ]
        # Additional refs as extra references
        for i, ref_path in enumerate(refs[1:], start=2):
            ref_mime = mimetypes.guess_type(ref_path.name)[0] or DEFAULT_MIME
            ref_img = types.Image(image_bytes=ref_path.read_bytes(), mime_type=ref_mime)
            reference_images.append(
                types.RawReferenceImage(reference_id=i, reference_image=ref_img)
            )
        response = client.models.edit_image(
            model=model,
            prompt=prompt,
            reference_images=reference_images,
        )
    else:
        response = client.models.generate_images(
            model=model,
            prompt=prompt,
        )

    generated = getattr(response, "generated_images", None)
    if not generated:
        raise RuntimeError("Model returned no generated images.")

    image_obj = generated[0].image
    image_bytes = getattr(image_obj, "image_bytes", None)
    if image_bytes is None:
        # Some versions expose .data as base64
        raw = getattr(image_obj, "data", None)
        if isinstance(raw, str):
            image_bytes = base64.b64decode(raw)
        elif isinstance(raw, bytes):
            image_bytes = raw
    if not image_bytes:
        raise RuntimeError("Could not extract image bytes from generate_images response.")

    return image_bytes, ""


def main() -> int:
    args = parse_args()
    aspect_ratio = ASPECT_RATIOS[args.type]
    prompt = effective_prompt(args.prompt, args.mode)
    model = resolve_codex_model(args.model) if args.provider == "codex" else (args.model or DEFAULT_MODEL)

    output_path = build_output_path(
        output_arg=args.output,
        image_type=args.type,
        topic=args.name or "image",
        explicit_name=args.name,
        ext=".png",
    )
    ensure_output_available(output_path)

    with tempfile.TemporaryDirectory(prefix="draw-refs-") as tmp:
        tmp_dir = Path(tmp)
        refs = [resolve_ref(raw, tmp_dir) for raw in args.ref]

        if args.provider == "codex":
            final_path = request_codex_image(
                prompt=prompt,
                refs=refs,
                image_type=args.type,
                model=model,
                output_path=output_path,
            )
            response_text = ""
        else:
            api_key = resolve_api_key()
            if not api_key:
                print(
                    "[ERROR] No ZENMUX_API_KEY found. Set it as env var, in .env.local, or in ~/.config/see/api_key",
                    file=sys.stderr,
                )
                return 1

            genai, types = load_genai()
            # OpenAI image models via ZenMux can take longer; bump timeout to 5 minutes.
            timeout = 300 if _uses_generate_images_api(model) else 120
            client = genai.Client(
                api_key=api_key,
                vertexai=True,
                http_options=types.HttpOptions(api_version="v1", base_url=args.base_url, timeout=timeout * 1000),
            )

            if _uses_generate_images_api(model):
                image_bytes, response_text = _run_generate_images(
                    client=client, model=model, prompt=prompt, refs=refs, types=types,
                )
                final_path = output_path if output_path.suffix else output_path.with_suffix(".png")
                final_path.parent.mkdir(parents=True, exist_ok=True)
                _write_new_bytes(final_path, image_bytes)
            else:
                response = client.models.generate_content(
                    model=model,
                    contents=build_contents(prompt=prompt, types=types, refs=refs),
                    config=types.GenerateContentConfig(
                        response_modalities=["TEXT", "IMAGE"],
                        image_config=types.ImageConfig(aspect_ratio=aspect_ratio),
                    ),
                )
                final_path_str, response_text = render_response(response=response, output_path=output_path)
                final_path = Path(final_path_str)

        meta_path = metadata_path_for(final_path)
        metadata = {
            "created_at": datetime.now().isoformat(timespec="seconds"),
            "type": args.type,
            "aspect_ratio": aspect_ratio,
            "prompt": prompt,
            "raw_prompt": args.prompt,
            "refs": [str(path) for path in refs],
            "provider": args.provider,
            "mode": args.mode,
            "model": model,
            "base_url": args.base_url if args.provider == "zenmux" else resolve_codex_base_url(),
            "output_path": str(final_path),
            "response_text": response_text,
        }
        meta_path.parent.mkdir(parents=True, exist_ok=True)
        with meta_path.open("x", encoding="utf-8") as handle:
            handle.write(json.dumps(metadata, ensure_ascii=False, indent=2))

    print(f"output_path={final_path}")
    print(f"metadata_path={meta_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())