import os
import pickle
from pathlib import Path
import streamlit as st
import streamlit_authenticator as stauth
from streamlit_feedback import streamlit_feedback
import traceback
import uuid
from dataset_loader import load_dataset
from dotenv import load_dotenv

load_dotenv()

if not os.path.isdir("spsetia_chatbot_data"):
    load_dataset()

from spsetia_chatbot_data.utils import (
    run_conversation,
    record_db,
    hashed_passwords,
    greeting,
    save_feedback
)


st.set_page_config(page_title="S P Setia Chatbot", page_icon=":robot:", layout="wide")


# --- USER AUTHENTICATION ---
names = ["S P Setia Admin"]
usernames = ["spsetia_admin"]

authenticator = stauth.Authenticate(
    names,
    usernames,
    hashed_passwords,
    "spsetia_chatbot",
    "eqiddnksadjw",
    cookie_expiry_days=30,
)

name, authentication_status, username = authenticator.login("Login", "main")

if authentication_status == False:
    st.error("Username/password is incorrect")

if authentication_status == None:
    st.warning("Please enter your username and password")

if authentication_status:
    st.sidebar.title(f"S P Setia Chatbot")
    st.sidebar.header(f"Welcome {name}")
    authenticator.logout("Logout", "sidebar")
    if st.sidebar.button("Refresh Chat"):
        st.session_state.messages = [
            {
                "role": "assistant",
                "content": greeting,
            }
        ]
    if "messages" not in st.session_state:
        st.session_state.messages = [
            {
                "role": "assistant",
                "content": greeting,
            }
        ]

    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    if "fbk" not in st.session_state:
        st.session_state.fbk = str(uuid.uuid4())

    def display_answer():
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.write(message["content"])

    def create_answer(prompt):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

        if st.session_state.messages[-1]["role"] != "assistant":
            try:
                with st.chat_message("assistant"):
                    with st.spinner("Thinking..."):
                        response = run_conversation(
                            st.session_state.messages, prompt, username
                        )
                        qa_id = response[4]
                        prompt_tokens = response[5]
                        completion_tokens = response[6]
                        if response[3]:
                            if isinstance(response[0], list):
                                answer = ""
                                for item in response[0]:
                                    answer += f"""{item[0]}\n"""
                                    st.text(item[0])
                                    for image in item[1]:
                                        if image:
                                            st.image(image)
                                reference = f"""  \n\nReference  \nfolder: {response[2]["folder"]}  \nfile: {response[2]["file"]}"""
                                st.text(reference)
                                answer += reference
                                record_db(
                                    content=answer,
                                    qa_id=qa_id,
                                    msg_mode=1,
                                    recipient_id=username,
                                    prompt_tokens=prompt_tokens,
                                    completion_tokens=completion_tokens,
                                )
                            else:
                                answer = response[0]
                                reference = f"""  \n\nReference  \nfolder: {response[2]["folder"]}  \nfile: {response[2]["file"]}"""
                                answer += reference
                                st.text(answer)
                                record_db(
                                    content=answer,
                                    qa_id=qa_id,
                                    msg_mode=1,
                                    recipient_id=username,
                                    prompt_tokens=prompt_tokens,
                                    completion_tokens=completion_tokens,
                                )
                        else:
                            answer = response[0]
                            st.text(answer)
                            if images := response[1]:
                                if isinstance(images, list):
                                    for image in images:
                                        st.image(image)
                            record_db(
                                content=answer,
                                qa_id=qa_id,
                                msg_mode=1,
                                recipient_id=username,
                                prompt_tokens=prompt_tokens,
                                completion_tokens=completion_tokens,
                            )

                message = {"role": "assistant", "content": answer}
                st.session_state.messages.append(message)
                st.session_state.chat_history.append(
                    {
                        "question": prompt,
                        "answer": answer,
                    }
                )

            except:
                st.error("An error occurred. Please try again.")
                traceback.print_exc()

    display_answer()

    if prompt := st.chat_input():
        create_answer(prompt)

    def fbcb(response):
        data = st.session_state.chat_history[-1]
        question = data["question"]
        answer = data["answer"]
        print(question, answer, response)
        save_feedback(question, answer, response)
        st.success("Thank you for your feedback!")
        st.session_state.fbk = str(uuid.uuid4())

    streamlit_feedback(
        feedback_type="thumbs",
        optional_text_label="[Optional]",
        align="flex-start",
        key=st.session_state.fbk,
        on_submit=fbcb,
    )

# {'type': 'thumbs', 'score': '👎', 'text': 'bad'}
# {'type': 'thumbs', 'score': '👍', 'text': 'good'}