#!/usr/bin/env python3
import argparse
import base64
import json
import os
import sys
import time
import urllib.request
import urllib.error


def load_config():
    candidates = [
        os.environ.get("OPENCLAW_CONFIG"),
        "/opt/openclaw/.openclaw/openclaw.json",
        os.path.expanduser("~/.openclaw/openclaw.json"),
    ]
    for path in candidates:
        if path and os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
    raise RuntimeError("openclaw config not found")


def resolve_provider(cfg):
    providers = cfg.get("models", {}).get("providers", {}) or {}
    if "clawkai" in providers:
        provider = providers["clawkai"]
    elif providers:
        provider = next(iter(providers.values()))
    else:
        raise RuntimeError("no providers configured")
    base_url = provider.get("baseUrl") or provider.get("base_url")
    api_key = provider.get("apiKey") or provider.get("api_key")
    if not base_url or not api_key:
        raise RuntimeError("provider baseUrl or apiKey missing")
    return base_url.rstrip("/"), api_key


def resolve_model(cfg, explicit):
    if explicit:
        return explicit
    image_model = (
        cfg.get("agents", {})
        .get("defaults", {})
        .get("imageModel", {})
        .get("primary")
    )
    if image_model and "/" in image_model:
        return image_model.split("/", 1)[1]
    return image_model or "gpt-image-1.5"


def build_url(base_url):
    if base_url.endswith("/v1"):
        return f"{base_url}/images/generations"
    return f"{base_url}/v1/images/generations"


def request_json(url, api_key, payload):
    req = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "User-Agent": "curl/8.5.0",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        return json.loads(resp.read().decode("utf-8"))


def fetch_image(url):
    req = urllib.request.Request(url, headers={"User-Agent": "curl/8.5.0"})
    with urllib.request.urlopen(req, timeout=300) as resp:
        return resp.read()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--model")
    parser.add_argument("--output")
    args = parser.parse_args()

    cfg = load_config()
    base_url, api_key = resolve_provider(cfg)
    model = resolve_model(cfg, args.model)

    payload = {"model": model, "prompt": args.prompt}
    response = request_json(build_url(base_url), api_key, payload)
    data = response.get("data") or []
    if not data:
        raise RuntimeError("image response missing data")

    item = data[0]
    image_b64 = item.get("b64_json")
    image_url = item.get("url")
    if not image_b64 and not image_url:
        raise RuntimeError("image response missing b64_json/url")

    if args.output:
        out_path = args.output
    else:
        stamp = int(time.time())
        out_path = f"/tmp/clawkai-image-{stamp}.png"

    if image_b64:
        with open(out_path, "wb") as f:
            f.write(base64.b64decode(image_b64))
    else:
        with open(out_path, "wb") as f:
            f.write(fetch_image(image_url))

    print(out_path)


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        sys.stderr.write(str(exc) + "\n")
        sys.exit(1)