from __future__ import annotations

import argparse
import base64
import binascii
from dataclasses import dataclass
import json
import mimetypes
import os
import re
import socket
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any

from runtime_config import (
    ConfigError,
    DEFAULT_BASE_URL,
    DEFAULT_MODEL,
    ROOT_DIR,
    build_api_url,
    infer_request_style,
    load_runtime_config,
    normalize_base_url,
    normalize_token,
    parse_env_file,
    runtime_config_summary,
    save_base_url as save_runtime_base_url,
    save_token as save_runtime_token,
)


SUPPORTED_ASPECT_RATIOS = (
    "1:1",
    "16:9",
    "9:16",
    "4:3",
    "3:4",
    "3:2",
    "2:3",
    "5:4",
    "4:5",
    "2:1",
    "1:2",
    "21:9",
    "9:21",
)
PIXEL_SIZE_PATTERN = re.compile(r"^(?P<width>\d{2,5})x(?P<height>\d{2,5})$")
OPENAI_RATIO_SIZE_MAP = {
    "1:1": "1024x1024",
    "16:9": "1536x864",
    "9:16": "864x1536",
    "4:3": "1536x1152",
    "3:4": "1152x1536",
    "3:2": "1536x1024",
    "2:3": "1024x1536",
    "5:4": "1280x1024",
    "4:5": "1024x1280",
    "2:1": "1536x768",
    "1:2": "768x1536",
    "21:9": "1680x720",
    "9:21": "720x1680",
}
SIZE_ENV_VARS = (
    "GPT_IMAGE_2_SIZE",
    "IMAGE_SIZE",
    "OPENAI_IMAGE_SIZE",
)
EDIT_SIZE_ENV_VARS = (
    "GPT_IMAGE_2_EDIT_SIZE",
    "IMAGE_EDIT_SIZE",
    "OPENAI_IMAGE_EDIT_SIZE",
)
EDIT_QUALITY_ENV_VARS = (
    "GPT_IMAGE_2_EDIT_QUALITY",
    "IMAGE_EDIT_QUALITY",
    "OPENAI_IMAGE_EDIT_QUALITY",
)
EDIT_BACKGROUND_ENV_VARS = (
    "GPT_IMAGE_2_EDIT_BACKGROUND",
    "IMAGE_EDIT_BACKGROUND",
    "OPENAI_IMAGE_EDIT_BACKGROUND",
)
EDIT_OUTPUT_FORMAT_ENV_VARS = (
    "GPT_IMAGE_2_EDIT_OUTPUT_FORMAT",
    "IMAGE_EDIT_OUTPUT_FORMAT",
    "OPENAI_IMAGE_EDIT_OUTPUT_FORMAT",
)
EDIT_INPUT_FIDELITY_ENV_VARS = (
    "GPT_IMAGE_2_EDIT_INPUT_FIDELITY",
    "IMAGE_EDIT_INPUT_FIDELITY",
    "OPENAI_IMAGE_EDIT_INPUT_FIDELITY",
)
REQUEST_STYLE_CHOICES = ("auto", "apicodex", "openai")
REQUEST_RETRY_ATTEMPTS = max(1, int(os.environ.get("IMAGE_RETRY_ATTEMPTS", "2")))
REQUEST_RETRY_DELAY_SECONDS = max(0.0, float(os.environ.get("IMAGE_RETRY_DELAY_SECONDS", "2")))


class ImageGenerationError(RuntimeError):
    pass


class ImageRequestError(ImageGenerationError):
    def __init__(self, message: str, *, status: int | None = None, timeout: bool = False):
        super().__init__(message)
        self.status = status
        self.timeout = timeout


@dataclass(frozen=True)
class ImageInputFile:
    filename: str
    mime_type: str
    data: bytes


@dataclass(frozen=True)
class RequestTarget:
    name: str
    base_url: str
    token: str
    request_style: str
    source: str


def normalize_size_value(raw_size: str) -> str:
    value = raw_size.strip().lower().replace(" ", "")
    match = PIXEL_SIZE_PATTERN.fullmatch(value)
    if match:
        width = int(match.group("width"))
        height = int(match.group("height"))
        if width < 64 or height < 64:
            raise ImageGenerationError(f"Unsupported size: {raw_size}")
        return f"{width}x{height}"
    return value


