#!/usr/bin/env python3
"""Generate images through the managed Soul Store image gateway."""

from __future__ import annotations

import argparse
import base64
import binascii
import json
import mimetypes
import os
import re
import sys
import tempfile
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import Request, urlopen

MAX_IMAGE_BYTES = 50 * 1024 * 1024
MAX_GATEWAY_RESPONSE_BYTES = 100 * 1024 * 1024
MODEL_ENV = "SOULSTORE_IMAGE_GENERATION_MODEL"
TIMEOUT_ENV = "SOULSTORE_IMAGE_GENERATION_TIMEOUT_SECONDS"
RUNTIME_HOME_ENV = "XIGUA_RUNTIME_HOME_DIR"
WORKSPACE_ENV = "XIGUA_WORKSPACE_DIR"
DEFAULT_TIMEOUT_SECONDS = 180
SUPPORTED_IMAGE_SIZES = ("1024x1024", "1536x1024", "1024x1536")
SUPPORTED_IMAGE_ASPECT_RATIOS = {
    "1024x1024": 1.0,
    "1536x1024": 1.5,
    "1024x1536": 2 / 3,
}


class GenerateImageError(RuntimeError):
    pass


def _resolve_timeout_seconds() -> int:
    raw = os.environ.get(TIMEOUT_ENV, "").strip()
    if not raw:
        return DEFAULT_TIMEOUT_SECONDS
    try:
        timeout_seconds = int(raw)
    except ValueError as exc:
        raise GenerateImageError(f"{TIMEOUT_ENV} must be an integer") from exc
    if not 1 <= timeout_seconds <= 3600:
        raise GenerateImageError(f"{TIMEOUT_ENV} must be between 1 and 3600")
    return timeout_seconds


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    prompt_group = parser.add_mutually_exclusive_group(required=True)
    prompt_group.add_argument("--prompt", help="Complete image prompt")
    prompt_group.add_argument(
        "--prompt-file",
        help="UTF-8 prompt file, or '-' to read from standard input",
    )
    parser.add_argument("--size", help="Requested image size, for example 1024x1024")
    parser.add_argument("--quality", choices=("low", "medium", "high"))
    parser.add_argument("--count", type=int, default=1)
    parser.add_argument("--output-dir", default="generated-images")
    args = parser.parse_args()
    if not 1 <= args.count <= 4:
        parser.error("--count must be between 1 and 4")
    return args


def _read_prompt(args: argparse.Namespace) -> str:
    if args.prompt is not None:
        prompt = args.prompt
    elif args.prompt_file == "-":
        prompt = sys.stdin.read()
    else:
        try:
            prompt = Path(args.prompt_file).read_text(encoding="utf-8")
        except OSError as exc:
            raise GenerateImageError(f"cannot read prompt file: {exc}") from exc
    prompt = prompt.strip()
    if not prompt:
        raise GenerateImageError("image prompt cannot be empty")
    return prompt


def _resolve_image_size(raw_size: str | None) -> tuple[str | None, str | None]:
    if raw_size is None:
        return None, None
    requested_size = raw_size.strip().lower().replace("×", "x").replace(" ", "")
    if requested_size == "auto":
        return requested_size, requested_size
    if requested_size in SUPPORTED_IMAGE_SIZES:
        return requested_size, requested_size
    match = re.fullmatch(r"(\d+)x(\d+)", requested_size)
    if not match or int(match.group(1)) == 0 or int(match.group(2)) == 0:
        supported = ", ".join(("auto", *SUPPORTED_IMAGE_SIZES))
        raise GenerateImageError(f"image size must be WIDTHxHEIGHT; supported native sizes: {supported}")
    aspect_ratio = int(match.group(1)) / int(match.group(2))
    resolved_size = min(
        SUPPORTED_IMAGE_SIZES,
        key=lambda size: abs(aspect_ratio - SUPPORTED_IMAGE_ASPECT_RATIOS[size]),
    )
    return requested_size, resolved_size


