import threading
import time
import os
from multiprocessing import Queue
import json
import random

taskQueue = Queue()
resultQueue = Queue()


def callServerStimu(prompt, model="deepseek-ai/DeepSeek-R1"):
    # 这里模拟等待远程服务器的结果
    time.sleep(random.randint(2, 5))
    success = True
    msg = "OK"
    think = ""
    result = ""
    count = 0
    # 重试机制
    while count < 3:
        count += 1
        try:
            if prompt[0] == "能":
                # 模拟异常
                print("第 " + str(count) + " 次重试：" + prompt)
                1 / 0
            think = 'reasoning_content of ' + prompt
            result = 'content of ' + prompt
            break
        except Exception as e:
            success = False
            msg = "API调用异常:" + str(e) + prompt
    return success, msg, think, result


# 生产者
def produceTask(fileName, targetName):
    # 万一崩溃重启怎么办？读取结果文件，把所有已完成的Id记录下来
    # 后面再遇到已完成的样本，就不用放任务队列了
    finishedSet = set()
    if os.path.exists(targetName):
        with open(targetName, "r", encoding="utf-8") as file:
            for line in file:
                # print(line)
                data = json.loads(line.strip())
                if "id" in data:
                    finishedSet.add(data["id"])

    with open(fileName, "r", encoding="utf-8") as file:
        #内存友好型的读取文件
        # 不友好的 data = file.read()
        for line in file:   # 逐行读取
            data = json.loads(line.strip())
            if data['id'] not in finishedSet:
                while True:
                    # 这个队列在内存里，不要无限放数据
                    # 如果替换成消息队列，要根据实际情况看允许放多少数据
                    # 队列友好型的
                    if taskQueue.qsize() < 20:
                        taskQueue.put(data)
                        print("队列长度："+str(taskQueue.qsize())+" 生产者正在生产任务："+str(data["id"]))
                        break
                    else:
                        time.sleep(1)
                        print("生产者等待中 ,队列长度 = "+str(taskQueue.qsize()))
            else:
                print("这条数据已经处理过了： "+str(str(data["id"])))
                pass


# 消费者
def consumTask(threadIndex):
    while True:
        data = taskQueue.get(block=True)
        success, msg, think, result = callServerStimu(data["instruction"])
        # print("消费者: "+str(threadIndex)+" 已经处理完数据： "+str(data["id"]))
        resultQueue.put([data, success, msg, think, result])


def askBatchMT(fileName, targetName, processorCount):
    # 启动生产者，可以是多个
    p = threading.Thread(target=produceTask, args=(fileName, targetName,))
    p.start()
    # 启动消费者，可以是多个
    # IO密集型任务最好用线程，具体数量要看访问服务器的吞吐量
    # 计算密集型任务python必须用进程，具体数量根据本机cpu、内存、gpu和任务的复杂度来确定
    for i in range(processorCount):
        p = threading.Thread(target=consumTask, args=(i,))
        p.start()

    # 接收结果
    with open(targetName, 'w', encoding='utf-8') as target:
        while True:
            data, success, msg, think, result = resultQueue.get(block=True)
            if success:
                data["output"] = "<think>\n" + think + "\n</think>\n\n\n" + result
                target.write(json.dumps(data, ensure_ascii=False) + "\n")
                target.flush()
                #print("数据入库：" + str(data["id"]), success, msg, think, result)
            else:
                #print("警告，结果错误：" + str(data["id"]), success, msg, think, result)
                pass


if __name__ == "__main__":
    fileName = "../data/testDataTemp/queryData.jsonl"
    targetName = "../data/testDataTemp/answerData_producer_consumer.jsonl"
    processorCount = 10 # 10个线程, 10个消费者
    askBatchMT(fileName, targetName, processorCount)