import gradio as gr
from gradio.components import label

from model import MLModel
from dotenv import load_dotenv
from os.path import join, dirname

dotenv_path = join(dirname(__file__), ".env")
load_dotenv(dotenv_path)


URL = "localhost:8000"
model = MLModel()


def save(image_path: str):
    try:
        id = model.save(image_path)
    except Exception as e:
        return str(e), None
    return f"保存しました。ID: {id}", None


def search(prompt: str):
    urls = model.search(prompt)
    result = [None] * 5
    for i, url in enumerate(urls[:5]):
        result[i] = url

    return result


with gr.Blocks() as app:
    # Rowでレイアウトを定義
    gr.Label("自然言語画像検索(左から検索ランキングが高い)")
    input = gr.Textbox(placeholder="可愛いワンコ", label="検索")
    with gr.Row():
        output1 = gr.Image(type="filepath", label="ランキング1")
        output2 = gr.Image(type="filepath", label="ランキング2")
        output3 = gr.Image(type="filepath", label="ランキング3")
        output4 = gr.Image(type="filepath", label="ランキング4")
        output5 = gr.Image(type="filepath", label="ランキング5")
    btn = gr.Button("検索")
    btn.click(
        fn=search, inputs=input, outputs=[output1, output2, output3, output4, output5]
    )

    gr.Label("画像保存")
    save_image = gr.Image(type="filepath")
    result = gr.Text(label="upload結果")
    save_btn = gr.Button("保存")
    save_btn.click(fn=save, inputs=save_image, outputs=[result, save_image])
app.launch()