from omegaconf import OmegaConf
import streamlit as st
import os
from PIL import Image
import re
from translate import Translator


from llama_index.indices.managed.vectara import VectaraIndex
from llama_index.core.agent import ReActAgent
from llama_index.llms.openai import OpenAI
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.utils import print_text

learning_styles = ['traditional', 'inquiry based']
languages = {'English': 'en', 'Spanish': 'es', 'French': 'fr', 'German': 'de', 'Arabic': 'ar', 'Chinese': 'zh-cn', 
             'Hebrew': 'he', 'Hindi': 'hi', 'Italian': 'it', 'Japanese': 'ja', 'Korean': 'ko', 'Portuguese': 'pt'}
initial_prompt = "How can I help you today?"

def launch_bot():
    def reset():
        cfg = st.session_state.cfg
        llm = OpenAI(model="gpt-4o", temperature=0)        
        tr_prompt = Translator(to_lang=languages[cfg.language]).translate(initial_prompt)
        print(tr_prompt)
        st.session_state.messages = [{"role": "assistant", "content": tr_prompt, "avatar": "🦖"}]
        vectara = VectaraIndex(vectara_api_key=cfg.api_key, 
                               vectara_customer_id=cfg.customer_id, 
                               vectara_corpus_id=cfg.corpus_id)
        vectara_tool = QueryEngineTool(
            query_engine = vectara.as_query_engine(summary_enabled=True, summary_num_results=5, summary_response_lang = languages[cfg.language],
                                                   summary_prompt_name="vectara-summary-ext-24-05-large"),
            metadata = ToolMetadata(name="Vectara", 
                                    description="Vectara Query Engine that is able to answer any questions about the Justice Harvard class."),
        )
        llm = OpenAI(model="gpt-4o", temperature=0)
        st.session_state.agent = ReActAgent.from_tools(
            tools=[vectara_tool], llm=llm, 
            context = f'''
                You are a teacher assistant at Justice Harvard course. You are helping a student with his questions.
                The student is student who is {cfg.student_age} years old, you personalize your assistance to the student's age, 
                and rephrase your answer if needed to fit the {cfg.style} learning style.
                ''',
            verbose=True
        )

    if 'cfg' not in st.session_state:
        cfg = OmegaConf.create({
            'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
            'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
            'api_key': str(os.environ['VECTARA_API_KEY']),
            'style': learning_styles[0],
            'language': 'English',
            'student_age': 21
        })
        st.session_state.cfg = cfg
        st.session_state.style = learning_styles[0]
        st.session_state.language = 'English'
        st.session_state.student_age = 21
        reset()

    cfg = st.session_state.cfg
    st.set_page_config(page_title="Teaching Assistant", layout="wide")

    # left side content
    with st.sidebar:
        image = Image.open('Vectara-logo.png')
        st.image(image, width=250)
        st.markdown(f"## Welcome to Justice Harvard.\n\n\n")

        st.markdown("\n")
        cfg.style = st.selectbox('Learning Style:', learning_styles)
        if st.session_state.style != cfg.style:
            st.session_state.style = cfg.style
            reset()

        st.markdown("\n")
        cfg.language = st.selectbox('Language:', languages.keys())
        if st.session_state.language != cfg.language:
            st.session_state.langage = cfg.language
            reset()

        st.markdown("\n") 
        cfg.student_age = st.number_input(
            'Student age:',  min_value=13, value=cfg.student_age,
            step=1, format='%i'
        )
        if st.session_state.student_age != cfg.student_age:
            st.session_state.student_age = cfg.student_age
            reset()

        st.markdown("\n\n")
        if st.button('Start Over'):
            reset()

        st.markdown("---")
        st.markdown(
            "## How this works?\n"
            "This app was built with [Vectara](https://vectara.com).\n\n"
            "It demonstrates the use of Agentic Chat functionality with Vectara"
        )
        st.markdown("---")


    if "messages" not in st.session_state.keys():
        reset()

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"], avatar=message["avatar"]):
            st.write(message["content"])

    # User-provided prompt
    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt, "avatar": '🧑‍💻'})
        with st.chat_message("user", avatar='🧑‍💻'):
            print_text(f"Starting new question: {prompt}\n", color='green')
            st.write(prompt)

    # Generate a new response if last message is not from assistant
    if st.session_state.messages[-1]["role"] != "assistant":
        with st.chat_message("assistant", avatar='🤖'):
            with st.spinner('Thinking...'):
                res = st.session_state.agent.chat(prompt)
                cleaned = re.sub(r'\[\d+\]', '', res.response)
                st.write(cleaned)
            message = {"role": "assistant", "content": cleaned, "avatar": '🤖'}
            st.session_state.messages.append(message)
    
if __name__ == "__main__":
    launch_bot()