import os
import pickle
from pathlib import Path
import streamlit as st
import streamlit_authenticator as stauth
import traceback
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
)


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}
        ]

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

    if prompt := st.chat_input():
        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)
        except:
            st.error("An error occurred. Please try again.")
            traceback.print_exc()