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

ROOT = Path(__file__).resolve().parents[1]
ENV_PATH = ROOT / '.env'
OUTPUT_DIR = ROOT / 'output'


def load_env(path: Path):
    if not path.exists():
        return
    for line in path.read_text(encoding='utf-8').splitlines():
        line = line.strip()
        if not line or line.startswith('#') or '=' not in line:
            continue
        k, v = line.split('=', 1)
        os.environ.setdefault(k.strip(), v.strip())


def post_json(url: str, payload: dict, headers: dict) -> dict:
    data = json.dumps(payload).encode('utf-8')
    req = urllib.request.Request(url, data=data, headers={**headers, 'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=180) as resp:
        raw = resp.read().decode('utf-8', errors='replace')
    return json.loads(raw)


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


def save_b64(data: str, out_path: Path):
    if ',' in data and data.split(',', 1)[0].startswith('data:'):
        data = data.split(',', 1)[1]
    out_path.write_bytes(base64.b64decode(data))


def main():
    load_env(ENV_PATH)

    parser = argparse.ArgumentParser(description='Generate image with Doubao/Ark Seedream')
    parser.add_argument('--prompt', required=True)
    parser.add_argument('--negative-prompt', default='')
    parser.add_argument('--width', type=int, default=1920)
    parser.add_argument('--height', type=int, default=1920)
    parser.add_argument('--size', default='', help='Ark size preset like 2K/4K; overrides width/height when set')
    parser.add_argument('--model', default=os.getenv('ARK_IMAGE_MODEL', 'seedream-5.0-lite'))
    parser.add_argument('--base-url', default=os.getenv('ARK_BASE_URL', 'https://ark.cn-beijing.volces.com/api/v3'))
    parser.add_argument('--api-key', default=os.getenv('ARK_API_KEY') or os.getenv('DOUBAO_API_KEY', ''))
    args = parser.parse_args()

    if not args.api_key:
        print(json.dumps({'success': False, 'error': 'Missing DOUBAO_API_KEY in .env or --api-key'}))
        sys.exit(2)

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    ts = int(time.time())
    out_path = OUTPUT_DIR / f'seedream_{ts}.png'

    # This follows the OpenAI-compatible image generation shape used by Ark-style APIs.
    payload = {
        'model': args.model,
        'prompt': args.prompt,
        'size': args.size or f'{args.width}x{args.height}'
    }
    if args.negative_prompt:
        payload['negative_prompt'] = args.negative_prompt

    url = args.base_url.rstrip('/') + '/images/generations'
    headers = {
        'Authorization': f'Bearer {args.api_key}'
    }

    try:
        result = post_json(url, payload, headers)
    except urllib.error.HTTPError as e:
        body = e.read().decode('utf-8', errors='replace')
        print(json.dumps({'success': False, 'error': f'HTTP {e.code}', 'body': body}, ensure_ascii=False))
        sys.exit(1)
    except Exception as e:
        print(json.dumps({'success': False, 'error': str(e)}, ensure_ascii=False))
        sys.exit(1)

    data = result.get('data') or result.get('images') or []
    if not data:
        print(json.dumps({'success': False, 'error': 'No image data returned', 'raw': result}, ensure_ascii=False))
        sys.exit(1)

    first = data[0]
    if isinstance(first, str):
        download(first, out_path)
    elif isinstance(first, dict):
        if first.get('url'):
            download(first['url'], out_path)
        elif first.get('b64_json'):
            save_b64(first['b64_json'], out_path)
        elif first.get('base64'):
            save_b64(first['base64'], out_path)
        else:
            print(json.dumps({'success': False, 'error': 'Unknown image payload shape', 'raw': result}, ensure_ascii=False))
            sys.exit(1)
    else:
        print(json.dumps({'success': False, 'error': 'Unsupported response type', 'raw': result}, ensure_ascii=False))
        sys.exit(1)

    print(json.dumps({'success': True, 'image_path': str(out_path), 'raw': result}, ensure_ascii=False))


if __name__ == '__main__':
    main()