import streamlit as st
import uuid
from typing import List, Dict
import os
from dataset_loader import load_dataset

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

from r_logic_data.agent import Agent


class StreamlitChatbot:
    def __init__(self):
        self.agent = Agent()

    def initialize_session_state(self):
        if "messages" not in st.session_state:
            st.session_state.messages = []
        if "recipient_id" not in st.session_state:
            st.session_state.recipient_id = str(uuid.uuid4())

    def display_chat_history(self):
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

    def get_chatbot_response(self, user_input: str) -> str:
        item = {
            "content": user_input,
            "recipient_id": st.session_state.recipient_id,
            "sender_id": "chatbot",
            "history_id": f"history_{len(st.session_state.messages)}",
            "qa_id": f"qa_{len(st.session_state.messages)}",
        }

        response = self.agent.answer(
            question=user_input,
            recipient_id=st.session_state.recipient_id,
            sender_id="chatbot",
            item=item,
            mode="chat",
            platform=1,  # Assuming 1 is for Streamlit
            history=st.session_state.messages,
        )

        return response[0]  # Returning just the answer text

    def run(self):
        st.title("R-Logic Computer Repair Chatbot")
        self.initialize_session_state()

        self.display_chat_history()

        if user_input := st.chat_input(
            "Ask about your repair status or any questions..."
        ):
            st.session_state.messages.append({"role": "user", "content": user_input})
            with st.chat_message("user"):
                st.markdown(user_input)

            response = self.get_chatbot_response(user_input)
            st.session_state.messages.append({"role": "assistant", "content": response})
            with st.chat_message("assistant"):
                st.markdown(response)


if __name__ == "__main__":
    chatbot = StreamlitChatbot()
    chatbot.run()