"""Image generation through the sub2api relay (image group).

Measured 2026-10-30: the relay's image group (group 4) serves
``gpt-image-2.5-flare`` at ``/v1/images/generations`` and returns a
download URL (not ``b64_json``) for the finished PNG. Production
task keys live in group 2 (``codex-system``), whose
``allow_image_generation`` is False, so a task's own gateway key
cannot generate images -- this tool uses a separate system-level
image key, ``IMAGE_GEN_API_KEY``, like the search keys. Never fall
back to the task user's gateway key: it 404s with
``model_not_found`` and the failure looks like a broken tool.

Env (read by the MCP process, never printed in errors):
- ``IMAGE_GEN_BASE_URL``: relay origin, no trailing ``/v1``.
  Defaults to ``SUB2API_BASE_URL`` (local gateway origin,
  reachable from task containers via NO_PROXY), else
  ``https://sub2.diante.me``.
- ``IMAGE_GEN_API_KEY``: system image-group key. Missing key is a
  configuration error, reported as such, never as a model error.
- ``IMAGE_GEN_MODEL``: default model, default
  ``gpt-image-2.5-flare`` (verified 2026-10-30: 200, real PNG
  1024x1024, 829K, via the relay).

Network notes for task containers: the API call to the local
gateway origin must bypass the egress proxy (gateway host in
NO_PROXY); downloading the returned image URL goes through the
egress proxy, so the image CDN host (observed:
``*.qinglanapi.top``) must be on the Tinyproxy whitelist, or the
download step fails after a successful generation -- the tool
reports that distinctly.
"""

import base64
import json
import os
import struct
import time
import urllib.error
import urllib.request

DEFAULT_MODEL = "gpt-image-2.5-flare"
DEFAULT_BASE_URL = "https://sub2.diante.me"
ALLOWED_SIZES = ("1024x1024", "1536x1024", "1024x1536", "auto")
API_TIMEOUT = 180
DOWNLOAD_TIMEOUT = 120


class ImageGenError(RuntimeError):
    """Image generation failure. Messages never contain API keys."""


def _base_url():
    base = (os.environ.get("IMAGE_GEN_BASE_URL", "").strip()
            or os.environ.get("SUB2API_BASE_URL", "").strip()
            or DEFAULT_BASE_URL)
    return base.rstrip("/")


def _api_key():
    return os.environ.get("IMAGE_GEN_API_KEY", "").strip()


def _default_model():
    return os.environ.get("IMAGE_GEN_MODEL", "").strip() or DEFAULT_MODEL


def _detect_format(raw):
    """Return (format, width, height) from magic bytes; width/height
    only for PNG (IHDR), otherwise None -- the caller verifies with
    PIL in shell when it needs exact dimensions."""
    if raw[:8] == b"\x89PNG\r\n\x1a\n" and len(raw) >= 24:
        width, height = struct.unpack(">II", raw[16:24])
        return "png", width, height
    if raw[:3] == b"\xff\xd8\xff":
        return "jpeg", None, None
    if raw[:4] == b"RIFF" and raw[8:12] == b"WEBP":
        return "webp", None, None
    if raw[:6] in (b"GIF87a", b"GIF89a"):
        return "gif", None, None
    return "unknown", None, None


