"""MiniMax image adapter — POST /image_generation, base64 decode to save."""
import base64
import json
import os
import sys

import requests


class MiniMaxImageAdapter:
    @staticmethod
    def _ratio_to_size(ratio_str, max_dim=2048):
        """Parse '16:9' style ratio and compute (width, height) capped at max_dim."""
        parts = ratio_str.split(":")
        if len(parts) != 2:
            return 1024, 1024
        try:
            rw, rh = float(parts[0]), float(parts[1])
        except ValueError:
            return 1024, 1024
        # Scale so the longer side = max_dim
        scale = max_dim / max(rw, rh)
        w = round(rw * scale)
        h = round(rh * scale)
        return w, h

    def generate(self, prompt, config, **kwargs):
        base_url = config["base_url"].rstrip("/")
        api_key = config["api_key"]
        model = config["model"]
        out_dir = kwargs["out_dir"]

        # 默认 1024x1024；传了 aspect_ratio 则按 2048 上限换算 width/height
        aspect_ratio = kwargs.get("aspect_ratio")
        if aspect_ratio and aspect_ratio != "1:1":
            w, h = self._ratio_to_size(aspect_ratio, max_dim=2048)
            size_params = {"width": w, "height": h}
        else:
            size_params = {"width": 1024, "height": 1024}

        endpoint = f"{base_url}/image_generation"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": model,
            "prompt": prompt,
            "response_format": "base64",
            "n": 1,
            **size_params,
        }

        print(f"[minimax_image] generating with model={model} size={size_params}...", file=sys.stderr)
        resp = requests.post(endpoint, headers=headers, json=payload, timeout=180)
        data = resp.json()

        if data.get("base_resp", {}).get("status_code") != 0:
            raise RuntimeError(f"MiniMax error: {data.get('base_resp', {}).get('status_msg')}")

        image_list = data.get("data", {}).get("image_base64", [])
        if not image_list:
            raise RuntimeError("no image data in response")

        saved = []
        for i, img_b64 in enumerate(image_list):
            decoded = base64.b64decode(img_b64)
            ext = "png" if decoded[:8].startswith(b"\x89PNG") else "jpg"
            fname = kwargs.get("output_path") or os.path.join(out_dir, f"image_{i+1}.{ext}")
            with open(fname, "wb") as f:
                f.write(decoded)
            saved.append(fname)
            print(f"[minimax_image] saved: {fname}", file=sys.stderr)

        return saved
