from utils.image_utils import UrlTest
from utils.caption_utils import ImageCaptioning
from utils.topic_generation import TopicGenerator
import streamlit as st

topic_generator = TopicGenerator()
img_caption = ImageCaptioning()
path_out = UrlTest()


def return_caption(img):
    capt = img_caption.get_caption(img)
    st.image(image=img, caption=capt, width=250, height=250)
    generated_topics = topic_generator.generate_topics(capt)
    st.write(f"Topic: {generated_topics}")


def main():
    st.title("Topic/Title Generator")

    # User input
    user_input = st.selectbox(label="Text Input or Image Input", options=["Text", "Image", "Image URL"])
    if user_input == "Text":
        text_input = st.text_input(label="Put in your Idea, Let's generate a matching Topic/Title 🤗🤗")
        generated_topics = topic_generator.generate_topics(text_input)
        st.write(f"Topic:{generated_topics}")
    elif user_input == "Image":
        img_input = st.file_uploader(label="Drop an Image you have been admiring, Let's see what we can do🤔🤔",
                                     type=["jpg", "png", "jpeg"],
                                     accept_multiple_files=True)
        for in_img in img_input:
            if in_img is not None:
                img_loaded = path_out.load_image(in_img)
                return_caption(img_loaded)

    elif user_input == "Image URL":
        url_input = st.text_input(label="Do you have a link to the Image you would like to drop, Go Ahead and We got "
                                        "you covered😉😉")
        url_img = path_out.check_url(url_input)
        return_caption(url_img)


if __name__ == "__main__":
    main()