def is_supported_size(size: str) -> bool:
    return size in SUPPORTED_ASPECT_RATIOS or PIXEL_SIZE_PATTERN.fullmatch(size) is not None


def resolve_api_size(requested_size: str, request_style: str) -> str:
    if request_style == "openai" and requested_size in OPENAI_RATIO_SIZE_MAP:
        return OPENAI_RATIO_SIZE_MAP[requested_size]
    return requested_size


def lookup_env_setting(env_map: dict[str, str], env_vars: tuple[str, ...]) -> str | None:
    for env_name in env_vars:
        value = os.environ.get(env_name, "").strip()
        if value:
            return value
    for env_name in env_vars:
        value = env_map.get(env_name, "").strip()
        if value:
            return value
    return None


def resolve_requested_size(
    size: str | None,
    request_style: str,
    has_refs: bool,
    env_map: dict[str, str],
) -> tuple[str, str]:
    if size and size.strip():
        requested_size = normalize_size_value(size)
    else:
        env_vars = EDIT_SIZE_ENV_VARS if has_refs else SIZE_ENV_VARS
        env_size = lookup_env_setting(env_map, env_vars)
        requested_size = normalize_size_value(env_size) if env_size else "1:1"

    if not is_supported_size(requested_size):
        raise ImageGenerationError(f"Unsupported size: {requested_size}")
    return requested_size, resolve_api_size(requested_size, request_style)


def resolve_optional_edit_field(
    explicit_value: str | None,
    env_map: dict[str, str],
    env_vars: tuple[str, ...],
) -> str | None:
    if explicit_value is not None and explicit_value.strip():
        return explicit_value.strip()
    value = lookup_env_setting(env_map, env_vars)
    return value.strip() if value else None


def build_env_template() -> str:
    return "\n".join(
        [
            "# Required: apicodex-compatible image API token",
            "IMAGE_API_KEY=your_apicodex_token_here",
            "",
            "# Optional: host root or full /v1 path are both accepted",
            f"IMAGE_BASE_URL={DEFAULT_BASE_URL}/v1",
            "",
            "# Optional: defaults shown below",
            f"IMAGE_MODEL={DEFAULT_MODEL}",
            "IMAGE_REQUEST_STYLE=auto",
            "IMAGE_OUTPUT_DIR=./output",
            "",
        ]
    )


def write_env_template(
    *,
    env_file: str | Path | None = None,
    overwrite: bool = False,
) -> Path:
    execution_root = Path.cwd()
    if env_file:
        target_path = Path(env_file).expanduser()
        if not target_path.is_absolute():
            target_path = (execution_root / target_path).resolve()
    else:
        target_path = execution_root / ".env.example"

    if target_path.exists() and not overwrite:
        raise ConfigError(f"Config template already exists: {target_path}. Use --force to overwrite.")

    target_path.parent.mkdir(parents=True, exist_ok=True)
    target_path.write_text(build_env_template(), encoding="utf-8")
    return target_path

def read_json_response(request: urllib.request.Request, timeout: int) -> dict[str, Any]:
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            body = response.read().decode("utf-8")
    except urllib.error.HTTPError as error:
        body = error.read().decode("utf-8", errors="replace")
        raise ImageRequestError(f"HTTP {error.code}: {body}", status=error.code) from error
    except urllib.error.URLError as error:
        if is_timeout_error(error):
            raise ImageRequestError(f"Timeout: {error}", timeout=True) from error
        raise ImageRequestError(f"Network error: {error}") from error
    except (TimeoutError, socket.timeout) as error:
        raise ImageRequestError(f"Timeout: {error}", timeout=True) from error

    if not body:
        return {}
    try:
        return json.loads(body)
    except json.JSONDecodeError as error:
        raise ImageGenerationError(f"API returned non-JSON response: {body[:500]}") from error


def request_json(method: str, url: str, token: str, payload: dict[str, Any] | None = None, timeout: int = 300) -> dict[str, Any]:
    data = None
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
        "User-Agent": "gpt-image-2-generator/1.0",
    }
    if payload is not None:
        data = json.dumps(payload, ensure_ascii=False).encode("utf-8")

    request = urllib.request.Request(url, data=data, headers=headers, method=method)
    return read_json_response(request, timeout)


