import gradio as gr
import os
import sys
import json
import requests
import random
from tenacity import retry, wait_fixed, stop_after_attempt

# 模型配置和API URL
MODEL = "o1-preview"
API_URL = os.getenv("API_URL")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

print(API_URL)
print(OPENAI_API_KEY)

# 异常处理函数
def exception_handler(exception_type, exception, traceback):
    print("%s: %s" % (exception_type.__name__, exception))
sys.excepthook = exception_handler
sys.tracebacklimit = 0

# 重试机制：如果请求失败，最多重试5次，每次等待2秒
@retry(stop=stop_after_attempt(5), wait=wait_fixed(2))
def call_openai_api(payload, headers):
    response = requests.post(API_URL, headers=headers, json=payload, stream=True)
    response.raise_for_status()  # 如果返回状态码不是200，抛出异常
    return response

# 处理API请求的核心函数
def predict(inputs, top_p, temperature, chat_counter, chatbot, history):
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": f"{inputs}"}],
        "temperature": temperature,
        "top_p": top_p,
        "n": 1,
        "stream": True,
        "presence_penalty": 0,
        "frequency_penalty": 0,
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {OPENAI_API_KEY}"
    }

    if chat_counter != 0:
        messages = []
        for i, data in enumerate(history):
            role = 'user' if i % 2 == 0 else 'assistant'
            message = {"role": role, "content": data}
            messages.append(message)

        messages.append({"role": "user", "content": inputs})
        payload = {
            "model": MODEL,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "n": 1,
            "stream": True,
            "presence_penalty": 0,
            "frequency_penalty": 0,
        }

    chat_counter += 1
    history.append(inputs)
    token_counter = 0
    partial_words = ""
    counter = 0

    try:
        # 使用重试机制的 API 请求
        response = call_openai_api(payload, headers)

        for chunk in response.iter_lines():
            if counter == 0:
                counter += 1
                continue
            if chunk.decode():
                chunk = chunk.decode()
                if len(chunk) > 12 and "content" in json.loads(chunk[6:])['choices'][0]['delta']:
                    partial_words += json.loads(chunk[6:])['choices'][0]["delta"]["content"]
                    if token_counter == 0:
                        history.append(" " + partial_words)
                    else:
                        history[-1] = partial_words
                    token_counter += 1
                    yield [(history[i], history[i + 1]) for i in range(0, len(history) - 1, 2)], history, chat_counter  # 更新chatbot内容

    except Exception as e:
        print(f'Error encountered: {e}')
    yield [(history[i], history[i + 1]) for i in range(0, len(history) - 1, 2)], history, chat_counter

# 重置输入框
def reset_textbox():
    return gr.update(value='', interactive=False)

# Gradio 界面部分
title = """<h1 align="center">OpenAI-O1-Preview: Personal Version</h1>"""
description = """This app allows a single user to interact with an OpenAI GPT-4 Turbo model."""

with gr.Blocks() as demo:
    gr.HTML(title)
    
    chatbot = gr.Chatbot()  # 对话框
    inputs = gr.Textbox(placeholder="Type your input here", label="Input")
    state = gr.State([])  # 保存历史对话
    chat_counter = gr.Number(value=0, visible=False, precision=0)  # 计数器

    # 参数控制
    with gr.Accordion("Advanced Parameters", open=False):
        top_p = gr.Slider(minimum=0, maximum=1.0, value=1.0, step=0.05, interactive=True, label="Top-p")
        temperature = gr.Slider(minimum=0, maximum=5.0, value=1.0, step=0.1, interactive=True, label="Temperature")

    # 处理用户提交的输入
    inputs.submit(reset_textbox, [], [inputs], queue=False)
    inputs.submit(predict, [inputs, top_p, temperature, chat_counter, chatbot, state], [chatbot, state, chat_counter], queue=False)

# 启动Gradio应用
demo.launch(share=False)