from fal_client import AsyncClient as FalClient
from openrouter.components import ChatMessageContentItemTypedDict
from typing import Any, Dict, List

from ai.tool import Tool
from ai.utils import *

FAL_MODEL = "fal-ai/bytedance/seedream/v4.5/text-to-image"


class GenerateImageTool(Tool):
    def __init__(self, fal_token: str):
        self._client = FalClient(key=fal_token)
    
    @property
    def name(self) -> str:
        return "generate_image"
    
    @property
    def description(self) -> str:
        return "Генерация изображения по описанию"
    
    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Подробное описание сцены на английском языке БЕЗ технических параметров "
                                   "(соотношение сторон, разрешение)"
                },
                "aspect_ratio": {
                    "type": "string",
                    "enum": ["1:1", "4:3", "3:4", "16:9", "9:16"],
                    "description": "Соотношение сторон"
                }
            },
            "required": ["prompt"]
        }
    
    async def execute(self, args: Dict[str, Any], artifacts: Dict[str, Any]) -> List[ChatMessageContentItemTypedDict]:
        prompt = args.get("prompt", "")
        aspect_ratio = args.get("aspect_ratio", None)
        
        aspect_ratio_size_map = {
            "1:1": "square",
            "4:3": "landscape_4_3",
            "3:4": "portrait_4_3",
            "16:9": "landscape_16_9",
            "9:16": "portrait_16_9",
            "9:20": "portrait_16_9"
        }
        image_size = aspect_ratio_size_map.get(aspect_ratio, "landscape_4_3")
        print(f"Генерация изображения {image_size}: {prompt}")
        
        arguments = {
            "prompt": prompt,
            "image_size": image_size,
            "enable_safety_checker": False
        }
        
        try:
            result = await self._client.run(FAL_MODEL, arguments=arguments)
            if "images" not in result:
                raise RuntimeError("Неожиданный ответ от сервера.")
            image_url = result["images"][0]["url"]
            
            from utils import download_file
            artifacts["generated_image_hires"] = await download_file(image_url)
            artifacts["generated_image"] = compress_image(artifacts["generated_image_hires"], 1280)
            
            return serialize_message_content(
                text="Изображение сгенерировано и будет показано пользователю.",
                image=None
            )
        except Exception as e:
            print(f"Ошибка генерации изображения: {e}")
            return serialize_message_content(text=f"Не удалось сгенерировать изображение: {e}")