import os
import time
import gradio as gr
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext, load_index_from_storage, PromptTemplate
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
from llama_index.core.llms import ChatMessage, MessageRole
from llama_index.core.chat_engine.types import ChatMode
import base64

#from theme import CustomTheme

#Trainigsdaten
path_modulhandbuch = "./hm_daten"
path_persist = os.path.join(path_modulhandbuch, "persist")

Settings.llm = OpenAI(temperature=0.1, model="gpt-4o-mini")

if not os.path.exists(path_persist):
    documents = SimpleDirectoryReader("./hm_daten/").load_data()
    index = VectorStoreIndex.from_documents(documents)
    index.storage_context.persist(persist_dir=path_persist)
else:
    storage_context = StorageContext.from_defaults(persist_dir=path_persist)
    index = load_index_from_storage(storage_context)



#prompt für anfragen
custom_system_prompt = """
"We have provided context information below. \n"
#    "---------------------\n"
#    "{context_str}"
#    "\n---------------------\n"
You are an expert assistant providing detailed and accurate information.
You help students finding important infromation about their university.
You always answer professional but friendly, encouraging, youthful und funny.
If the question is in german give answer in german , else give answer in english: {query_str}
"""

chat_engine = index.as_chat_engine(
chat_mode=ChatMode.CONDENSE_PLUS_CONTEXT,
system_prompt=custom_system_prompt,
streaming=True
)


def response(message, history):
    chat_history = []
    for i, msg in enumerate(history):
        if i % 2 == 0:
            history_message = ChatMessage(role=MessageRole.ASSISTANT, content=msg["content"])
        else:
            history_message = ChatMessage(role=MessageRole.USER, content=msg["content"])
        chat_history.append(history_message)

    streaming_response = chat_engine.stream_chat(message, chat_history=chat_history)

    answer = ""
    for text in streaming_response.response_gen:
        time.sleep(0.05)
        answer += text
        yield answer


#theme = CustomTheme()   

#background
with open("./images/bg_hell.jpg", "rb") as image_file:
    encoded_string = base64.b64encode(image_file.read()).decode()


custom_css = f"""
.gradio-container {{
background: url("data:image/jpeg;base64,{encoded_string}") !important;
background-size: cover !important;
background-position: center !important;
max-width: 100% !important;
height: auto !important;
}}
"""
#test01

def main():
    with gr.Blocks(css = custom_css, css_paths = "./style.css") as demo:
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                gr.Image("./images/hochi6.JPG", show_label = False, show_download_button = False, show_share_button = False, show_fullscreen_button = False)
            with gr.Column(scale=8):
                gr.Markdown("")
        with gr.Row(equal_height=True):
            with gr.Column(scale=1, variant = "default"):
                gr.Markdown("")
            with gr.Column(scale=1):
                gr.Image("./images/scroll_test.jpg",  show_label = False, show_download_button = False, show_share_button = False, show_fullscreen_button = False)
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    value=[{"role": "assistant", "content": "Hi. Du schon wieder. wie kann ich dir helfen?"}],
                    type="messages",
                    show_label=False,
                    avatar_images=("./images/avatar_images/human_2.png", "./images/avatar_images/hochi.PNG"),
                    elem_id="CHATBOT"
                )
                chat_interface = gr.ChatInterface(
                fn=response,
                chatbot=chatbot,
                type="messages"
                )
            with gr.Column(scale=1, variant = "default"):
                gr.Markdown("")
            
        
    demo.launch(inbrowser= True)

main()