def _load_gateway_config() -> tuple[str, str, dict[str, str]]:
    runtime_home = Path(os.environ.get(RUNTIME_HOME_ENV, "~/.xigua-agent")).expanduser()
    config_path = runtime_home / "config.json"
    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise GenerateImageError(f"Xigua runtime config not found: {config_path}") from exc
    except (OSError, json.JSONDecodeError) as exc:
        raise GenerateImageError(f"cannot read Xigua runtime config: {exc}") from exc

    providers = config.get("providers")
    if not isinstance(providers, dict):
        raise GenerateImageError("Xigua runtime config has no provider registry")
    items = providers.get("items")
    if not isinstance(items, list) or not items:
        raise GenerateImageError("Xigua runtime config has no enabled provider")
    default_name = str(providers.get("default") or "").strip()
    provider = next(
        (
            item
            for item in items
            if isinstance(item, dict)
            and item.get("enabled", True)
            and str(item.get("name") or "").strip() == default_name
        ),
        None,
    )
    if provider is None:
        provider = next(
            (item for item in items if isinstance(item, dict) and item.get("enabled", True)),
            None,
        )
    if provider is None:
        raise GenerateImageError("Xigua runtime config has no enabled provider")

    base_url = str(provider.get("baseURL") or provider.get("baseUrl") or "").strip()
    if not base_url:
        raise GenerateImageError("managed provider has no baseURL")
    api_key_env = str(provider.get("apiKeyEnv") or "").strip()
    if not api_key_env:
        raise GenerateImageError("managed provider has no apiKeyEnv")
    managed_credential = os.environ.get(api_key_env, "").strip()
    if not managed_credential:
        raise GenerateImageError(f"managed gateway credential is missing from {api_key_env}")

    headers: dict[str, str] = {}
    configured_headers = provider.get("headers")
    if isinstance(configured_headers, dict):
        headers.update(
            {
                str(key): str(value)
                for key, value in configured_headers.items()
                if str(key).strip() and isinstance(value, (str, int, float, bool))
            }
        )
    headers["Authorization"] = f"Bearer {managed_credential}"
    headers["Content-Type"] = "application/json"
    return f"{base_url.rstrip('/')}/images/generations", api_key_env, headers


def _request_generation(
    endpoint: str,
    headers: dict[str, str],
    payload: dict[str, Any],
    timeout_seconds: int,
) -> dict[str, Any]:
    request = Request(
        endpoint,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers=headers,
        method="POST",
    )
    try:
        with urlopen(request, timeout=timeout_seconds) as response:
            body = response.read(MAX_GATEWAY_RESPONSE_BYTES + 1)
    except HTTPError as exc:
        detail = exc.read(64 * 1024).decode("utf-8", errors="replace").strip()
        raise GenerateImageError(
            f"image gateway returned HTTP {exc.code}: {detail or exc.reason}"
        ) from exc
    except (OSError, URLError) as exc:
        raise GenerateImageError(f"cannot reach image gateway: {exc}") from exc
    if len(body) > MAX_GATEWAY_RESPONSE_BYTES:
        raise GenerateImageError("image gateway response exceeds the 100 MiB safety limit")
    try:
        result = json.loads(body)
    except json.JSONDecodeError as exc:
        raise GenerateImageError("image gateway returned invalid JSON") from exc
    if not isinstance(result, dict):
        raise GenerateImageError("image gateway returned an invalid result")
    return result


