import streamlit as st
from streamlit_option_menu import option_menu
# from chat import Chat
from convo import Convo
from remainder import rem
from transformers import pipeline


visualqna = pipeline(model="dandelin/vilt-b32-finetuned-vqa")


def load_image():
    with st.sidebar:
        if img := st.text_input("Enter Image URL") or st.selectbox("Select Image", ("https://images.unsplash.com/photo-1593466144596-8abd50ad2c52?ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D&auto=format&fit=crop&w=3434&q=80", "https://images.unsplash.com/photo-1566438480900-0609be27a4be?ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D&auto=format&fit=crop&w=3394&q=80")):
            if st.button("Load Image"):
                st.write("Image Uploaded!")
                st.image(img)
        else:
            st.warning("Please enter an image URL and click 'Load Image' before asking a question.")
    return img



def visual_qna():
    st.title("Visual Q&A")
    img = load_image()
    if img:
        if query := st.chat_input("Enter your message"):
            response = visualqna(question=query, image=img)
            with st.chat_message("assistant"):
                st.write(response)
    else:
        st.warning("Please enter an image URL and click 'Load Image' before asking a question.")
    

def homepage():
    st.title("Home")   
    # st.header("Home Page")
    st.subheader("Welcome to the Home Page")

def dashboard():
    # st.title("Dashboard")
    # st.header("Dashboard")

    with st.sidebar:
        selected = option_menu("Menu", ["Home", "Dashboar","Remainder","VISUAL Q&A","Conversation","Logout"])
                    
    if selected == "Home":
        homepage()
        
    elif selected == "Dashboard":
        "gfjfvjhvjhv"
    # elif selected == "Chat":
    #     Chat()
    elif selected == "Conversation": 
        Convo()
    elif selected == "Logout":
        st.session_state["user"] = "visitor"
        st.experimental_rerun()
    elif selected == "Remainder":
        rem()
    elif selected == 'Visual Q&A':
        visual_qna()