import os
import json
import asyncio
from typing import Optional
from utils.tools import string_to_json, save_json_file
from tools.banana_pro import generate_image_from_prompt_and_images
# from tools.gemini import generate_image_from_prompt_and_images
from tools.text_generator import media_captioner
from tools.image_generator import image_generator

prompt_front = \
"""
根据以下描述生成角色{identifier}的全身正面肖像，背景为纯白色。角色应位于画面中央，占据画面大部分空间。目光直视前方，双臂自然垂于身体两侧。表情自然。
特征：{features}
风格：{style}
"""

prompt_side = \
"""
根据提供的正面肖像生成角色{identifier}的全身侧面肖像，背景为纯白色。角色应位于图像中央，占据画面大部分空间。面向左侧。自然站立，双臂放松垂于身体两侧。
"""

prompt_back = \
"""
根据提供的正面肖像生成角色{identifier}的全身背面肖像，背景为纯白色。角色应位于图像中央，占据画面大部分空间。不应显示任何面部特征。自然站立，双臂放松垂于身体两侧。
"""


class CharacterPortraitsGenerator:
    def __init__(self):
        pass

    def generate_front_portrait(
        self,
        character: dict,
        refer_image_path: list[str] = None,
        style: Optional[str] = None
    ):
        features = "(静态特征)" + character.get("static_features", "") + ";(动态特征)" + character.get("dynamic_features", "")
        prompt = prompt_front.format(identifier=character["identifier_in_scene"], features=features, style=style)

        # result = await image_generator.generate_without_refer(
        #     prompt=prompt,
        # )

        # result.save_url("./output_front.png")
        result = generate_image_from_prompt_and_images(
            prompt=prompt,
            image_paths=refer_image_path if refer_image_path else []
        )

        return result

    def generate_side_portrait(
        self,
        character: dict,
        front_image_path: list[str],
    ):
        prompt = prompt_side.format(identifier=character["identifier_in_scene"])
        
        print(f"参考图片：{front_image_path}")
        # result = await image_generator.generate(
        #     prompt=prompt,
        #     image_url=front_image_path
        # )

        # result.save_url("./output_side.png")
        result = generate_image_from_prompt_and_images(
            prompt=prompt,
            image_paths=front_image_path
        )

        return result

    def generate_back_portrait(
        self,
        character: dict,
        front_image_path: list[str],
    ):
        prompt = prompt_back.format(identifier=character["identifier_in_scene"])
        
        print(f"参考图片：{front_image_path}")
        # result = await image_generator.generate(
        #     prompt=prompt,
        #     image_url=front_image_path
        # )

        # result.save_url("./output_back.png")
        result = generate_image_from_prompt_and_images(
            prompt=prompt,
            image_paths=front_image_path
        )

        return result

    def generate_full_portraits(
        self,
        character: dict,
        style: Optional[str] = None
    ):
        # front_portrait = await self.generate_front_portrait(character, style)
        # side_portrait = await self.generate_side_portrait(character, [front_portrait.data])
        # back_portrait = await self.generate_back_portrait(character, [front_portrait.data])
        front_portrait = self.generate_front_portrait(character, style)
        side_portrait = self.generate_side_portrait(character, [front_portrait.data])
        back_portrait = self.generate_back_portrait(character, [front_portrait.data])

        return {
            "front_portrait": front_portrait,
            "side_portrait": side_portrait,
            "back_portrait": back_portrait
        }

character_portraits_generator = CharacterPortraitsGenerator()

if __name__ == "__main__":
    
    
    with open("characters.json", "r", encoding="utf-8") as f:
        characters = json.load(f)

    character = characters["characters"][0]

    # front_image_url = asyncio.run(character_portraits_generator.generate_front_portrait(character, style="写实"))
    # side_image_url = asyncio.run(character_portraits_generator.generate_side_portrait(character, [front_image_url.data]))
    # back_image_url = asyncio.run(character_portraits_generator.generate_back_portrait(character, [front_image_url.data]))
    # print(front_image_url, side_image_url, back_image_url)

    full_portraits = asyncio.run(character_portraits_generator.generate_full_portraits(character, style="写实"))
    print(full_portraits)