import base64
import io

import numpy as np
import requests
import torch
from PIL import Image

from .aimlapi_shared import aspect_ratio_to_size, tensors_to_data_uri_list, validate_api_key


class AIMLAPIImageEditor:
    """
    Image-to-Image editing using Google Nano Banana Pro Edit or ByteDance Seedream 4.5.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "api_key": ("STRING", {
                    "multiline": False,
                    "default": "",
                    "password": True,
                    "tooltip": "Your AI/ML API key from https://aimlapi.com/app/keys",
                }),
                "model": (["google/nano-banana-pro-edit", "bytedance/seedream-4-5"], {
                    "default": "google/nano-banana-pro-edit",
                    "tooltip": "Select the image editing model",
                }),
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": "Transform the scene to autumn with colorful falling leaves",
                    "tooltip": "Text description of how to edit the image",
                }),
                "image": ("IMAGE", {
                    "tooltip": "Input image(s) to edit (supports batching)",
                }),
            },
            "optional": {
                "aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "21:9", "5:4", "4:5"], {
                    "default": "1:1",
                    "tooltip": "Aspect ratio (Nano Banana only)",
                }),
                "resolution": (["1K", "2K", "4K"], {
                    "default": "1K",
                    "tooltip": "Output image resolution",
                }),
                "num_images": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 4,
                    "step": 1,
                    "tooltip": "Number of images to generate (Nano Banana only, 1-4)",
                }),
                "seed": ("INT", {
                    "default": -1,
                    "min": -1,
                    "max": 2147483647,
                    "tooltip": "Seed for reproducibility (Seedream only, -1 for random)",
                }),
                "guidance_scale": ("FLOAT", {
                    "default": 3.5,
                    "min": 0.0,
                    "max": 20.0,
                    "step": 0.1,
                    "tooltip": "Guidance scale (Seedream only)",
                }),
                "watermark": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Add invisible watermark (Seedream only)",
                }),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images",)
    FUNCTION = "edit_image"
    CATEGORY = "AI/ML API"

    def edit_image(self, api_key, model, prompt, image,
                   aspect_ratio="1:1", resolution="1K", num_images=1,
                   seed=-1, guidance_scale=3.5, watermark=False):
        api_key = validate_api_key(api_key)

        url = "https://api.aimlapi.com/v1/images/generations"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

        image_urls = tensors_to_data_uri_list(image, limit=14)
        if len(image_urls) > 14:
            image_urls = image_urls[:14]

        payload = {
            "model": model,
            "prompt": prompt,
            "image_urls": image_urls,
            "response_format": "url",
        }

        if model == "google/nano-banana-pro-edit":
            payload["aspect_ratio"] = aspect_ratio
            payload["resolution"] = resolution
            payload["num_images"] = num_images
        elif model == "bytedance/seedream-4-5":
            payload["size"] = aspect_ratio_to_size(aspect_ratio, resolution)
            payload["watermark"] = watermark
            payload["guidance_scale"] = guidance_scale
            if seed != -1:
                payload["seed"] = seed

        try:
            response = requests.post(url, headers=headers, json=payload, timeout=120)
            if response.status_code != 200:
                self._raise_image_error(response)

            result = response.json()
            image_data_list = None
            if "data" in result and result["data"]:
                image_data_list = result["data"]
            elif "images" in result and result["images"]:
                image_data_list = result["images"]

            if not image_data_list:
                raise ValueError(f"No images returned from API. Response: {result}")

            images = []
            for img_data in image_data_list:
                output_image = self._download_image(img_data)
                image_np = np.array(output_image).astype(np.float32) / 255.0
                image_tensor = torch.from_numpy(image_np)[None,]
                images.append(image_tensor)

            output = torch.cat(images, dim=0) if len(images) > 1 else images[0]
            return (output,)

        except requests.exceptions.RequestException as exc:
            raise ValueError(f"Network error: {exc}")

    @staticmethod
    def _raise_image_error(response):
        error_msg = f"API request failed with status {response.status_code}"
        try:
            error_data = response.json()
            if "error" in error_data:
                error_msg += f": {error_data['error']}"
            elif "message" in error_data:
                error_msg += f": {error_data['message']}"
            else:
                error_msg += f": {error_data}"
        except Exception:
            error_msg += f": {response.text[:300]}"

        if response.status_code == 403:
            error_msg += "\n\nTroubleshooting 403 Forbidden:\n"
            error_msg += "- Verify your API key is correct at https://aimlapi.com/app/keys\n"
            error_msg += "- Check if your API key has sufficient credits\n"
            error_msg += "- Ensure your API key hasn't expired\n"
            error_msg += "- Try generating a new API key"
        elif response.status_code == 401:
            error_msg += "\n\nAPI key is invalid or missing. Get a valid key from https://aimlapi.com/app/keys"
        elif response.status_code == 429:
            error_msg += "\n\nRate limit exceeded. Please wait a moment and try again."
        elif response.status_code >= 500:
            error_msg += "\n\nAPI server error. Please try again later."

        raise ValueError(error_msg)

    @staticmethod
    def _download_image(img_data):
        if "url" in img_data:
            img_response = requests.get(img_data["url"], timeout=30)
            img_response.raise_for_status()
            image = Image.open(io.BytesIO(img_response.content))
        elif "b64_json" in img_data and img_data["b64_json"]:
            image_bytes = base64.b64decode(img_data["b64_json"])
            image = Image.open(io.BytesIO(image_bytes))
        else:
            raise ValueError(f"No image URL or base64 data in response: {img_data}")

        if image.mode != "RGB":
            image = image.convert("RGB")
        return image