import argparse
import base64
import json
import os
import re
import sys
import time
import urllib.request
from pathlib import Path

BASE_URL = os.environ.get('XIAKE_BASE_URL') or os.environ.get('XIAKE_API_URL') or 'https://yunwu.ai/v1'
BASE_URL = BASE_URL.rstrip('/')
ENDPOINT = f"{BASE_URL}/images/generations"
MODEL = 'gpt-image-2'
SUPPORTED_SIZES = {
    '1024x1024', '1536x1024', '1024x1536', '2048x2048', '2048x1152', '3840x2160', '2160x3840', 'auto'
}


def get_api_key():
    for key in ('XIAKE_API', 'XIAKE_API_KEY', 'XI_API_KEY', 'OPENAI_API_KEY'):
        value = os.environ.get(key)
        if value:
            return value
    raise SystemExit('Missing API key in env: XIAKE_API / XIAKE_API_KEY / XI_API_KEY / OPENAI_API_KEY')


def infer_size(user_size: str | None, prompt: str) -> str:
    text = f"{user_size or ''} {prompt or ''}".lower()
    if user_size and user_size in SUPPORTED_SIZES:
        return user_size
    if re.search(r'4k.*(竖|vertical)|2160x3840|9:16', text):
        return '2160x3840'
    if re.search(r'4k.*(横|landscape)|3840x2160|16:9', text):
        return '3840x2160'
    if re.search(r'2k.*(竖|vertical)', text):
        return '1024x1536'
    if re.search(r'2k.*(横|landscape)|2048x1152', text):
        return '2048x1152'
    if re.search(r'2k.*(方|square)|2048x2048', text):
        return '2048x2048'
    # 竖版内容
    if re.search(r'手机壁纸|wallpaper|竖屏壁纸|vertical wallpaper|9:16', text):
        return '2160x3840'
    if re.search(r'小红书|竖版海报|竖图海报|poster|vertical|portrait|3:4|4:5|封面', text):
        return '1024x1536'

    # 横版内容
    if re.search(r'youtube|youtube封面|视频封面|横版海报|横图|banner|公众号头图|ppt配图|landscape|16:9', text):
        return '1536x1024'

    # 方图内容
    if re.search(r'头像|logo|图标|icon|square|1:1|方图', text):
        return '1024x1024'

    # 默认尺寸
    return '1024x1024'


def normalize_format(fmt: str | None) -> str:
    if not fmt:
        return 'png'
    fmt = fmt.lower()
    return fmt if fmt in {'png', 'jpeg', 'webp'} else 'png'


def normalize_quality(q: str | None) -> str:
    if not q:
        return 'auto'
    q = q.lower()
    return q if q in {'low', 'medium', 'high', 'auto'} else 'auto'


def download(url: str, out_path: Path):
    req = urllib.request.Request(url, headers={'User-Agent': 'OpenClaw/xiakeyun-image'})
    with urllib.request.urlopen(req, timeout=180) as resp:
        out_path.write_bytes(resp.read())


def deep_find_image(payload):
    if isinstance(payload, dict):
        for key in ('url', 'b64_json', 'b64', 'image_base64'):
            if key in payload and payload[key]:
                return key, payload[key]
        for value in payload.values():
            found = deep_find_image(value)
            if found:
                return found
    elif isinstance(payload, list):
        for item in payload:
            found = deep_find_image(item)
            if found:
                return found
    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prompt', required=True)
    ap.add_argument('--size')
    ap.add_argument('--quality', default='auto')
    ap.add_argument('--format', default='png')
    ap.add_argument('--n', type=int, default=1)
    desktop = str(Path.home() / 'Desktop')
    ap.add_argument('--outdir', default=desktop)
    args = ap.parse_args()

    n = min(max(args.n, 1), 10)
    size = infer_size(args.size, args.prompt)
    fmt = normalize_format(args.format)
    quality = normalize_quality(args.quality)
    api_key = get_api_key()

    body = {
        'model': MODEL,
        'prompt': args.prompt,
        'n': n,
        'size': size,
        'quality': quality
    }

    req = urllib.request.Request(
        ENDPOINT,
        data=json.dumps(body).encode('utf-8'),
        headers={
            'Content-Type': 'application/json',
            'Accept': 'application/json',
            'Authorization': f'Bearer {api_key}',
            'User-Agent': 'OpenClaw-xiakeyun-image/1.0',
            'Connection': 'keep-alive'
        },
        method='POST',
    )

    def should_not_retry(msg: str) -> bool:
        msg = msg.lower()
        keywords = ['余额', 'balance', 'insufficient', 'quota', 'api key', 'invalid key']
        return any(k in msg for k in keywords)

    retries = 1
    attempt = 0
    payload = None

    while attempt <= retries:
        try:
            print(json.dumps({'status': 'submit_task', 'attempt': attempt+1}, ensure_ascii=False))
            with urllib.request.urlopen(req) as resp:
                payload = json.loads(resp.read().decode('utf-8', errors='replace'))
                break
        except urllib.error.HTTPError as e:
            detail = e.read().decode('utf-8', errors='replace')
            if should_not_retry(detail):
                raise SystemExit(f'HTTP {e.code}: {detail}')
            if attempt >= retries:
                raise SystemExit(f'HTTP {e.code}: {detail}')
        except Exception as e:
            msg = str(e)
            if should_not_retry(msg):
                raise SystemExit(msg)
            if attempt >= retries:
                raise SystemExit(msg)
        attempt += 1
        time.sleep(1)

    if payload is None:
        raise SystemExit('Image task submission failed')

    task_id = payload.get('task_id') or payload.get('id')
    if not task_id:
        # 如果接口直接返回图片则继续原流程
        pass
    else:
        status_url = f"{BASE_URL}/images/generations/{task_id}"
        start = time.time()
        while True:
            if time.time() - start > 240:
                raise SystemExit('Image generation polling timeout (240s)')

            print(json.dumps({'status': 'polling', 'task_id': task_id}, ensure_ascii=False))

            poll_req = urllib.request.Request(
                status_url,
                headers={
                    'Accept': 'application/json',
                    'Authorization': f'Bearer {api_key}',
                    'User-Agent': 'OpenClaw-xiakeyun-image/1.0'
                }
            )

            with urllib.request.urlopen(poll_req) as resp:
                poll_payload = json.loads(resp.read().decode('utf-8', errors='replace'))

            status = poll_payload.get('status') or poll_payload.get('state')

            if status in ('succeeded', 'success', 'completed', 'done'):
                payload = poll_payload
                break

            if status in ('failed', 'error'):
                raise SystemExit(f"Image generation failed: {poll_payload}")

            time.sleep(3)

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    ts = int(time.time())

    found = deep_find_image(payload)
    if not found:
        print(json.dumps({'ok': False, 'size': size, 'response': payload}, ensure_ascii=False))
        raise SystemExit('No image url/base64 found in response')

    kind, value = found
    ext = 'jpg' if fmt == 'jpeg' else fmt
    out_path = outdir / f'xiakeyun_{ts}.{ext}'

    if kind == 'url':
        download(value, out_path)
    else:
        out_path.write_bytes(base64.b64decode(value))

    print(json.dumps({
        'image': str(out_path.resolve()),
        'size': size
    }, ensure_ascii=False))


if __name__ == '__main__':
    main()
