#!/usr/bin/env python3
"""AI 出图台 command-line client. Needs only Python 3 (no extra packages).

Set once:
    GPU_URL = https://win-8845.tail062337.ts.net      (the site address)
    GPU_KEY = gw_xxxxxxxx                             (API key from the admin)

Examples:
    python gpu.py image "一位穿红裙的女人站在海边，电影感" -n 2 --size portrait
    python gpu.py edit photo.jpg "把背景换成夜晚的城市" --confirm
    python gpu.py edit base.png "把图1的脸换成图2的脸" --ref face.png --confirm
    python gpu.py video first.png "她转身走向镜头" --model h3 --seconds 5 --confirm
    python gpu.py video job:1016203011ab12:1016203011ab12_1.png "镜头缓慢推进" --model wan
    python gpu.py list
    python gpu.py get <job id>          (download a finished job's files)

Each command waits for the job and downloads the results into the current folder
(use --no-wait to just submit).
"""
import argparse
import base64
import json
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

URL = os.environ.get("GPU_URL", "").rstrip("/")
KEY = os.environ.get("GPU_KEY", "")
for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass


def call(path, payload=None, raw=False):
    req = urllib.request.Request(URL + path, headers={"X-API-Key": KEY, "Content-Type": "application/json"},
                                 data=None if payload is None else json.dumps(payload).encode())
    try:
        with urllib.request.urlopen(req, timeout=300) as r:
            data = r.read()
    except urllib.error.HTTPError as e:
        try:
            msg = json.loads(e.read()).get("error")
        except Exception:
            msg = str(e)
        sys.exit(f"错误：{msg}")
    return data if raw else json.loads(data)


CONFIRM = False


def image_ref(arg):
    """A local file is uploaded first; 'job:...' / 'upload:...' refs pass through."""
    if arg.startswith(("job:", "upload:")):
        return arg
    if not CONFIRM:
        sys.exit("上传本地图片需要加 --confirm，表示确认：图片为 AI 生成，或来自本人同意的成年人")
    p = Path(arg)
    if not p.is_file():
        sys.exit(f"找不到文件 {arg}")
    print(f"上传 {p.name} …")
    up = call("/api/upload", {"filename": p.name, "consent": True,
                              "data": base64.b64encode(p.read_bytes()).decode()})
    return "upload:" + up["id"]


def wait_and_download(job, outdir="."):
    jid = job["id"]
    print(f"任务号 {jid}，等待中 …")
    last = None
    while True:
        j = call(f"/api/jobs/{jid}")
        st = j["status"]
        msg = st + (f"（排第 {j['queue_position']} 个）" if j.get("queue_position") else "")
        if msg != last:
            print("  状态：" + msg)
            last = msg
        if st in ("done", "failed", "cancelled"):
            break
        time.sleep(5)
    if st != "done":
        sys.exit(f"任务没有成功：{j.get('error') or st}")
    download(j, outdir)


def download(j, outdir="."):
    for name in j.get("outputs", []):
        data = call(f"/api/files/{j['id']}/{name}?download=1", raw=True)
        dest = Path(outdir) / name
        dest.write_bytes(data)
        print(f"  已下载 {dest}  ({len(data) // 1024} KB)")


def main():
    if not URL or not KEY:
        sys.exit("先设置环境变量 GPU_URL 和 GPU_KEY（见脚本开头说明）")
    ap = argparse.ArgumentParser(description="AI 出图台")
    sub = ap.add_subparsers(dest="cmd", required=True)
    a = sub.add_parser("image", help="文字出图")
    a.add_argument("prompt"); a.add_argument("-n", "--count", type=int, default=1)
    a.add_argument("--size", default="portrait", choices=["portrait", "landscape", "square"])
    a.add_argument("--seed")
    e = sub.add_parser("edit", help="改图 / 换脸")
    e.add_argument("image"); e.add_argument("prompt"); e.add_argument("--ref", help="第二张参考图")
    e.add_argument("-n", "--count", type=int, default=1); e.add_argument("--seed")
    v = sub.add_parser("video", help="图生视频")
    v.add_argument("image"); v.add_argument("prompt")
    v.add_argument("--model", default="h3", choices=["h3", "wan"])
    v.add_argument("--size", default="portrait", choices=["portrait", "landscape"])
    v.add_argument("--seconds", type=float, default=5, help="2–30 秒"); v.add_argument("--preset", default="fast", choices=["fast", "turbo"])
    v.add_argument("--lora", help="H3 风格/动作 LoRA 文件名（用 python gpu.py loras 查看）")
    v.add_argument("--lora-strength", type=float, default=1.0)
    v.add_argument("--seed")
    sub.add_parser("list", help="列出最近的任务")
    sub.add_parser("loras", help="列出可用的 H3 LoRA")
    g = sub.add_parser("get", help="下载某个任务的成品"); g.add_argument("job")
    for p in (a, e, v):
        p.add_argument("--no-wait", action="store_true")
    for p in (e, v):
        p.add_argument("--confirm", action="store_true", help="确认上传的图为 AI 生成或来自本人同意的成年人")
    args = ap.parse_args()
    global CONFIRM
    CONFIRM = getattr(args, "confirm", False)

    if args.cmd == "list":
        for j in call("/api/jobs")["jobs"][:30]:
            print(f"{j['id']}  {j['type']:5}  {j['status']:9}  {(j['params'].get('prompt') or '')[:50]}")
        return
    if args.cmd == "loras":
        for n in call("/api/loras")["h3"]:
            print(n)
        return
    if args.cmd == "get":
        return download(call(f"/api/jobs/{args.job}"))
    if args.cmd == "image":
        body = {"type": "t2i", "prompt": args.prompt, "count": args.count, "size": args.size, "seed": args.seed}
    elif args.cmd == "edit":
        body = {"type": "edit", "image": image_ref(args.image), "prompt": args.prompt, "count": args.count,
                "seed": args.seed, "image2": image_ref(args.ref) if args.ref else None}
    else:
        body = {"type": args.model, "image": image_ref(args.image), "prompt": args.prompt, "size": args.size,
                "seconds": args.seconds, "preset": args.preset, "seed": args.seed,
                "lora": args.lora, "lora_strength": args.lora_strength}
    job = call("/api/jobs", body)
    if args.no_wait:
        print(f"已提交，任务号 {job['id']}")
    else:
        wait_and_download(job)


if __name__ == "__main__":
    main()