"""Linghang image API example. Python 3.8+, pip install requests.

models is read-only. generate/edit may incur charges. No automatic retries.
Set LINGHANG_API_KEY locally; never publish a real key with this file.
"""
import argparse
import base64
import os
import re
import sys
import uuid
from pathlib import Path
from urllib.parse import quote, urlsplit

import requests

BASE_URL = "https://api.tianjinlinghang.com/v1"
MAX_REFERENCE_IMAGES = 9
RATIOS = ["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"]


def normalize_base(value):
    value = value.strip().rstrip("/")
    url = urlsplit(value)
    if (url.scheme != "https" or not url.hostname or url.username or url.password
            or url.query or url.fragment or not url.path.endswith("/v1")):
        raise ValueError("Use an HTTPS base URL ending in /v1, without credentials/query/fragment.")
    return value


def protocol(model):
    if re.search(r"gemini.*image|(?:nano-)?banana", model, re.I):
        return "gemini"
    if re.search(r"grok.*image", model, re.I):
        return "grok"
    if re.match(r"gpt-image-", model, re.I):
        return "gpt"
    raise ValueError("This example supports GPT Image, Gemini image and Grok image only.")


def raster_type(data):
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png", "png"
    if data.startswith(b"\xff\xd8\xff"):
        return "image/jpeg", "jpg"
    if data.startswith(b"RIFF") and data[8:12] == b"WEBP":
        return "image/webp", "webp"
    raise ValueError("Expected PNG, JPEG or WebP bytes.")


def references(paths):
    if len(paths) > MAX_REFERENCE_IMAGES:
        raise ValueError("At most %d reference images." % MAX_REFERENCE_IMAGES)
    result, total = [], 0
    for name in paths:
        path = Path(name)
        with path.open("rb") as source:
            data = source.read(10 * 1024 * 1024 + 1)
        total += len(data)
        if len(data) > 10 * 1024 * 1024 or total > 100 * 1024 * 1024:
            raise ValueError("References exceed 10MB per file or 100MB total.")
        mime, ext = raster_type(data)
        result.append(("reference-%d.%s" % (len(result) + 1, ext), data, mime))
    return result


def validate_params(args, kind):
    if kind == "gpt":
        if args.aspect_ratio != "auto" or args.resolution != "auto":
            raise ValueError("GPT uses --size WIDTHxHEIGHT, not Gemini aspect/resolution flags.")
        if args.size == "auto":
            return
        flexible = re.match(r"gpt-image-(?:2(?:[.-]|$)|[24]k-adobe$)", args.model, re.I)
        if not flexible:
            if args.size not in ["1024x1024", "1536x1024", "1024x1536"]:
                raise ValueError("This GPT model supports only the standard image sizes.")
            return
        match = re.fullmatch(r"(\d+)x(\d+)", args.size)
        if not match:
            raise ValueError("Use --size WIDTHxHEIGHT or auto.")
        w, h = map(int, match.groups())
        if (min(w, h) <= 0 or max(w, h) > 3840 or w % 16 or h % 16
                or max(w / h, h / w) > 3 or not 655360 <= w * h <= 8294400):
            raise ValueError("Invalid GPT Image 2 dimensions; see the size limits in the guide.")
        tier = re.search(r"(?:-|image-)([124])k(?:-|$)", args.model, re.I)
        if tier and w * h > {"1": 1572864, "2": 4194304, "4": 8294400}[tier[1]]:
            raise ValueError("Dimensions exceed the named upstream resolution tier.")
    elif kind == "gemini":
        if args.size != "auto":
            raise ValueError("Gemini uses --aspect-ratio and --resolution, not --size.")
        flash = re.search(r"^gemini-3\.1-flash-image(?:-|$)|^nano-banana-2$", args.model, re.I)
        pro = re.search(r"^gemini-3(?:\.1)?-pro-image(?:-|$)|^nano-banana-pro$", args.model, re.I)
        ratios = RATIOS + (["1:4", "4:1", "1:8", "8:1"] if flash else [])
        sizes = ["auto", "512", "1K", "2K", "4K"] if flash else ["auto", "1K", "2K", "4K"] if pro else ["auto"]
        if args.aspect_ratio != "auto" and args.aspect_ratio not in ratios:
            raise ValueError("Unsupported aspect ratio for this Gemini model.")
        if args.resolution not in sizes:
            raise ValueError("Unsupported resolution for this Gemini model.")
    elif args.size != "auto" or args.aspect_ratio != "auto" or args.resolution != "auto":
        raise ValueError("The current site Grok adapter supports automatic size only.")


def request_json(session, method, url, **kwargs):
    # No redirects or retries: a failed generation may already have incurred cost.
    response = session.request(method, url, timeout=(15, 600 if method == "POST" else 15),
                               allow_redirects=False, **kwargs)
    request_id = response.headers.get("X-Oneapi-Request-Id", "")
    suffix = " Request ID: " + request_id if re.fullmatch(r"[\w-]{1,100}", request_id) else ""
    if not 200 <= response.status_code < 300:
        raise RuntimeError("HTTP %s.%s Check usage/error logs before resubmitting." % (response.status_code, suffix))
    try:
        result = response.json()
    except ValueError:
        raise RuntimeError("Expected a JSON response." + suffix) from None
    if not isinstance(result, dict):
        raise RuntimeError("Invalid API response shape." + suffix)
    return result, suffix


def list_models(session, base):
    body, _ = request_json(session, "GET", base + "/models")
    if not isinstance(body.get("data"), list):
        raise ValueError("No data[] model list returned.")
    return sorted({row["id"] for row in body["data"] if isinstance(row, dict)
                   and isinstance(row.get("id"), str)
                   and re.search(r"image|dall-e|flux|imagen|seedream|ideogram|banana", row["id"], re.I)
                   and not re.search(r"video|veo|omni", row["id"], re.I)})


