import time

from tqdm import tqdm
import random
import requests
import json
import threading
from threading import Thread


api_key_list = [
    "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",
    "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",
    "sk-nYbapOeC5VmSReJB1JgEr3BlbkFJnOo2J9qFJRKXrOSNiYFO",
    "sk-tOy3uBFkPsg9uVWTpDOor3BlbkFJkbXgo0sHAubK8VWyaeso",
    "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",
    "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",
    "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",
    "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",
    "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",
    "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"
                ]

lock = threading.RLock()

def request_api_chatgpt(api_key, prompt):
    try:
        OPENAI_API_KEY = api_key
        url = "https://api.openai.com/v1/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {OPENAI_API_KEY}"
        }
        data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "user", "content": prompt},
            ],
            "temperature": 0.5
        }
        response = requests.post(url,
                                 headers=headers,
                                 data=json.dumps(data),
                                 timeout=240)

        res = response.json()
        text = res["choices"][0]["message"]["content"]

        lock.acquire()
        api_key_list.append(api_key)


        with open("/home/majiahui/mulu_ner/data/prompt_small_gen.txt", mode="a") as f:
            f.write(prompt)
            f.write("**************")
            f.write(text)
            f.write("\n")
        lock.release()
        
    except:
        time.sleep(5)
        lock.acquire()
        api_key_list.append(api_key)
        lock.release()



if __name__ == '__main__':
    with open("./data/prompt_shuffle.txt", encoding="utf-8") as f:
        text = f.read()

    text_list = text.split("\n")

    index = 0
    while True:
        if index == len(text_list):
            break

        if api_key_list == []:
            time.sleep(1)
            continue
        else:
            api_key = api_key_list.pop(0)
            prompt = text_list[index]

            t = Thread(target=request_api_chatgpt, args=(api_key, prompt))
            t.start()
            lock.acquire()
            index += 1
            print(index)
            lock.release()