def _post_json(url, key, body, request=None):
    if request is not None:
        return request("POST", url, body, key)
    data = json.dumps(body).encode()
    req = urllib.request.Request(url, data=data, method="POST", headers={
        "Authorization": "Bearer " + key, "Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=API_TIMEOUT) as resp:
            return resp.status, json.loads(resp.read().decode() or "{}")
    except urllib.error.HTTPError as exc:
        raw = exc.read().decode()
        try:
            payload = json.loads(raw or "{}")
        except ValueError:
            payload = {"raw": raw[:400]}
        return exc.code, payload
    except Exception as exc:  # noqa: BLE001 - surfaced as tool error
        return 0, {"error": str(exc)[:300]}


def _download(url, download=None):
    if download is not None:
        return download(url)
    req = urllib.request.Request(url, headers={
        "User-Agent": "codex-imagegen/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=DOWNLOAD_TIMEOUT) as resp:
            return resp.read()
    except Exception as exc:  # noqa: BLE001 - surfaced as tool error
        raise ImageGenError("image download failed: %s" % exc)


def _validate_common(prompt, size, output_filename):
    if not prompt or not str(prompt).strip():
        raise ImageGenError("prompt is required")
    if size not in ALLOWED_SIZES:
        raise ImageGenError("unsupported size: %r (allowed: %s)"
                            % (size, ", ".join(ALLOWED_SIZES)))
    key = _api_key()
    if not key:
        raise ImageGenError(
            "image generation not configured: IMAGE_GEN_API_KEY "
            "is not set in the MCP environment")
    name = output_filename or ("generated-%d.png" % int(time.time()))
    if (os.path.basename(name) != name or name in (".", "..")
            or "/" in name or "\\" in name):
        raise ImageGenError(
            "output_filename must be a plain basename, got %r" % name)
    return key, name


def _save_image_response(status, data, name, use_model, size,
                         download, cwd, verb):
    if status != 200:
        err = (data or {}).get("error") or {}
        msg = err.get("message") if isinstance(err, dict) else str(err)
        raise ImageGenError(
            "image %s failed: HTTP %s: %s"
            % (verb, status, (msg or data or "unknown error")))
    items = (data or {}).get("data") or []
    if not items or not isinstance(items[0], dict):
        raise ImageGenError("image %s returned no image data" % verb)
    first = items[0]
    b64 = first.get("b64_json")
    if b64:
        try:
            raw = base64.b64decode(b64)
        except Exception as exc:  # noqa: BLE001
            raise ImageGenError("invalid b64_json in response: %s" % exc)
    elif first.get("url"):
        raw = _download(first["url"], download=download)
    else:
        raise ImageGenError(
            "image %s returned neither url nor b64_json "
            "(keys: %s)" % (verb, ", ".join(sorted(first.keys()))))
    if not raw:
        raise ImageGenError("downloaded image is empty")
    fmt, width, height = _detect_format(raw)
    workdir = os.path.abspath(cwd or os.getcwd())
    path = os.path.join(workdir, name)
    with open(path, "wb") as fh:
        fh.write(raw)
    result = {"path": path, "filename": name, "format": fmt,
              "bytes": len(raw), "model": use_model, "size": size,
              "revised_prompt": first.get("revised_prompt"),
              "usage": data.get("usage")}
    if width is not None:
        result["width"] = width
        result["height"] = height
    return result


def generate(prompt, output_filename=None, size="1024x1024",
             model=None, request=None, download=None, cwd=None):
    """Generate one image and save it to the current working
    directory (the task worktree when run as a Codex MCP server).

    ``output_filename`` must be a plain basename (no separators,
    no ``..``); default is ``generated-<epoch>.png``. Returns a
    dict with path, format, dimensions (PNG only), byte size and
    the relay's usage block. Raises ImageGenError on any failure;
    the message never contains the API key.
    """
    key, name = _validate_common(prompt, size, output_filename)
    use_model = model or _default_model()
    status, data = _post_json(
        _base_url() + "/v1/images/generations", key,
        {"model": use_model, "prompt": str(prompt), "size": size},
        request=request)
    return _save_image_response(status, data, name, use_model, size,
                                download, cwd, "generation")


def _content_type_for(filename):
    lower = filename.lower()
    if lower.endswith(".jpg") or lower.endswith(".jpeg"):
        return "image/jpeg"
    if lower.endswith(".webp"):
        return "image/webp"
    return "image/png"


def _build_multipart(fields, file_field, filename, file_bytes):
    """Build a multipart/form-data body. Returns (body, boundary).
    Pure function for testability: the real post uses it too."""
    import uuid

    boundary = "----codeximagegen" + uuid.uuid4().hex
    body = b""
    for name, value in fields:
        body += ("--" + boundary + "\r\n").encode()
        body += ('Content-Disposition: form-data; '
                 'name="%s"\r\n\r\n' % name).encode()
        body += value.encode() + b"\r\n"
    body += ("--" + boundary + "\r\n").encode()
    body += ('Content-Disposition: form-data; name="%s"; '
             'filename="%s"\r\n' % (file_field, filename)).encode()
    body += ("Content-Type: %s\r\n\r\n"
             % _content_type_for(filename)).encode()
    body += file_bytes + b"\r\n"
    body += ("--" + boundary + "--\r\n").encode()
    return body, boundary


def _post_multipart(url, key, body, content_type,
                    request_multipart=None):
    if request_multipart is not None:
        return request_multipart(url, key, body, content_type)
    req = urllib.request.Request(url, data=body, method="POST", headers={
        "Authorization": "Bearer " + key,
        "Content-Type": content_type})
    try:
        with urllib.request.urlopen(req, timeout=API_TIMEOUT) as resp:
            return resp.status, json.loads(resp.read().decode() or "{}")
    except urllib.error.HTTPError as exc:
        raw = exc.read().decode()
        try:
            payload = json.loads(raw or "{}")
        except ValueError:
            payload = {"raw": raw[:400]}
        return exc.code, payload
    except Exception as exc:  # noqa: BLE001 - surfaced as tool error
        return 0, {"error": str(exc)[:300]}


MAX_INPUT_BYTES = 20 * 1024 * 1024


def edit(prompt, input_filename, output_filename=None,
         size="1024x1024", model=None, request_multipart=None,
         download=None, cwd=None):
    """Edit an existing image in the working directory and save
    the result as a NEW file (the input is never overwritten).

    ``input_filename`` and ``output_filename`` must both be plain
    basenames in the working directory. The relay's
    ``/v1/images/edits`` endpoint is multipart/form-data (verified
    2026-10-30: 200, URL PNG 1024x1024, input image tokens billed
    in usage). Returns the same dict shape as ``generate``.
    """
    key, name = _validate_common(prompt, size, output_filename)
    if (os.path.basename(input_filename or "") != input_filename
            or input_filename in (".", "..")
            or "/" in (input_filename or "")
            or "\\" in (input_filename or "")):
        raise ImageGenError(
            "input_filename must be a plain basename, got %r"
            % input_filename)
    workdir = os.path.abspath(cwd or os.getcwd())
    input_path = os.path.join(workdir, input_filename)
    if not os.path.isfile(input_path):
        raise ImageGenError("input image not found: %r" % input_filename)
    with open(input_path, "rb") as fh:
        file_bytes = fh.read()
    if len(file_bytes) > MAX_INPUT_BYTES:
        raise ImageGenError(
            "input image too large: %d bytes (max %d)"
            % (len(file_bytes), MAX_INPUT_BYTES))
    if os.path.abspath(os.path.join(workdir, name)) == input_path:
        raise ImageGenError(
            "output_filename must differ from input_filename")
    use_model = model or _default_model()
    body, boundary = _build_multipart(
        [("model", use_model), ("prompt", str(prompt)),
         ("size", size)],
        "image", input_filename, file_bytes)
    status, data = _post_multipart(
        _base_url() + "/v1/images/edits", key, body,
        "multipart/form-data; boundary=" + boundary,
        request_multipart=request_multipart)
    result = _save_image_response(status, data, name, use_model, size,
                                  download, cwd, "edit")
    result["input_filename"] = input_filename
    return result


# ------------------------------------------------------- codex-home glue

def ensure_codex_home_imagegen(codex_home_dir, command_args=None,
                               script_dir=None):
    """Seed imagegen MCP files + config into a tenant codex-home.

    Mirrors ``ensure_codex_home_websearch``: copies ``imagegen.py``
    and ``imagegen_mcp.py`` only when missing or identical (a
    tenant-modified script is never silently overwritten, it is
    reported in ``script_differs``), and appends the
    ``[mcp_servers.imagegen]`` section only when absent. Keys are
    never written here -- they arrive via
    ``ensure_imagegen_mcp_env`` from worker env.
    """
    import shutil

    home = os.path.abspath(codex_home_dir)
    os.makedirs(home, exist_ok=True)
    here = os.path.abspath(script_dir or os.path.dirname(__file__))
    script_changed = False
    script_differs = []
    for name in ("imagegen.py", "imagegen_mcp.py"):
        src = os.path.join(here, name)
        dst = os.path.join(home, name)
        if not os.path.exists(src):
            continue
        if not os.path.exists(dst):
            shutil.copyfile(src, dst)
            script_changed = True
        else:
            with open(src, "rb") as fh:
                src_bytes = fh.read()
            with open(dst, "rb") as fh:
                dst_bytes = fh.read()
            if src_bytes != dst_bytes:
                script_differs.append(name)
                continue
    config_path = os.path.join(home, "config.toml")
    if command_args is None:
        command_args = ["python3",
                        os.path.join(home, "imagegen_mcp.py")]
    section = ("\n[mcp_servers.imagegen]\n"
               "command = %s\n"
               "args = [%s]\n"
               % (json.dumps(command_args[0]),
                  ", ".join(json.dumps(a) for a in command_args[1:])))
    config_changed = False
    existing = ""
    if os.path.exists(config_path):
        with open(config_path, encoding="utf-8") as fh:
            existing = fh.read()
    if "[mcp_servers.imagegen]" not in existing:
        with open(config_path, "a", encoding="utf-8") as fh:
            if existing and not existing.endswith("\n"):
                fh.write("\n")
            fh.write(section)
        config_changed = True
    return {"config_path": config_path,
            "script_path": os.path.join(home, "imagegen_mcp.py"),
            "config_changed": config_changed,
            "script_changed": script_changed,
            "script_differs": script_differs}


def ensure_imagegen_mcp_env(codex_home_dir, env):
    """Write [mcp_servers.imagegen.env] into a tenant config.toml.

    Same rationale as ``ensure_websearch_mcp_env``: Codex's MCP
    subprocess does not inherit the Codex process env, so the env
    must live in config.toml. Idempotent in-place replacement of
    the env subsection; other sections untouched. Returns True
    when the file was written.
    """
    if not env:
        return False
    home = os.path.abspath(codex_home_dir)
    config_path = os.path.join(home, "config.toml")
    try:
        existing = ""
        if os.path.exists(config_path):
            with open(config_path, encoding="utf-8") as fh:
                existing = fh.read()
        if "[mcp_servers.imagegen]" not in existing:
            existing += ("\n[mcp_servers.imagegen]\n"
                         "command = \"python3\"\n"
                         "args = [\"/var/lib/codex-app-server/"
                         "codex-home/imagegen_mcp.py\"]\n")
        env_lines = ["[mcp_servers.imagegen.env]"]
        for key in sorted(env):
            env_lines.append("%s = %s" % (key, json.dumps(env[key])))
        env_block = "\n".join(env_lines) + "\n"
        lines = existing.splitlines(keepends=True)
        start = end = None
        for i, line in enumerate(lines):
            if line.strip() == "[mcp_servers.imagegen.env]":
                start = i
            elif start is not None and line.startswith("["):
                end = i
                break
        if start is not None:
            if end is None:
                end = len(lines)
            lines[start:end] = [env_block]
            new_text = "".join(lines)
        else:
            insert_at = None
            in_section = False
            for i, line in enumerate(lines):
                if line.strip() == "[mcp_servers.imagegen]":
                    in_section = True
                elif in_section and line.startswith("["):
                    insert_at = i
                    break
            if insert_at is None:
                insert_at = len(lines)
            lines.insert(insert_at, env_block)
            new_text = "".join(lines)
        if new_text == existing:
            return False
        with open(config_path, "w", encoding="utf-8") as fh:
            fh.write(new_text)
        return True
    except Exception:
        return False


def build_mcp_env(worker_environ, proxy_env=None):
    """Build the [mcp_servers.imagegen.env] dict from worker env.

    The MCP process only sees config.toml env, never the worker's
    own env, so IMAGE_GEN_BASE_URL must be resolved HERE (worker
    IMAGE_GEN_BASE_URL, else SUB2API_BASE_URL, else absent and
    imagegen.py falls back to the public relay domain -- which is
    NOT on the Tinyproxy whitelist, so a missing base here caused
    a 403 Filtered in the first production verification, 2026-10-30).
    proxy_env supplies HTTP(S)_PROXY for downloading the returned
    image URL; NO_PROXY gains the gateway host for direct API calls.
    Never injects the task user's gateway key.
    """
    env = dict(proxy_env) if proxy_env else {}
    for key in ("IMAGE_GEN_BASE_URL", "IMAGE_GEN_API_KEY",
                "IMAGE_GEN_MODEL"):
        if worker_environ.get(key):
            env[key] = worker_environ[key]
    base = (worker_environ.get("IMAGE_GEN_BASE_URL", "")
            or worker_environ.get("SUB2API_BASE_URL", ""))
    if base:
        env["IMAGE_GEN_BASE_URL"] = base
    host = base.split("://", 1)[-1].split("/", 1)[0]
    host = host.split(":", 1)[0]
    parts = [p for p in env.get("NO_PROXY", "").split(",") if p]
    for item in ("localhost", "127.0.0.1", host):
        if item and item not in parts:
            parts.append(item)
    env["NO_PROXY"] = ",".join(parts)
    env["no_proxy"] = ",".join(parts)
    return env