from igf_api.IgfClient import IgfClient
from piapi_ai import client
import asyncio
from dotenv import dotenv_values
from telegram import Bot
from ai_bot.handlers.FluxHandler import FluxHandler

config = dotenv_values(".env")


def send_msg(igf_task_dict: dict, piapi_task_dict: dict):
    if piapi_task_dict['data']['output'] is not None:
        if piapi_task_dict['data']['model'] == "Qubico/flux1-dev":
            asyncio.run(FluxHandler.send_task_result(igf_task=igf_task_dict, piapi_task=piapi_task_dict))


if __name__ == "__main__":
    igf_client = IgfClient()
    tasks = igf_client.piapiTask.get_new_tasks()

    for task in tasks:
        piapi_task = client.flux.get_task(task['task_id'])
        send_msg(igf_task_dict=task, piapi_task_dict=piapi_task)