def request_multipart_json(
    url: str,
    token: str,
    fields: dict[str, Any],
    files: list[ImageInputFile],
    image_field_name: str = "image[]",
    timeout: int = 300,
) -> dict[str, Any]:
    boundary = f"----gpt-image-2-{uuid.uuid4().hex}"
    body = bytearray()

    def add(value: str | bytes) -> None:
        if isinstance(value, str):
            body.extend(value.encode("utf-8"))
        else:
            body.extend(value)

    for name, value in fields.items():
        add(f"--{boundary}\r\n")
        add(f'Content-Disposition: form-data; name="{name}"\r\n\r\n')
        add(str(value))
        add("\r\n")

    for image_file in files:
        add(f"--{boundary}\r\n")
        add(
            f'Content-Disposition: form-data; name="{image_field_name}"; '
            f'filename="{image_file.filename}"\r\n'
        )
        add(f"Content-Type: {image_file.mime_type}\r\n\r\n")
        add(image_file.data)
        add("\r\n")

    add(f"--{boundary}--\r\n")
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": f"multipart/form-data; boundary={boundary}",
        "Content-Length": str(len(body)),
        "User-Agent": "gpt-image-2-generator/1.0",
    }
    request = urllib.request.Request(url, data=bytes(body), headers=headers, method="POST")
    return read_json_response(request, timeout)


def is_timeout_error(error: urllib.error.URLError) -> bool:
    reason = getattr(error, "reason", None)
    if isinstance(reason, (TimeoutError, socket.timeout)):
        return True
    return "timed out" in str(reason).lower() or "timeout" in str(error).lower()


def is_failover_eligible_error(error: BaseException) -> bool:
    if isinstance(error, ImageRequestError):
        if error.status in {502, 503, 504, 524}:
            return True
        if error.timeout:
            return True
    message = str(error).lower()
    return "timed out" in message or "timeout" in message


def read_first_config_line(path: Path) -> str:
    if not path.exists():
        return ""
    for raw_line in path.read_text(encoding="utf-8-sig").splitlines():
        line = raw_line.strip()
        if line and not line.startswith("#"):
            return line
    return ""


def load_skill_fallback_target(runtime_config: Any) -> RequestTarget | None:
    skill_config_dir = ROOT_DIR / "config"
    base_url = normalize_base_url(read_first_config_line(skill_config_dir / "base_url.txt"))
    token = normalize_token(read_first_config_line(skill_config_dir / "token.txt"))
    request_style_raw = read_first_config_line(skill_config_dir / "request_style.txt").strip().lower()

    if not base_url or not token:
        return None

    request_style = request_style_raw if request_style_raw in {"apicodex", "openai"} else infer_request_style(base_url)
    primary_token = normalize_token(runtime_config.token or "")
    if runtime_config.base_url == base_url and primary_token == token:
        return None

    return RequestTarget(
        name="skill-fallback",
        base_url=base_url,
        token=token,
        request_style=request_style,
        source=str(skill_config_dir),
    )


def build_request_targets(runtime_config: Any, request_style: str) -> list[RequestTarget]:
    targets = [
        RequestTarget(
            name="primary",
            base_url=runtime_config.base_url,
            token=runtime_config.token or "",
            request_style=request_style,
            source=runtime_config.base_url_source,
        )
    ]
    fallback_target = load_skill_fallback_target(runtime_config)
    if fallback_target is not None:
        targets.append(fallback_target)
    return targets


def format_request_error(error: BaseException) -> str:
    if isinstance(error, ImageRequestError):
        if error.status is not None:
            return f"HTTP {error.status}"
        if error.timeout:
            return "timeout"
    return error.__class__.__name__


def should_retry_request(error: BaseException) -> bool:
    if isinstance(error, ImageRequestError):
        if error.status in {429, 500, 502, 503, 504, 524}:
            return True
        if error.timeout:
            return True
    message = str(error).lower()
    return "timed out" in message or "timeout" in message


