import base64
import numpy as np
import random
import logging
import os
import json
import openai
from openai import OpenAIError
from time import sleep


def _throttled_openai_chat_completion_create(
    client,
    model: str,
    messages,
    temperature: float,
    max_tokens: int,
    top_p: float,
    json_format: bool = False,
    n: int = 1,
):
    """同步的OpenAI聊天补全函数，支持限流与重试"""
    for _ in range(10):  # 进行10次尝试
        try:
            if json_format:
                return client.chat.completions.create(
                    model=model,
                    messages=messages,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    top_p=top_p,
                    n=n,
                    response_format={"type": "json_object"},
                )
            else:
                return client.chat.completions.create(
                    model=model,
                    messages=messages,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    top_p=top_p,
                    n=n,
                )
        except openai.RateLimitError as e:
            print("Rate limit exceeded, retrying...")
            sleep(random.randint(10, 20))  # 增加重试等待时间
        except openai.BadRequestError as e:
            print(e)
            return None
        except OpenAIError as e:
            print(e)
            sleep(random.randint(5, 10))
    return None


def generate_from_openai_chat_completion(
    client,
    messages,
    engine_name: str,
    temperature: float = 1.0,
    max_tokens: int = 512,
    top_p: float = 1.0,
    requests_per_minute: int = 100,
    json_format: bool = False,
    n: int = 1,
):
    """同步生成OpenAI聊天补全"""
    delay = 60.0 / requests_per_minute
    sleep(delay)  # 简单的限流处理

    responses = [
        _throttled_openai_chat_completion_create(
            client,
            model=engine_name,
            messages=message,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            json_format=json_format,
            n=n,
        )
        for message in messages
    ]

    empty_dict = {
        "question": "",
        "options": {
            "A": "",
            "B": "",
            "C": "",
            "D": "",
        },
        "distractors": {
            "E": "",
            "F": "",
            "G": "",
        },
        "correct_answer": ""
    }
    empty_str = ""
    outputs = []
    for response in responses:
        if n == 1:
            if json_format:
                if response and response.choices[0] and response.choices[0].message and response.choices[0].message.content:
                    outputs.append(json.loads(response.choices[0].message.content))
                else:
                    outputs.append(empty_dict)
            else:
                if response and response.choices[0] and response.choices[0].message and response.choices[0].message.content:
                    outputs.append(response.choices[0].message.content)
                else:
                    outputs.append(empty_str)
        else:
            if json_format:
                outputs.append([
                    json.loads(response.choices[i].message.content) if response and response.choices[i].message.content else empty_dict
                    for i in range(n)
                ])
            else:
                outputs.append([
                    response.choices[i].message.content if response and response.choices[i].message.content else empty_str
                    for i in range(n)
                ])
    return outputs