"""即梦(Jimeng) image adapter — Volcengine SDK, async workflow.
Uses CVSync2AsyncSubmitTask → CVSync2AsyncGetResult polling.
Supports text-to-image via jimeng_t2i_v40."""
import json
import os
import sys
import time

import requests
from volcengine.visual.VisualService import VisualService


class JimengImageAdapter:
    POLL_INTERVAL = 5    # 图片生成比视频快，5 秒轮询
    MAX_WAIT = 120       # 最长等 2 分钟

    # 比例 → 推荐的 2K 分辨率（官方文档中的推荐值）
    RATIO_MAP = {
        "1:1":  (2048, 2048),
        "4:3":  (2304, 1728),
        "3:4":  (1728, 2304),
        "3:2":  (2496, 1664),
        "2:3":  (1664, 2496),
        "16:9": (2560, 1440),
        "9:16": (1440, 2560),
        "21:9": (3024, 1296),
        "9:21": (1296, 3024),
    }

    def generate(self, prompt, config, **kwargs):
        access_key = config["access_key"]
        secret_key = config["secret_key"]
        req_key = config.get("model", "jimeng_t2i_v40")
        out_dir = kwargs["out_dir"]
        aspect_ratio = kwargs.get("aspect_ratio", "1:1")

        vs = VisualService()
        vs.set_ak(access_key)
        vs.set_sk(secret_key)

        # ═══ Step 1: 提交任务 ═══
        body = {
            "req_key": req_key,
            "prompt": prompt,
            "force_single": True,  # 强制输出单图，确保智能体调用行为可预测
        }

        # 宽高映射
        if aspect_ratio in self.RATIO_MAP:
            w, h = self.RATIO_MAP[aspect_ratio]
            body["width"] = w
            body["height"] = h
        else:
            # 不传 width/height，让模型根据 prompt 自动判断
            body["size"] = 2048 * 2048  # 2K 默认面积

        print(f"[jimeng_image] 提交文生图任务 req_key={req_key} ratio={aspect_ratio}...", file=sys.stderr)
        try:
            data = vs.cv_sync2async_submit_task(body)
        except Exception as e:
            if hasattr(e, 'args') and len(e.args) > 0 and isinstance(e.args[0], bytes):
                try:
                    err_json = json.loads(e.args[0].decode('utf-8'))
                    code = err_json.get("code", 0)
                    msg = err_json.get("message", "")
                    raise RuntimeError(f"即梦 API 错误: {msg} (code={code})")
                except RuntimeError:
                    raise
                except Exception:
                    pass
            raise RuntimeError(f"即梦提交失败: {str(e)}")

        if data.get("code") != 10000:
            raise RuntimeError(f"即梦 API 错误: {data.get('message')} (code={data.get('code')})")

        task_id = data["data"]["task_id"]
        print(f"[jimeng_image] task_id={task_id}, 开始轮询...", file=sys.stderr)

        # ═══ Step 2: 轮询状态 ═══
        image_urls = None
        binary_data = None
        elapsed = 0
        while elapsed < self.MAX_WAIT:
            time.sleep(self.POLL_INTERVAL)
            elapsed += self.POLL_INTERVAL

            query_body = {
                "req_key": req_key,
                "task_id": task_id,
                "req_json": json.dumps({"return_url": True}),
            }
            try:
                qdata = vs.cv_sync2async_get_result(query_body)
            except Exception as e:
                if hasattr(e, 'args') and len(e.args) > 0 and isinstance(e.args[0], bytes):
                    try:
                        err_json = json.loads(e.args[0].decode('utf-8'))
                        code = err_json.get("code", 0)
                        if code in (50429, 50430, 50500, 50501):
                            print(f"[jimeng_image] 轮询 {elapsed}s: 可重试错误 code={code}", file=sys.stderr)
                            continue
                        raise RuntimeError(f"即梦查询失败: {err_json.get('message')} (code={code})")
                    except RuntimeError:
                        raise
                    except Exception:
                        pass
                raise RuntimeError(f"即梦查询异常: {str(e)}")

            if qdata.get("code") != 10000:
                print(f"[jimeng_image] 轮询 {elapsed}s: code={qdata.get('code')} msg={qdata.get('message')}", file=sys.stderr)
                if qdata.get("code") in (50429, 50430, 50500, 50501):
                    continue
                raise RuntimeError(f"即梦状态获取失败: {qdata.get('message')} (code={qdata.get('code')})")

            status = qdata.get("data", {}).get("status", "")
            print(f"[jimeng_image] 轮询 {elapsed}s: status={status}", file=sys.stderr)

            if status == "done":
                image_urls = qdata["data"].get("image_urls", [])
                binary_data = qdata["data"].get("binary_data_base64", [])
                break
            elif status in ("in_queue", "generating"):
                continue
            elif status in ("not_found", "expired"):
                raise RuntimeError(f"即梦任务异常: status={status}")

        if not image_urls and not binary_data:
            raise RuntimeError(f"即梦图片生成超时 ({self.MAX_WAIT}s), task_id={task_id}")

        print(f"[jimeng_image] 生成完成, 下载中...", file=sys.stderr)

        # ═══ Step 3: 下载图片 ═══
        import base64
        saved = []

        if image_urls:
            for i, url in enumerate(image_urls):
                fname = kwargs.get("output_path") or os.path.join(out_dir, f"image_{i}.png")
                r = requests.get(url, timeout=60)
                r.raise_for_status()
                with open(fname, "wb") as f:
                    f.write(r.content)
                saved.append(fname)
                print(f"[jimeng_image] 已保存: {fname}", file=sys.stderr)
        elif binary_data:
            for i, b64 in enumerate(binary_data):
                if not b64:
                    continue
                fname = kwargs.get("output_path") or os.path.join(out_dir, f"image_{i}.png")
                with open(fname, "wb") as f:
                    f.write(base64.b64decode(b64))
                saved.append(fname)
                print(f"[jimeng_image] 已保存 (base64): {fname}", file=sys.stderr)

        if not saved:
            raise RuntimeError("即梦未返回任何图片数据")

        return saved