def submit_request_for_target(
    target: RequestTarget,
    *,
    runtime_config: Any,
    env_map: dict[str, str],
    prompt: str,
    size: str | None,
    n: int,
    has_refs: bool,
    image_files: list[ImageInputFile],
    timeout: int,
    quality: str | None,
    background: str | None,
    output_format: str | None,
    input_fidelity: str | None,
) -> tuple[dict[str, Any], str, str, str]:
    generations_api_url = build_api_url(target.base_url, "images/generations")
    edits_api_url = build_api_url(target.base_url, "images/edits")
    payload: dict[str, Any] = {
        "model": runtime_config.model,
        "prompt": prompt,
    }
    requested_size, api_size = resolve_requested_size(size, target.request_style, has_refs, env_map)
    payload["size"] = api_size
    if target.request_style == "apicodex" or n != 1:
        payload["n"] = n

    if has_refs:
        edit_quality = resolve_optional_edit_field(quality, env_map, EDIT_QUALITY_ENV_VARS)
        edit_background = resolve_optional_edit_field(background, env_map, EDIT_BACKGROUND_ENV_VARS)
        edit_output_format = resolve_optional_edit_field(output_format, env_map, EDIT_OUTPUT_FORMAT_ENV_VARS)
        edit_input_fidelity = resolve_optional_edit_field(input_fidelity, env_map, EDIT_INPUT_FIDELITY_ENV_VARS)
        if edit_quality:
            payload["quality"] = edit_quality
        if edit_background:
            payload["background"] = edit_background
        if edit_output_format:
            payload["output_format"] = edit_output_format
        if edit_input_fidelity:
            payload["input_fidelity"] = edit_input_fidelity
        image_field_name = "image[]" if target.request_style == "apicodex" or len(image_files) != 1 else "image"
        response = request_multipart_json(
            edits_api_url,
            target.token,
            payload,
            image_files,
            image_field_name=image_field_name,
            timeout=timeout,
        )
    else:
        response = request_json("POST", generations_api_url, target.token, payload, timeout=timeout)

    return response, target.request_style, requested_size, api_size


def safe_filename(filename: str, fallback: str, mime_type: str) -> str:
    candidate = re.sub(r"[^A-Za-z0-9._-]+", "_", Path(filename).name).strip("._")
    if not candidate:
        candidate = fallback
    if not Path(candidate).suffix:
        candidate += mimetypes.guess_extension(mime_type) or ".png"
    return candidate


def read_data_uri_image(value: str, index: int) -> ImageInputFile:
    match = re.match(r"^data:(?P<mime>image/[^;,]+)(?P<params>(?:;[^,]+)*),(?P<data>.+)$", value, flags=re.DOTALL)
    if not match or ";base64" not in match.group("params").lower():
        raise ImageGenerationError("Reference image data URI must be a base64 image")

    mime_type = match.group("mime")
    try:
        data = base64.b64decode(match.group("data"), validate=True)
    except (ValueError, binascii.Error) as error:
        raise ImageGenerationError("Reference image data URI contains invalid base64") from error
    return ImageInputFile(
        filename=safe_filename("", f"reference-{index}", mime_type),
        mime_type=mime_type,
        data=data,
    )


def download_reference_image(url: str, index: int, timeout: int) -> ImageInputFile:
    request = urllib.request.Request(url, headers={"User-Agent": "gpt-image-2-generator/1.0"})
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            data = response.read()
            content_type = response.headers.get("Content-Type")
    except urllib.error.HTTPError as error:
        body = error.read().decode("utf-8", errors="replace")
        raise ImageGenerationError(f"Reference image URL returned HTTP {error.code}: {body[:500]}") from error
    except urllib.error.URLError as error:
        raise ImageGenerationError(f"Reference image URL could not be downloaded: {error}") from error

    mime_type = (content_type or "").split(";")[0].strip().lower()
    if not mime_type.startswith("image/"):
        mime_type = mimetypes.guess_type(urllib.parse.urlparse(url).path)[0] or "image/png"
    filename = safe_filename(Path(urllib.parse.urlparse(url).path).name, f"reference-{index}", mime_type)
    return ImageInputFile(filename=filename, mime_type=mime_type, data=data)


def read_local_reference_image(image_ref: str, index: int) -> ImageInputFile:
    value = image_ref.strip()
    path = Path(value).expanduser()
    if not path.is_absolute():
        cwd_path = Path.cwd() / path
        root_path = ROOT_DIR / path
        path = cwd_path if cwd_path.exists() else root_path

    if not path.exists():
        raise ImageGenerationError(f"Reference image not found: {image_ref}")

    mime_type = mimetypes.guess_type(str(path))[0] or "image/png"
    return ImageInputFile(
        filename=safe_filename(path.name, f"reference-{index}", mime_type),
        mime_type=mime_type,
        data=path.read_bytes(),
    )