def _decode_base64_image(value: str) -> bytes:
    encoded = value.strip()
    if encoded.startswith("data:"):
        if "," not in encoded:
            raise GenerateImageError("image gateway returned an invalid data URL")
        encoded = encoded.split(",", 1)[1]
    if len(encoded) > (MAX_IMAGE_BYTES * 4 // 3) + 16:
        raise GenerateImageError("generated image exceeds the 50 MiB safety limit")
    try:
        data = base64.b64decode(encoded, validate=True)
    except (ValueError, binascii.Error) as exc:
        raise GenerateImageError("image gateway returned invalid base64 data") from exc
    if not data:
        raise GenerateImageError("image gateway returned an empty image")
    if len(data) > MAX_IMAGE_BYTES:
        raise GenerateImageError("generated image exceeds the 50 MiB safety limit")
    return data


def _download_image(url: str) -> tuple[bytes, str | None]:
    parsed = urlparse(url)
    if parsed.scheme not in {"http", "https"}:
        raise GenerateImageError("image gateway returned an unsupported image URL")
    request = Request(url, headers={"Accept": "image/*"})
    try:
        with urlopen(request, timeout=180) as response:
            content_type = response.headers.get_content_type()
            data = response.read(MAX_IMAGE_BYTES + 1)
    except HTTPError as exc:
        raise GenerateImageError(f"generated image download returned HTTP {exc.code}") from exc
    except (OSError, URLError) as exc:
        raise GenerateImageError(f"cannot download generated image: {exc}") from exc
    if not data:
        raise GenerateImageError("generated image download was empty")
    if len(data) > MAX_IMAGE_BYTES:
        raise GenerateImageError("generated image exceeds the 50 MiB safety limit")
    return data, content_type


def _image_extension(data: bytes, content_type: str | None = None) -> str:
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        return ".png"
    if data.startswith(b"\xff\xd8\xff"):
        return ".jpg"
    if data.startswith((b"GIF87a", b"GIF89a")):
        return ".gif"
    if data.startswith(b"RIFF") and data[8:12] == b"WEBP":
        return ".webp"
    return {
        "image/png": ".png",
        "image/jpeg": ".jpg",
        "image/gif": ".gif",
        "image/webp": ".webp",
    }.get(str(content_type or "").lower(), ".png")


def _resolve_output_dir(raw_output_dir: str) -> tuple[Path, Path]:
    workspace = Path(os.environ.get(WORKSPACE_ENV, os.getcwd())).expanduser().resolve()
    output_dir = Path(raw_output_dir).expanduser()
    if not output_dir.is_absolute():
        output_dir = workspace / output_dir
    output_dir = output_dir.resolve()
    try:
        output_dir.relative_to(workspace)
    except ValueError as exc:
        raise GenerateImageError("output directory must stay inside the Xigua workspace") from exc
    output_dir.mkdir(parents=True, exist_ok=True)
    return workspace, output_dir


def _save_image(output_dir: Path, data: bytes, index: int, content_type: str | None) -> Path:
    extension = _image_extension(data, content_type)
    with tempfile.NamedTemporaryFile(
        mode="wb",
        prefix=f"generated-{index + 1}-",
        suffix=extension,
        dir=output_dir,
        delete=False,
    ) as handle:
        handle.write(data)
        temporary_path = Path(handle.name)
    final_path = temporary_path.with_name(temporary_path.name.replace("generated-", "image-", 1))
    temporary_path.replace(final_path)
    return final_path


def _extract_and_save_images(
    result: dict[str, Any],
    workspace: Path,
    output_dir: Path,
) -> list[dict[str, Any]]:
    items = result.get("data")
    if not isinstance(items, list) or not items:
        raise GenerateImageError("image gateway returned no images")
    files: list[dict[str, Any]] = []
    for index, item in enumerate(items):
        if not isinstance(item, dict):
            raise GenerateImageError("image gateway returned an invalid image item")
        content_type: str | None = None
        if isinstance(item.get("b64_json"), str):
            data = _decode_base64_image(item["b64_json"])
        elif isinstance(item.get("url"), str):
            data, content_type = _download_image(item["url"])
        else:
            raise GenerateImageError("image gateway returned neither b64_json nor url")
        path = _save_image(output_dir, data, index, content_type)
        files.append(
            {
                "path": str(path),
                "workspacePath": str(path.relative_to(workspace)),
                "bytes": len(data),
                "mimeType": content_type or mimetypes.guess_type(path.name)[0] or "image/png",
            }
        )
    return files


def main() -> int:
    args = _parse_args()
    try:
        prompt = _read_prompt(args)
        model = os.environ.get(MODEL_ENV, "").strip()
        if not model:
            raise GenerateImageError(
                "this instance has no image-generation model; bind one in Desktop first"
            )
        endpoint, _api_key_env, headers = _load_gateway_config()
        timeout_seconds = _resolve_timeout_seconds()
        requested_size, resolved_size = _resolve_image_size(args.size)
        workspace, output_dir = _resolve_output_dir(args.output_dir)
        payload: dict[str, Any] = {
            "model": model,
            "prompt": prompt,
            "n": args.count,
        }
        if resolved_size:
            payload["size"] = resolved_size
        if args.quality:
            payload["quality"] = args.quality
        result = _request_generation(endpoint, headers, payload, timeout_seconds)
        files = _extract_and_save_images(result, workspace, output_dir)
        artifacts = [
            {
                "type": "image",
                "path": item["workspacePath"],
                "name": Path(item["path"]).name,
                "size": item["bytes"],
                "mimeType": item["mimeType"],
            }
            for item in files
        ]
        print(
            json.dumps(
                {
                    "ok": True,
                    "model": model,
                    **({
                        "requestedSize": requested_size,
                        "resolvedSize": resolved_size,
                    } if requested_size and requested_size != resolved_size else {}),
                    "files": files,
                    "artifacts": artifacts,
                },
                ensure_ascii=False,
            )
        )
        return 0
    except GenerateImageError as exc:
        print(json.dumps({"ok": False, "error": str(exc)}, ensure_ascii=False), file=sys.stderr)
        return 1


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