import gradio as gr
from utils import extract_data, remove_id_and_ext
from prodiapy import Custom
import os

pipe = Custom(os.getenv('PRODIA_API_KEY'))


model_list = pipe.constant("/sd/models")
model_names = {}

for model_name in model_list:
    name_without_ext = remove_id_and_ext(model_name)
    model_names[name_without_ext] = model_name


def update_btn_start():
    return [
        gr.update(visible=False),
        gr.update(visible=True)
    ]


def update_btn_end():
    return [
        gr.update(visible=True),
        gr.update(visible=False)
    ]


def switch_to_t2i():
    return gr.Tabs.update(selected="t2i")


def send_to_txt2img(image):
    try:
        text = image.info['parameters']
        data = extract_data(text)

        if data['model'] in model_names:
            model = gr.update(value=model_names[data['model']])
        else:
            model = gr.update()

        result = [
            gr.update(value=data['prompt']),
            gr.update(value=data['negative_prompt']) if data['negative_prompt'] is not None else gr.update(),
            gr.update(value=int(data['steps'])) if data['steps'] is not None else gr.update(),
            gr.update(value=int(data['seed'])) if data['seed'] is not None else gr.update(),
            gr.update(value=float(data['cfg_scale'])) if data['cfg_scale'] is not None else gr.update(),
            gr.update(value=int(data['w'])) if data['w'] is not None else gr.update(),
            gr.update(value=int(data['h'])) if data['h'] is not None else gr.update(),
            gr.update(value=data['sampler']) if data['sampler'] is not None else gr.update(),
            model
        ]
        return result

    except Exception as e:
        print(e)
        return