from app.utils.chatgpt_prompt import ChatgptService
from app.utils.prompts.prompts import image_autotag_prompt
from app.utils.prompts.schemas import SpaceType, Subject, Camera, MotionRecommended

import asyncio

async def autotag_image(image_url : str) -> list[str]: #tag_list
    chatgpt = ChatgptService()
    image_input_data = {
        "img_url" : image_url,
        "space_type" : list(SpaceType),
        "subject" : list(Subject),
        "camera" : list(Camera),
        "motion_recommended" : list(MotionRecommended)
    }

    image_result = await chatgpt.generate_structured_output(image_autotag_prompt, image_input_data, image_url, False)   
    return image_result

async def autotag_images(image_url_list : list[str]) -> list[dict]: #tag_list
    chatgpt = ChatgptService()
    image_input_data_list = [{ 
        "img_url" : image_url,
        "space_type" : list(SpaceType),
        "subject" : list(Subject),
        "camera" : list(Camera),
        "motion_recommended" : list(MotionRecommended)
    }for image_url in image_url_list]
    
    image_result_tasks = [chatgpt.generate_structured_output(image_autotag_prompt, image_input_data, image_input_data['img_url'], False, silent = True) for image_input_data in image_input_data_list] 
    image_result_list = await asyncio.gather(*image_result_tasks, return_exceptions=True)
    MAX_RETRY = 3 # 하드코딩, 어떻게 처리할지는 나중에
    for _ in range(MAX_RETRY):
        failed_idx = [i for i, r in enumerate(image_result_list) if isinstance(r, Exception)]
        print("Failed", failed_idx)
        if not failed_idx:
            break
        retried = await asyncio.gather(
            *[chatgpt.generate_structured_output(image_autotag_prompt, image_input_data_list[i], image_input_data_list[i]['img_url'], False, silent=True) for i in failed],
            return_exceptions=True
        )
        for i, result in zip(failed_idx, retried):
            image_result_list[i] = result

    print("Failed", failed_idx)
    return image_result_list