"""Agnes 原生协议适配器。

模型发现： GET  {base_url}/models
出图：     POST {base_url}/images/generations
尺寸：     size 档位（1K/2K/3K/4K）+ ratio 宽高比
返回：     data[0].b64_json 或 data[0].url
注意：     response_format 必须放在 extra_body 内（顶层不允许）
"""

from __future__ import annotations

import re

from .base import Capabilities, ImageResult, ProviderAdapter, ProviderError

RATIOS = ("1:1", "3:4", "4:3", "16:9", "9:16", "2:3", "3:2", "21:9")

# 统一像素尺寸 → Agnes 档位
_TIER_BY_MAX = (
    (1024, "1K"),
    (1536, "1K"),
    (2048, "2K"),
    (3072, "3K"),
    (4096, "4K"),
)


def size_to_tier_and_ratio(size: str, ratio: str) -> tuple[str, str]:
    """把统一的 WxH（可选带 ratio）翻译成 Agnes 的档位 + 宽高比。"""
    m = re.match(r"^\s*(\d+)\s*[xX×]\s*(\d+)\s*$", size or "")
    tier = "1K"
    derived = ""
    if m:
        w, h = int(m.group(1)), int(m.group(2))
        long_side = max(w, h)
        for limit, t in _TIER_BY_MAX:
            if long_side <= limit:
                tier = t
                break
        else:
            tier = "4K"
        derived = _nearest_ratio(w, h)
    elif size and size.strip().upper() in ("1K", "2K", "3K", "4K"):
        tier = size.strip().upper()
    if ratio:
        r = ratio.strip()
        if r in RATIOS:
            return tier, r
        raise ProviderError("不支持的宽高比: " + r, status=400)
    return tier, (derived or "1:1")


def _nearest_ratio(w: int, h: int) -> str:
    target = w / h if h else 1.0
    best, best_diff = "1:1", 1e9
    for r in RATIOS:
        a, b = r.split(":")
        value = int(a) / int(b)
        diff = abs(value - target)
        if diff < best_diff:
            best, best_diff = r, diff
    return best


class AgnesAdapter(ProviderAdapter):
    name = "agnes"
    capabilities = Capabilities(
        supports_image_input=True,
        size_hint="1K~4K / 8 比例",
        kind="image",
    )

    def generate(
        self,
        model: str,
        prompt: str,
        size: str,
        ratio: str,
        n: int,
        images: list[str],
    ) -> list[ImageResult]:
        tier, r = size_to_tier_and_ratio(size, ratio)
        extra: dict = {"response_format": "b64_json"}
        if images:
            extra["image"] = images
        body: dict = {
            "model": model,
            "prompt": prompt,
            "size": tier,
            "ratio": r,
            "n": n,
            "extra_body": extra,
        }
        resp = self._request("POST", "/images/generations", json_body=body, timeout=600.0)
        payload = self._json(resp)
        results = self.collect_results(payload)
        out: list[ImageResult] = []
        for item in results:
            out.append(self.download(item.url) if item.url else item)
        return out