import openai
import streamlit as st
import os,time
from dotenv import load_dotenv
load_dotenv()
st.set_page_config(page_title="Pak Case Law ChatBot")
st.header("Wellcome To Pakistani AI Lawyer!")

class OpenAIClient:
    def __init__(self):
        self.client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        self.thread_id = None
        self.file_id = ["file-CwRdM2JgGKEVZP4rhoKSvgpD"]
        self.thread = self.client.beta.threads.create()
      
    def upload_file(self):
        file=self.client.files.create(
        file=open("1964_35_THE_WEST_PAKISTAN_FAMILY_COURTS_ACT_1964.pdf", "rb"),
        purpose="assistants")
        self.file_id.append(file.id)
        print(file.id)
        print("File Uploaded")
      
    def main(self,user_message):
      lawyer = self.client.beta.assistants.create(
        instructions="""You are an Pakistani case law provider, and you have access to files to answer client questions about Pakistan case laws. Always response with info from either of the files.\
Please always make sure to follow these instructions that are delimited in triple backticks:\
```
1 - Behave like a Pakistani lawyer.\
2 - Always remember that you are also suitable for specific legal issues.\
3 - Providing helpful information, drafting documents, reviewing documents, or completing tasks based on user input.
""",
        name="Pakistani Lawyer",
        tools=[{"type": "retrieval"}],
        model="gpt-3.5-turbo-1106",
        file_ids=self.file_id)

      message = self.client.beta.threads.messages.create(
        thread_id= self.thread.id,
        role="user",
        content=user_message,
        file_ids= self.file_id
      )
      _run = self.client.beta.threads.runs.create(
        thread_id= self.thread.id,
        assistant_id= lawyer.id
      )

      while True:
        run = self.client.beta.threads.runs.retrieve(
        thread_id=self.thread.id,
        run_id=_run.id
      ).status
        
        time.sleep(3)
        if run == "completed":
          print(run)
          return self.get_response()
        elif run == "queued":
          continue
        elif run == "failed" or run == "cancelled" or run == "expired":
          print(run)
          break
        elif run == "in_progress":
          continue
        else:
          break
          
    def get_response(self):
      messages = self.client.beta.threads.messages.list(
        thread_id= self.thread.id
      )
      return messages.data[0].content[0].text.value
      
    def file_handler(self):
      return self.client.files.retrieve("file-CwRdM2JgGKEVZP4rhoKSvgpD")
    def delete_file(self):
      return self.client.files.list()
  
try:
  client = OpenAIClient()
  input = st.text_input("Input: ",key='input')
  response=client.main(input)
  submit = st.button("Ask the Question")
  if submit:
    st.subheader("The Response is")
    st.write(response)
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)