def image_ref_to_input_file(image_ref: str, index: int, timeout: int) -> ImageInputFile:
    value = image_ref.strip()
    lowered = value.lower()
    if lowered.startswith("data:"):
        return read_data_uri_image(value, index)
    if lowered.startswith(("http://", "https://")):
        return download_reference_image(value, index, timeout)
    return read_local_reference_image(value, index)


def save_json(path: Path, data: dict[str, Any]) -> None:
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def unique_path(path: Path) -> Path:
    if not path.exists():
        return path
    for index in range(2, 1000):
        candidate = path.with_name(f"{path.stem}-{index}{path.suffix}")
        if not candidate.exists():
            return candidate
    raise ImageGenerationError(f"Cannot create unique output path for {path}")


def save_b64_image(image_b64: str, output_dir: Path, stem: str, index: int, mime: str | None = None) -> Path:
    extension = mimetypes.guess_extension(mime or "image/png") or ".png"
    output_path = unique_path(output_dir / f"{stem}-{index}{extension}")
    output_path.write_bytes(base64.b64decode(image_b64))
    return output_path


def save_data_uri(data_uri: str, output_dir: Path, stem: str, index: int) -> Path:
    match = re.match(r"^data:(?P<mime>[^;]+);base64,(?P<data>.+)$", data_uri, flags=re.DOTALL)
    if not match:
        raise ImageGenerationError("Unsupported data URI image format")
    return save_b64_image(match.group("data"), output_dir, stem, index, match.group("mime"))


def extract_b64_items(data: Any) -> list[dict[str, Any]]:
    items: list[dict[str, Any]] = []

    def visit(value: Any) -> None:
        if isinstance(value, dict):
            if isinstance(value.get("b64_json"), str):
                items.append(value)
                return
            for child in value.values():
                visit(child)
        elif isinstance(value, list):
            for child in value:
                visit(child)

    visit(data)
    return items


def extract_image_values(data: Any) -> list[str]:
    values: list[str] = []

    def add(value: Any) -> None:
        if isinstance(value, str) and value.startswith(("http://", "https://", "data:image/")):
            values.append(value)
        elif isinstance(value, list):
            for item in value:
                add(item)

    candidates = []
    if isinstance(data, dict):
        candidates.append(data.get("url"))
        nested_data = data.get("data")
        if isinstance(nested_data, list):
            candidates.extend(item.get("url") for item in nested_data if isinstance(item, dict))
        if isinstance(nested_data, dict):
            result = nested_data.get("result")
            if isinstance(result, dict):
                images = result.get("images")
                if isinstance(images, list):
                    candidates.extend(item.get("url") for item in images if isinstance(item, dict))
                candidates.append(result.get("url"))
        result = data.get("result")
        if isinstance(result, dict):
            images = result.get("images")
            if isinstance(images, list):
                candidates.extend(item.get("url") for item in images if isinstance(item, dict))

    for candidate in candidates:
        add(candidate)
    return values


def extract_task_id(data: dict[str, Any]) -> str | None:
    for key in ("task_id", "id"):
        value = data.get(key)
        if isinstance(value, str) and value:
            return value

    nested_data = data.get("data")
    if isinstance(nested_data, dict):
        for key in ("task_id", "id"):
            value = nested_data.get(key)
            if isinstance(value, str) and value:
                return value
    return None


def get_task_status(task: dict[str, Any]) -> str | None:
    nested_data = task.get("data")
    if isinstance(nested_data, dict) and isinstance(nested_data.get("status"), str):
        return nested_data["status"]
    status = task.get("status")
    return status if isinstance(status, str) else None


def get_task_error(task: dict[str, Any]) -> str:
    nested_data = task.get("data")
    error = nested_data.get("error") if isinstance(nested_data, dict) else task.get("error")
    if isinstance(error, dict):
        return str(error.get("message") or error)
    if error:
        return str(error)
    return "task failed"


def guess_download_extension(url: str, content_type: str | None) -> str:
    suffix = Path(urllib.parse.urlparse(url).path).suffix.lower()
    if suffix in {".png", ".jpg", ".jpeg", ".webp", ".gif"}:
        return suffix
    if content_type:
        return mimetypes.guess_extension(content_type.split(";")[0].strip()) or ".png"
    return ".png"


