from openrouter.components import ChatMessageContentItemTypedDict
from replicate import Client as ReplicateClient
from typing import Any, Dict, List

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

REPLICATE_MODEL = "ultracoderru/nova-anime-xl-17:8f702486aa2852a08564ede8c83a7f58e52c83f6698e7be0e061d79c113dc88b"


class GenerateImageAnimeTool(Tool):
    def __init__(self, replicate_token: str):
        self._client = ReplicateClient(api_token=replicate_token)
    
    @property
    def name(self) -> str:
        return "generate_image_anime"
    
    @property
    def description(self) -> str:
        return "Генерация изображения в стиле аниме по описанию"
    
    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Положительный запрос"
                },
                "negative_prompt": {
                    "type": "string",
                    "description": "Отрицательный запрос"
                },
                "aspect_ratio": {
                    "type": "string",
                    "enum": ["1:1", "4:3", "3:4", "16:9", "9:16", "9:20"],
                    "description": "Соотношение сторон"
                }
            },
            "required": ["prompt", "negative_prompt"]
        }
    
    async def execute(self, args: Dict[str, Any], artifacts: Dict[str, Any]) -> List[ChatMessageContentItemTypedDict]:
        prompt = args.get("prompt", "")
        negative_prompt = args.get("negative_prompt", "")
        aspect_ratio = args.get("aspect_ratio", None)
        
        aspect_ratio_resolution_map = {
            "1:1": (1280, 1280),
            "4:3": (1280, 1024),
            "3:4": (1024, 1280),
            "16:9": (1280, 720),
            "9:16": (720, 1280),
            "9:20": (720, 1600)
        }
        width, height = aspect_ratio_resolution_map.get(aspect_ratio, (1280, 1024))
        print(f"Генерация аниме-изображения {width}x{height}:\n+ {prompt}\n- {negative_prompt}")
        
        arguments = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "add_recommended_tags": False,
            "width": width,
            "height": height,
            "guidance_scale": 4.5,
            "num_inference_steps": 20,
            "hires_enable": True,
            "hires_num_inference_steps": 30,
            "disable_safety_checker": True
        }
        
        try:
            outputs = await self._client.async_run(REPLICATE_MODEL, input=arguments)
            artifacts["generated_image_hires"] = await outputs[0].aread()
            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}")