import openai
import streamlit as st
import os
import time
from dotenv import load_dotenv
load_dotenv()
st.set_page_config(page_title="Pak Case Law ChatBot")
st.header("Welcome To Pakistani AI Lawyer!")
st.markdown("Currently the ChatBot is taking content only from this [file.](http://kpcode.kp.gov.pk/uploads/1964_35_THE_WEST_PAKISTAN_FAMILY_COURTS_ACT_1964.pdf)")

class OpenAIClient:
    def __init__(self):
        self.client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        self.thread_id = None
        self.file_id = ["file-Anvyhz2k0puV0NCjWH8mGsuk"]
        self.assistant_id = "asst_Z7pj5QGUjiAlVpJzqsX2SJkg"
        self.run_id = "run_A5wq14ZMQXBpaDnULFPky5Bb"

    def create_thread_id(self):
        if "thread_id" not in st.session_state:
            _id = self.client.beta.threads.create().id
            st.session_state["thread_id"] = _id
        
    def main(self, user_message):
        message = self.client.beta.threads.messages.create(
            st.session_state.get("thread_id"),
            role="user",
            content=user_message,
        )
        _run = self.client.beta.threads.runs.create(thread_id=st.session_state.get("thread_id"), assistant_id=self.assistant_id)
        while True:
            run = self.client.beta.threads.runs.retrieve(thread_id=st.session_state.get("thread_id"),
                                                        run_id=_run.id).status

            time.sleep(3)
            if run == "completed":
                return self.get_response()
            elif run == "queued":
                continue
            elif run == "failed" or run == "cancelled" or run == "expired":
                return run
            elif run == "in_progress":
                continue
            else:
                break

    def get_response(self):
        messages = self.client.beta.threads.messages.list(
            thread_id=st.session_state.get("thread_id")
        )
        return messages.data[0].content[0].text.value

try:
    client = OpenAIClient()
    client.create_thread_id()
    user_input = st.text_input("Input: ", key='input')
    submit = st.button("Ask the Question")
    if submit:
        st.subheader("The Response is")
        st.write(client.main(user_input))
except openai.AuthenticationError:
    st.write("Your API key or token was invalid, expired, or revoked.")
except openai.RateLimitError:
    st.write("You have hit your assigned rate limit. Please try again in a few minutes.")
except Exception as e:
    st.write(e)