import json
import mimetypes
import os
import time
import uuid
from pathlib import Path

import requests

BASE_URL = os.environ.get("BASE_URL", "https://ic.xshow.live").rstrip("/")
IMAGE = "source.png"
session = requests.Session()
session.cookies.set("ic_session", os.environ["IC_SESSION"])


def checked(response):
    if not response.ok:
        raise RuntimeError(f"HTTP {response.status_code}: {response.text[:1000]}")
    return response


loras = [
    {"name": "krea2/girlslike/girlslike_krea2_lyf.safetensors", "strength": 1.0},
    {"name": "krea2/sex/krea missionary.safetensors", "strength": 1.0},
]
with open(IMAGE, "rb") as source:
    created = checked(session.post(
        f"{BASE_URL}/api/proxy/v1/images/edits/tasks",
        headers={"X-IC-Operation-ID": str(uuid.uuid4())},
        data={
            "model": "krea2-turbo-v25-edit",
            "prompt": "编辑原图，同时保持主体身份、构图和光照自然",
            "loras": json.dumps(loras, ensure_ascii=False),
            "seed": "123456",
            "output_sizing": "source_ratio",
        },
        files={"image": (Path(IMAGE).name, source, mimetypes.guess_type(IMAGE)[0] or "image/png")},
        timeout=300,
    )).json()

task_id = created.get("id") or (created.get("data") or {}).get("id")
if not task_id:
    raise RuntimeError(f"创建响应缺少任务 ID: {created}")

deadline = time.time() + 6 * 60 * 60
while time.time() < deadline:
    state = checked(session.post(f"{BASE_URL}/api/proxy/v1/images/tasks/{task_id}", timeout=60)).json()
    status = str(state.get("status") or (state.get("data") or {}).get("status") or "pending").lower()
    print("status:", status)
    if status == "completed":
        result = checked(session.get(
            f"{BASE_URL}/api/proxy/v1/images/tasks/{task_id}/content",
            params={"index": 0},
            timeout=600,
        ))
        Path("krea2-turbo-edit-result.png").write_bytes(result.content)
        break
    if status in {"failed", "cancelled", "canceled", "expired"}:
        raise RuntimeError(f"任务失败: {state}")
    time.sleep(3)
else:
    raise TimeoutError(f"任务 {task_id} 超时")