import aiohttp
import asyncio
import json
from typing import Dict, Any


class ModalImageGenerator:
    def __init__(self, session: aiohttp.ClientSession):
        self.session = session
        self.base_url = (
            "https://allanwatts705--kolors-app-model-web-inference.modal.run/"
        )

    async def generate_image(
        self, prompt: str, reference_image_url: str, ip_adapter_weight: float = 0.4
    ) -> Dict[str, Any]:
        params = {
            "ip_adapter_weight": ip_adapter_weight,
            "reference_image_url": reference_image_url,
            "prompt": prompt,
        }

        async with self.session.get(self.base_url, params=params) as response:
            if response.status == 200:
                return await response.json()
            else:
                raise Exception(f"Failed to generate image: {response.status}")

    async def wait_for_image(
        self, result: Dict[str, Any], max_attempts: int = 60, delay: int = 5
    ) -> Dict[str, Any]:
        attempts = 0
        while attempts < max_attempts:
            if result.get("image", {}).get("url"):
                return result
            await asyncio.sleep(delay)
            attempts += 1

        raise Exception("Image generation timed out")

    async def generate_and_wait(
        self, prompt: str, reference_image_url: str, ip_adapter_weight: float = 0.4
    ) -> str:
        result = await self.generate_image(
            prompt, reference_image_url, ip_adapter_weight
        )
        final_result = await self.wait_for_image(result)
        return final_result["image"]["url"]