def download_image(url: str, output_dir: Path, stem: str, index: int, timeout: int) -> Path:
    request = urllib.request.Request(url, headers={"User-Agent": "gpt-image-2-generator/1.0"})
    with urllib.request.urlopen(request, timeout=timeout) as response:
        content_type = response.headers.get("Content-Type")
        extension = guess_download_extension(url, content_type)
        output_path = unique_path(output_dir / f"{stem}-{index}{extension}")
        output_path.write_bytes(response.read())
    return output_path


def save_image_values(values: list[str], output_dir: Path, stem: str, timeout: int) -> list[Path]:
    paths: list[Path] = []
    for index, value in enumerate(values, start=1):
        if value.startswith("data:image/"):
            paths.append(save_data_uri(value, output_dir, stem, index))
        else:
            paths.append(download_image(value, output_dir, stem, index, timeout))
    return paths


def generate_image(
    prompt: str,
    size: str | None = None,
    image_refs: list[str] | None = None,
    n: int = 1,
    output_dir: str | Path | None = None,
    timeout: int = 300,
    initial_poll_delay: float = 10,
    poll_interval: float = 5,
    max_wait: float = 300,
    config_dir: str | Path | None = None,
    env_file: str | Path | None = None,
    base_url: str | None = None,
    model: str | None = None,
    request_style: str | None = None,
    quality: str | None = None,
    background: str | None = None,
    output_format: str | None = None,
    input_fidelity: str | None = None,
) -> dict[str, Any]:
    prompt = prompt.strip()
    if not prompt:
        raise ImageGenerationError("Prompt is required")
    if n < 1:
        raise ImageGenerationError("n must be >= 1")

    runtime_config = load_runtime_config(
        execution_root=Path.cwd(),
        config_dir=config_dir,
        env_file=env_file,
        base_url_override=base_url,
        model_override=model,
        output_dir_override=output_dir,
        require_token=True,
    )
    env_map = parse_env_file(runtime_config.env_file)
    resolved_request_style = request_style.strip().lower() if request_style else runtime_config.request_style
    if resolved_request_style == "auto":
        resolved_request_style = runtime_config.request_style
    if resolved_request_style not in {"apicodex", "openai"}:
        raise ImageGenerationError(
            f"Unsupported request style: {resolved_request_style}. Use one of: {', '.join(REQUEST_STYLE_CHOICES)}"
        )

    refs = [ref.strip() for ref in (image_refs or []) if ref.strip()]
    output_path = runtime_config.output_dir
    output_path.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    stem = f"gpt-image-2-{timestamp}"
    request_type = "edit" if refs else "generation"
    image_files = [image_ref_to_input_file(ref, index, timeout) for index, ref in enumerate(refs, start=1)] if refs else []
    targets = build_request_targets(runtime_config, resolved_request_style)
    response: dict[str, Any] | None = None
    selected_target: RequestTarget | None = None
    requested_size = ""
    api_size = ""
    last_error: BaseException | None = None
    for target_index, target in enumerate(targets):
        for attempt in range(1, REQUEST_RETRY_ATTEMPTS + 1):
            try:
                response, resolved_request_style, requested_size, api_size = submit_request_for_target(
                    target,
                    runtime_config=runtime_config,
                    env_map=env_map,
                    prompt=prompt,
                    size=size,
                    n=n,
                    has_refs=bool(refs),
                    image_files=image_files,
                    timeout=timeout,
                    quality=quality,
                    background=background,
                    output_format=output_format,
                    input_fidelity=input_fidelity,
                )
                selected_target = target
                break
            except BaseException as error:
                last_error = error
                if attempt < REQUEST_RETRY_ATTEMPTS and should_retry_request(error):
                    print(
                        f"request attempt {attempt}/{REQUEST_RETRY_ATTEMPTS} to {target.base_url} failed with {format_request_error(error)}; retrying in {REQUEST_RETRY_DELAY_SECONDS:g}s",
                        file=sys.stderr,
                    )
                    time.sleep(REQUEST_RETRY_DELAY_SECONDS)
                    continue
                next_target = targets[target_index + 1] if target_index + 1 < len(targets) else None
                if next_target is not None and is_failover_eligible_error(error):
                    print(
                        f"request to {target.base_url} failed with {format_request_error(error)}; switching to fallback {next_target.base_url}",
                        file=sys.stderr,
                    )
                    break
                raise
        if response is not None and selected_target is not None:
            break
    if response is None or selected_target is None:
        raise ImageGenerationError(f"Image request failed: {last_error}")

    task_api_base = build_api_url(selected_target.base_url, "tasks/")
    response_path = output_path / f"{stem}-response.json"
    save_json(response_path, response)

    image_paths: list[Path] = []
    revised_prompts: list[str] = []

    b64_items = extract_b64_items(response)
    if b64_items:
        for index, item in enumerate(b64_items, start=1):
            image_paths.append(save_b64_image(item["b64_json"], output_path, stem, index))
            if isinstance(item.get("revised_prompt"), str):
                revised_prompts.append(item["revised_prompt"])
        return {
            "images": [str(path) for path in image_paths],
            "response_path": str(response_path),
            "revised_prompts": revised_prompts,
            "request_type": request_type,
            "mode": "sync-b64",
            "base_url": selected_target.base_url,
            "model": runtime_config.model,
            "request_style": resolved_request_style,
            "requested_size": requested_size,
            "api_size": api_size,
            "output_dir": str(output_path),
        }

    image_values = extract_image_values(response)
    if image_values:
        image_paths = save_image_values(image_values, output_path, stem, timeout)
        return {
            "images": [str(path) for path in image_paths],
            "response_path": str(response_path),
            "revised_prompts": revised_prompts,
            "request_type": request_type,
            "mode": "sync-url",
            "base_url": selected_target.base_url,
            "model": runtime_config.model,
            "request_style": resolved_request_style,
            "requested_size": requested_size,
            "api_size": api_size,
            "output_dir": str(output_path),
        }

    task_id = extract_task_id(response)
    if not task_id:
        raise ImageGenerationError(f"Unexpected response. Saved raw response to {response_path}")

    time.sleep(initial_poll_delay)
    deadline = time.time() + max_wait
    last_task: dict[str, Any] | None = None
    while time.time() <= deadline:
        last_task = request_json(
            "GET",
            urllib.parse.urljoin(task_api_base, urllib.parse.quote(task_id)),
            runtime_config.token or "",
            timeout=timeout,
        )
        status = get_task_status(last_task)
        if status == "completed":
            task_path = output_path / f"{stem}-task.json"
            save_json(task_path, last_task)

            b64_items = extract_b64_items(last_task)
            if b64_items:
                for index, item in enumerate(b64_items, start=1):
                    image_paths.append(save_b64_image(item["b64_json"], output_path, stem, index))
                    if isinstance(item.get("revised_prompt"), str):
                        revised_prompts.append(item["revised_prompt"])
            else:
                image_paths = save_image_values(extract_image_values(last_task), output_path, stem, timeout)

            if not image_paths:
                raise ImageGenerationError(f"Task completed but no image was found. Saved task response to {task_path}")
            return {
                "images": [str(path) for path in image_paths],
                "response_path": str(response_path),
                "task_response_path": str(task_path),
                "revised_prompts": revised_prompts,
                "request_type": request_type,
                "mode": "async-task",
                "base_url": selected_target.base_url,
                "model": runtime_config.model,
                "request_style": resolved_request_style,
                "requested_size": requested_size,
                "api_size": api_size,
                "output_dir": str(output_path),
            }
        if status == "failed":
            raise ImageGenerationError(get_task_error(last_task))
        time.sleep(poll_interval)

    if last_task is not None:
        task_path = output_path / f"{stem}-task-timeout.json"
        save_json(task_path, last_task)
        raise ImageGenerationError(f"Task polling timed out. Saved last task response to {task_path}")
    raise ImageGenerationError("Task polling timed out before the first task response")


