import openai
from openai import OpenAI
import tiktoken
from config import Config
from colorama import Fore, Style
import time
class AI_assistant:
    def __init__(self, cfg: Config):
        openai.api_key = cfg.open_ai_key
#        openai.proxy = cfg.open_ai_proxy
        self._chat_model = cfg.open_ai_chat_model
        self._use_stream = cfg.use_stream
        self._encoding = tiktoken.encoding_for_model('gpt-4-1106-preview')
        self._language = cfg.language
        self._temperature = cfg.temperature
        self.client = OpenAI(api_key=cfg.open_ai_key)
        self.assistant = None
        self.thread = None
        self.run = None
    def check_run(self, thread_id, run_id):
        while True:
            # Refresh the run object to get the latest status
            run = self.client.beta.threads.runs.retrieve(
                thread_id=thread_id,
                run_id=run_id
            )

            if run.status == "completed":
                print(f"{Fore.GREEN} Run is completed.{Style.RESET_ALL}")
                break
            elif run.status == "expired":
                print(f"{Fore.RED}Run is expired.{Style.RESET_ALL}")
                break
            else:
                print(f"{Fore.YELLOW} OpenAI: Run is not yet completed. Waiting...{run.status} {Style.RESET_ALL}")
                time.sleep(3)  # Wait for 1 second before checking again
    def create_assistant(self, name, instructions, tools,files):
        self.assistant = self.client.beta.assistants.create(
            name=name,
            instructions=instructions,
            tools=tools,
            model=self._chat_model,
            file_ids=files
        )
    def create_thread(self):
        self.thread = self.client.beta.threads.create()
    def add_message_to_thread(self, role, content):
        self.client.beta.threads.messages.create(
            thread_id=self.thread.id,
            role=role,
            content=content
        )
    def run_assistant(self, instructions):
        self.run = self.client.beta.threads.runs.create(
            thread_id=self.thread.id,
            assistant_id=self.assistant.id,
            instructions=instructions
        )
    def process_messages(self):
        messages = self.client.beta.threads.messages.list(thread_id=self.thread.id)
        total_price = 0
        ans = ""
        for msg in messages.data:
            role = msg.role
            content = msg.content[0].text.value
            if role == "user":
                total_price = total_price + self._num_tokens_from_string(content)/1000*0.01
            elif role == "assistant":
                total_price = total_price + self._num_tokens_from_string(content)/1000*0.03
                ans = content
        return total_price , ans
    def upload_file(self, file_path):
    # Upload the file to the thread
        if file_path != "":
            file = self.client.files.create(
                file=open(file_path, "rb"),
                purpose = 'assistants'
            )
            print("File successfully uploaded. File ID :" , file.id)

        return file.id
    def get_files(self):
        lists = self.client.files.list()
        files_id = []
        for list in lists:
            files_id.append(list.id)
        return files_id
    def delete_all_files(self):
        files_id = self.get_files()
        for id in files_id:
            self.client.files.delete(id=id)
    def _num_tokens_from_string(self, string: str) -> int:
        """Returns the number of tokens in a text string."""
        num_tokens = len(self._encoding.encode(string))
        return num_tokens