from openai import OpenAI
import os
import gradio as gr
import time




# Define the dropdown options
dropdown_options = ["1_BRD_Lending_v0.1.pdf", "2_BRD_Lending_v0.1.pdf", "3_BRD_Payments_v0.1.pdf"]


query = {
    dropdown_options[0] : "Generate test cases for the following process:  Loan Boarding Process, Loan Repayment Management, Loan Servicing Management, Disbursement Cancellation",
    dropdown_options[1] : "Generate test cases for the following process: Process Flows, Quick Data Entry, RBI Defaulters, Credit Approval Process",
    dropdown_options[2] : "Generate test cases for the following process: Fund Transfer and Payments, SBM Account Payment, FCY Payment , Credit Card Payment",
}



class Assistant:
  def __init__(self, path):
    self.client = OpenAI()
    print(str(path))
    # Upload a file with an "assistants" purpose
    self.file = self.client.files.create(
      file=open(str(path), "rb"),
      purpose='assistants'
    )
    
    # Add the file to the assistant
    self.assistant = self.client.beta.assistants.create(
      instructions="You are an helpful assistant to generating test scenarios and test cases for test design process automation.",
      model="gpt-4-1106-preview",
      tools=[{"type": "retrieval"}],
      file_ids=[self.file.id]
    )

    self.thread = self.client.beta.threads.create()

    self.query = query.get(path)


  def infer(self):
    
    
    
    message = self.client.beta.threads.messages.create(
      thread_id=self.thread.id,
      role="user",
      content=self.query,
    )

    # print(message)
    
    run = self.client.beta.threads.runs.create(
        thread_id = self.thread.id,
        assistant_id= self.assistant.id
    )

    while True:
        retreive =  self.client.beta.threads.runs.retrieve(
            thread_id = self.thread.id,
            run_id= run.id
        )
        print(retreive.status)
        if retreive.status=="completed":
            break
        else:
            time.sleep(1)
            continue

    # time.sleep(1)

    # run =  self.client.beta.threads.runs.retrieve(
    #     thread_id = thread.id,
    #     run_id= run.id
    # )
      
    messages = self.client.beta.threads.messages.list(
        thread_id =self.thread.id
    )
    print(messages)
    for message in messages.data:
        print(message.content[0].text.value)
        return  message.content[0].text.value

        


if __name__ == "__main__":

    
    # Define a function that takes the selected option as an argument
    def create_assistant(path):
        obj = Assistant(path)
        return obj.infer()
    



    demo_2 = gr.Interface(
        create_assistant,
        [gr.Dropdown(dropdown_options, label="path", info="Select BRDS doc"),],
        "text",
         )
    
    # # Create the Gradio interface with a dropdown
    # iface = gr.Interface(
    #     fn=create_assistant(inputs),
    #     inputs=gr.Dropdown(choices=dropdown_options, label="Select an option"),
    #     outputs="text"
    # )

    # gr.Dropdown(
    #         dropdown_options, label="Animal", info="Will add more animals later!"
    #     ),

    
    # Launch the interface
    print("Running main")
    
    # demo = gr.Interface(fn=obj.infer, inputs="text", outputs="text")
    demo_2.launch(show_api=False, share=True)