def save_token(token: str, config_dir: str | Path | None = None) -> Path:
    return save_runtime_token(token, execution_root=Path.cwd(), config_dir=config_dir)


def save_base_url(base_url: str, config_dir: str | Path | None = None) -> Path:
    return save_runtime_base_url(base_url, execution_root=Path.cwd(), config_dir=config_dir)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Generate or edit images with GPT Image 2 via apicodex-style or OpenAI-compatible APIs")
    parser.add_argument("--save-token", metavar="TOKEN", help="Save an API token to <cwd>/config/token.txt and exit")
    parser.add_argument("--save-base-url", metavar="URL", help="Save an API base URL to <cwd>/config/base_url.txt and exit")
    parser.add_argument("--print-env-template", action="store_true", help="Print a recommended .env template and exit")
    parser.add_argument("--init-config", action="store_true", help="Write a recommended env template file and exit")
    parser.add_argument("--prompt", required=False, default=None, help="Text prompt")
    parser.add_argument("--size", default=None, help="Size such as 2:3 or 1024x1536. If omitted, tries env defaults before falling back to 1:1.")
    parser.add_argument("--n", default=1, type=int, help="Number of images")
    parser.add_argument("--image-ref", action="append", default=[], help="Reference image URL, data URI, or local path")
    parser.add_argument("--output-dir", default=None, help="Output directory. Defaults to <cwd>/outputs.")
    parser.add_argument("--timeout", default=300, type=int, help="HTTP timeout in seconds")
    parser.add_argument("--initial-poll-delay", default=10, type=float, help="Seconds to wait before polling async tasks")
    parser.add_argument("--poll-interval", default=5, type=float, help="Async task poll interval in seconds")
    parser.add_argument("--max-wait", default=300, type=float, help="Maximum async task wait time in seconds")
    parser.add_argument("--config-dir", help="Config directory. Defaults to <cwd>/config.")
    parser.add_argument("--env-file", help="Env file. Defaults to <cwd>/.env.")
    parser.add_argument("--base-url", help="Override the API base URL for this request.")
    parser.add_argument("--model", help="Override the model for this request.")
    parser.add_argument("--request-style", choices=REQUEST_STYLE_CHOICES, default=None, help="Force request style. Defaults to auto/inferred from base URL.")
    parser.add_argument("--quality", default=None, help="Optional image edit quality value.")
    parser.add_argument("--background", default=None, help="Optional image edit background value.")
    parser.add_argument("--output-format", default=None, help="Optional image edit output format value.")
    parser.add_argument("--input-fidelity", default=None, help="Optional image edit input fidelity value.")
    parser.add_argument("--show-config", action="store_true", help="Print the resolved runtime configuration and exit")
    parser.add_argument("--force", action="store_true", help="Allow overwriting files when used with --init-config")
    return parser


