import fal_client
from PIL import Image
from typing import Dict, Any
import requests
from io import BytesIO

from weave_prompt import TextToImageModel

from dotenv import load_dotenv
load_dotenv()

class FalImageGenerator(TextToImageModel):
    """Handles image generation using fal_client."""
    
    def __init__(self, model_name: str = "fal-ai/flux-pro"):
        self.model_name = model_name
    
    def _on_queue_update(self, update):
        """Handle queue updates during image generation."""
        if isinstance(update, fal_client.InProgress):
            for log in update.logs:
                print(log["message"])
    
    def generate(self, prompt: str, **kwargs) -> Image.Image:
        """Generate an image from a text prompt using fal_client."""
        result = fal_client.subscribe(
            self.model_name,
            arguments={
                "prompt": prompt,
                **kwargs
            },
            with_logs=True,
            on_queue_update=self._on_queue_update,
        )
        print(result)
        
        return self._extract_image_from_result(result)
    
    def _extract_image_from_result(self, result: Dict[str, Any]) -> Image.Image:
        """Extract and download image from fal_client result."""
        if result and 'images' in result and len(result['images']) > 0:
            image_url = result['images'][0]['url']
            response = requests.get(image_url)
            response.raise_for_status()  # Raise an exception for bad status codes
            image = Image.open(BytesIO(response.content))
            return image
        else:
            raise ValueError("No image found in the result")