import os
import json
import asyncio
from typing import Optional, Dict
from taskflow import FileIOHandler
from api_modules.ark_image_client_async import AsyncArkImageClient
from taskflow import get_logger

io_handler = FileIOHandler()
logger = get_logger("examples.video_create.mcps.character_portrait")

prompt_portrait = \
"""
根据以下描述生成角色{identifier}的全身三视图肖像，背景为纯白色。从左往右依次为角色{identifier}的正面、侧面、背面；尽可能占满画面。
角色描述：{features}
风格：{style}
保持人物身材比例协调性，不要出现畸形。
"""

prompt_refer_portrait = \
"""
请生成图中人物的全身三视图肖像，背景为纯白色，从左往右依次为角色{identifier}的正面、侧面、背面；尽可能占满画面。
风格：{style}
保持人物身材比例协调性，不要出现畸形。
"""

async def gen_single_character_portrait(
    client: AsyncArkImageClient,
    character: dict,
    size: Optional[str] = "2048x2048",
    refer_image: Optional[list[str]] = None,
    style: Optional[str] = None
) -> str:
    features = "(静态特征)" + character.get("static_features", "") + ";(动态特征)" + character.get("dynamic_features", "")
    if refer_image is not None:
        prompt = prompt_refer_portrait.format(identifier=character["identifier_in_scene"], style=style)
    else:
        prompt = prompt_portrait.format(identifier=character["identifier_in_scene"], features=features, style=style)
    response = await client.generate_image(prompt=prompt, size=size, reference_image=refer_image)
    image_url = client.get_image_url(response)
    return image_url

async def gen_character_portrait(
    characters: dict,
    size: Optional[str] = "2048x2048",
    refer_image: Optional[list[str]] = None,
    style: Optional[str] = None
) -> list[str]:
    image_urls = []
    for character in characters["characters"]:
        image_url = await gen_single_character_portrait(client=client, character=character, size=size, refer_image=refer_image, style=style)
        image_urls.append(image_url)
    return image_urls

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

    character = characters["characters"][1]
    print(character)

    async def main():
        async with AsyncArkImageClient() as client:
            image_url = await gen_single_character_portrait(client=client, character=character, refer_image=None, style="写实")
            print(image_url)

    asyncio.run(main())