def main() -> int:
    args = build_parser().parse_args()

    if args.save_token:
        try:
            path = save_token(args.save_token, config_dir=args.config_dir)
            print(f"Token saved to {path}")
            return 0
        except (ConfigError, Exception) as error:
            print(f"ERROR: {error}", file=sys.stderr)
            return 1

    if args.save_base_url:
        try:
            path = save_base_url(args.save_base_url, config_dir=args.config_dir)
            print(f"Base URL saved to {path}")
            return 0
        except (ConfigError, Exception) as error:
            print(f"ERROR: {error}", file=sys.stderr)
            return 1

    if args.print_env_template:
        print(build_env_template(), end="")
        return 0

    if args.init_config:
        try:
            path = write_env_template(env_file=args.env_file or ".env.example", overwrite=args.force)
            print(f"Config template written to {path}")
            return 0
        except (ConfigError, Exception) as error:
            print(f"ERROR: {error}", file=sys.stderr)
            return 1

    if args.show_config:
        try:
            config = load_runtime_config(
                execution_root=Path.cwd(),
                config_dir=args.config_dir,
                env_file=args.env_file,
                base_url_override=args.base_url,
                model_override=args.model,
                output_dir_override=args.output_dir,
                require_token=False,
            )
            print(json.dumps(runtime_config_summary(config), ensure_ascii=False, indent=2))
            return 0
        except (ConfigError, Exception) as error:
            print(f"ERROR: {error}", file=sys.stderr)
            return 1

    if not args.prompt:
        print(
            "ERROR: --prompt is required (unless using --save-token / --save-base-url / --print-env-template / --init-config / --show-config)",
            file=sys.stderr,
        )
        return 1

    try:
        result = generate_image(
            prompt=args.prompt,
            size=args.size,
            image_refs=args.image_ref,
            n=args.n,
            output_dir=args.output_dir,
            timeout=args.timeout,
            initial_poll_delay=args.initial_poll_delay,
            poll_interval=args.poll_interval,
            max_wait=args.max_wait,
            config_dir=args.config_dir,
            env_file=args.env_file,
            base_url=args.base_url,
            model=args.model,
            request_style=args.request_style,
            quality=args.quality,
            background=args.background,
            output_format=args.output_format,
            input_fidelity=args.input_fidelity,
        )
    except (ConfigError, Exception) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        return 1

    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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