from PIL import Image
import gradio as gr

from config import (
    description,
    article,
)
from encode import get_embeddings
from database import search_vector

def search_images(values):
    image = Image.new("RGBA", values["composite"].size, (255, 255, 255, 255))
    image.paste(values["composite"], mask=values["composite"])
    embedding = get_embeddings([image])[0]
    results = search_vector(embedding, limit=100)

    _deduplicated = '\t'.join(dict.fromkeys(result.kanji for result in results))
    # TODO Format the results better
    # Huge boxes using the right font for each of them?

    return f"Results: {_deduplicated}"

# TODO FIND OUT HOW TO CHANGE THE DEFAULT EDITOR TAB?
input_image = gr.ImageEditor(
    label="Write the Kanji you want to search for",
    show_label=False,
    type="pil",
    brush=gr.Brush(
        default_size=16,
        color_mode="fixed",
        colors=["#000000", "#ffffff"],
    ),
)


output_box = gr.Textbox()


demo = gr.Interface(
    fn=search_images,
    inputs=[input_image],
    outputs=output_box,
    title="Kanji Lookup",
    description=description,
    article=article,
    examples="examples",
    # cache_examples=False,
    # live=True,
)

if __name__ == "__main__":
    demo.launch()