import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import json
from PIL import Image
import io
import base64
import numpy as np
import torch

class OneThingAILoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "api_key": ("STRING", {"default": "", "multiline": False}),
                "prompt": ("STRING", {"default": "", "multiline": True}),
                "model": ("STRING", {"default": "gpt-image-1"}),
                "num_images": ("INT", {"default": 1, "min": 1, "max": 10}),
                "image_size": (["1536x1024", "1024x1024", "1024x1536", "自定义"], {"default": "1024x1024"}),
                "custom_width": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64}),
                "custom_height": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64}),
                "quality": (["low", "medium", "high"], {"default": "medium"}),
                "retries": ("INT", {"default": 3, "min": 1, "max": 10}),
                "timeout": ("INT", {"default": 120, "min": 5, "max": 180}),
            },
            "optional": {
                "reference_image": ("IMAGE",),
                "reference_image_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
            }
        }
    
    RETURN_TYPES = ("IMAGE", )
    FUNCTION = "generate"
    CATEGORY = "OneThingAI/image generation"

    def generate(self, api_key, prompt, model="gpt-image-1", num_images=1, image_size="1024x1024", custom_width=1024, custom_height=1024, quality="medium", retries=3, timeout=8, reference_image=None, reference_image_weight=0.5):
        # API endpoint
        url = "https://api-model.onethingai.com/v1/images/generations"
        
        # Setup retry strategy
        retry_strategy = Retry(
            total=retries,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        
        # Create session with retry strategy
        session = requests.Session()
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("https://", adapter)
        
        # Request headers
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "User-Agent": "OneThingAI ComfyUI/1.0"
        }
        
        # Determine image size based on selection
        if image_size == "自定义":
            image_size = f"{custom_width}x{custom_height}"

        # Request payload
        payload = {
            "model": model,
            "prompt": prompt,
            "n": num_images,
            "size": image_size,
            "quality": quality,
        }

        # Add reference image if provided
        if reference_image is not None:
            # Convert the first image if it's a batch
            if len(reference_image.shape) == 4:
                image_to_encode = reference_image[0]
            else:
                image_to_encode = reference_image

            # Convert from torch tensor to PIL Image
            image_to_encode = (image_to_encode.cpu().numpy() * 255).astype(np.uint8)
            if len(image_to_encode.shape) == 3:
                image_to_encode = Image.fromarray(image_to_encode)
            
            # Convert to base64
            buffered = io.BytesIO()
            image_to_encode.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            
            # Add to payload
            payload["reference_image"] = img_str
            payload["reference_image_weight"] = reference_image_weight
            
        try:
            # Make API request with timeout
            response = session.post(
                url, 
                headers=headers, 
                json=payload,
                timeout=timeout
            )
            response.raise_for_status()
            
            # Parse response
            result = response.json()
            
            if "data" not in result:
                raise ValueError(f"Unexpected API response: {result}")
            
            # Process images
            images = []
            for img_data in result["data"]:
                if "b64_json" in img_data:
                    # Decode base64 image
                    image_bytes = base64.b64decode(img_data["b64_json"])
                    image = Image.open(io.BytesIO(image_bytes))
                    
                    # Convert to RGB if necessary
                    if image.mode != "RGB":
                        image = image.convert("RGB")

                    # Convert to numpy array
                    image_array = np.array(image)
                    images.append(image_array)
                    
            # Stack images if multiple were generated
            if len(images) > 1:
                final_image = np.stack(images)
            else:
                final_image = images[0]
            
            # Convert numpy array to torch tensor
            final_image = torch.from_numpy(final_image).float()/255.0
            final_image = final_image.unsqueeze(0)
            
            return (final_image,)
            
        except requests.exceptions.RequestException as e:
            raise RuntimeError(f"API request failed: {str(e)}")
        except json.JSONDecodeError:
            raise RuntimeError("Failed to parse API response")
        except Exception as e:
            raise RuntimeError(f"Error generating images: {str(e)}")
        finally:
            session.close()