from typing import Any, Dict

from replicate import Client as ReplicateClient

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

REPLICATE_MODEL = "bytedance/seedream-4.5"


class GenerateImageTool(Tool):
    def __init__(self, replicate_token: str):
        self._client = ReplicateClient(api_token=replicate_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]) -> str:
        prompt = args.get("prompt", "")
        aspect_ratio = args.get("aspect_ratio", "4:3")
        print(f"Генерация изображения {aspect_ratio}: {prompt}")

        arguments = {
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
            "disable_safety_checker": True
        }

        try:
            outputs: Any = 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 "Изображение сгенерировано и будет показано пользователю."
        except Exception as e:
            print(f"Ошибка генерации изображения: {e}")
            return f"Не удалось сгенерировать изображение: {e}"