def generate(session, base, args, refs):
    kind = protocol(args.model)
    validate_params(args, kind)
    if len(refs) > MAX_REFERENCE_IMAGES:
        raise ValueError("At most %d reference images." % MAX_REFERENCE_IMAGES)
    if kind == "grok" and refs:
        raise ValueError("Grok reference-image editing is not supported by the current site adapter.")
    if kind == "gemini":
        parts = [{"text": args.prompt}]
        parts += [{"inlineData": {"mimeType": mime, "data": base64.b64encode(data).decode("ascii")}}
                  for _, data, mime in refs]
        config = {"responseModalities": ["TEXT", "IMAGE"]}
        image_config = {}
        if args.aspect_ratio != "auto":
            image_config["aspectRatio"] = args.aspect_ratio
        if args.resolution != "auto":
            image_config["imageSize"] = args.resolution
        if image_config:
            config["imageConfig"] = image_config
        url = base[:-3] + "/v1beta/models/" + quote(args.model, safe="") + ":generateContent"
        return request_json(session, "POST", url, json={"contents": [{"role": "user", "parts": parts}], "generationConfig": config})
    body = {"model": args.model, "prompt": args.prompt, "n": 1, "response_format": "b64_json"}
    if kind == "gpt":
        body.update(size=args.size, quality="auto", output_format="png", moderation="auto")
    if refs:
        return request_json(session, "POST", base + "/images/edits",
                            data=body, files=[("image[]", ref) for ref in refs])
    return request_json(session, "POST", base + "/images/generations", json=body)


def decode_images(body, kind):
    if body.get("error"):
        raise ValueError("Upstream returned an error object; check usage logs.")
    encoded = []
    if kind == "gemini":
        if body.get("promptFeedback", {}).get("blockReason"):
            raise ValueError("Gemini safety block; no automatic retry.")
        for candidate in body.get("candidates", []):
            if candidate.get("finishReason") not in (None, "STOP", "MAX_TOKENS"):
                raise ValueError("Gemini generation did not finish successfully.")
            for part in candidate.get("content", {}).get("parts", []):
                if part.get("thought"):
                    continue
                inline = part.get("inlineData") or part.get("inline_data") or {}
                if inline.get("data"):
                    mime = inline.get("mimeType") or inline.get("mime_type")
                    if mime not in ("image/png", "image/jpeg", "image/webp"):
                        raise ValueError("Unsupported Gemini image MIME type.")
                    encoded.append((inline["data"], mime))
    else:
        encoded = [(row["b64_json"], None) for row in body.get("data", []) if row.get("b64_json")]
    if not encoded:
        raise ValueError("No inline image bytes returned (possibly text or URL only). Check logs; no automatic retry/download.")
    images = []
    for value, expected_mime in encoded:
        data = base64.b64decode(re.sub(r"\s", "", value), validate=True)
        actual_mime, extension = raster_type(data)
        if expected_mime and actual_mime != expected_mime:
            raise ValueError("Image bytes do not match the returned MIME type.")
        images.append((data, extension))
    return images


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["models", "generate", "edit"])
    parser.add_argument("--base-url", default=BASE_URL)
    parser.add_argument("--model")
    parser.add_argument("--prompt")
    parser.add_argument("--reference", action="append", default=[])
    parser.add_argument("--size", default="auto")
    parser.add_argument("--aspect-ratio", default="auto")
    parser.add_argument("--resolution", default="auto")
    parser.add_argument("--output-dir", default="images")
    args = parser.parse_args(argv)
    base = normalize_base(args.base_url)
    key = os.environ.get("LINGHANG_API_KEY", "").strip()
    if not key or "\r" in key or "\n" in key:
        raise ValueError("Set LINGHANG_API_KEY locally to a valid token for the chosen endpoint.")
    if args.action != "models":
        if not args.model or not args.prompt or not args.prompt.strip():
            raise ValueError("generate/edit require --model and --prompt.")
        if (args.action == "edit") != bool(args.reference):
            raise ValueError("Use edit with --reference; generate takes no reference files.")
        validate_params(args, protocol(args.model))
        if args.action == "edit" and protocol(args.model) == "grok":
            raise ValueError("The current Grok adapter does not support reference images.")
    refs = references(args.reference)
    with requests.Session() as session:
        session.mount("https://", requests.adapters.HTTPAdapter(max_retries=0))
        session.headers["Authorization"] = "Bearer " + key
        models = list_models(session, base)
        if args.action == "models":
            print("\n".join(models))
            return
        if args.model not in models:
            raise ValueError("Requested model is not in this token's current /v1/models list.")
        body, request_id = generate(session, base, args, refs)
    try:
        images = decode_images(body, protocol(args.model))
    except (ValueError, TypeError, KeyError, AttributeError):
        raise ValueError("No usable image returned or response validation failed." + request_id + " Check logs before resubmitting.") from None
    output = Path(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    for data, extension in images:
        path = output / ("linghang-" + uuid.uuid4().hex + "." + extension)
        with path.open("xb") as target:
            target.write(data)
        print(path.resolve())


if __name__ == "__main__":
    try:
        main()
    except requests.RequestException:
        print("Network request failed; check logs before resubmitting. No automatic retry.", file=sys.stderr)
        sys.exit(1)
    except (ValueError, RuntimeError, OSError) as error:
        text = str(error)
        secret = os.environ.get("LINGHANG_API_KEY", "").strip()
        print(text.replace(secret, "[REDACTED]") if secret else text, file=sys.stderr)
        sys.exit(1)