import aiohttp
import asyncio
import os
import time
import base64
import aiofiles
from pathlib import Path


class SimpleImageGenerator:
    def __init__(self, output_dir="output/images"):
        self.api_key = os.getenv("OPENAI_API_KEY")
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        if not self.api_key:
            raise ValueError("OPENAI_API_KEY environment variable required")

    async def generate_single_image(self, session, prompt):
        """Make a single API call to generate an image"""
        payload = {
            "model": "gpt-image-1",
            "prompt": prompt,
            "size": "1024x1536",
            "quality": "high",
        }
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        try:
            async with session.post(
                "https://api.openai.com/v1/images/generations",
                json=payload,
                headers=headers,
            ) as response:
                if response.status == 200:
                    base64_data = (await response.json())["data"][0]["b64_json"]
                    return await self._save_base64_image(base64_data)
                print(f"API error {response.status}")
                return None
        except Exception as e:
            print(f"Failed to generate image: {e}")
            return None

    async def _save_base64_image(self, base64_data):
        """Save base64 image data to output directory"""
        timestamp = str(time.time_ns())
        filepath = self.output_dir / f"{timestamp}.png"
        try:
            image_bytes = base64.b64decode(base64_data)
            async with aiofiles.open(filepath, "wb") as f:
                await f.write(image_bytes)
            return str(filepath)
        except Exception as e:
            print(f"Failed to save image: {e}")
            return None

    async def generate_images_parallel(self, prompts):
        """Generate multiple images in parallel using asyncio.gather()"""
        async with aiohttp.ClientSession() as session:
            # Create tasks for all prompts - this is the key parallel part!
            tasks = [self.generate_single_image(session, prompt) for prompt in prompts]

            # Execute all tasks concurrently
            results = await asyncio.gather(*tasks, return_exceptions=True)

            # Filter out failed requests
            successful_files = [
                filepath
                for filepath in results
                if filepath and isinstance(filepath, str)
            ]

            print(f"Generated {len(successful_files)}/{len(prompts)} images")
            return successful_files


# Usage example
async def main():
    generator = SimpleImageGenerator()

    prompts = [
        "A red cat sitting on a blue chair",
        "A sunset over mountains",
        "A robot reading a book",
    ]

    # This will make all 3 API calls simultaneously!
    saved_files = await generator.generate_images_parallel(prompts)

    for i, filepath in enumerate(saved_files):
        print(f"Image {i + 1} saved to: {filepath}")


# Run the example
if __name__ == "__main__":
    asyncio.run(main())