import gradio as gr
from huggingface_hub import InferenceClient, HfHubHTTPError, InferenceTimeoutError
import httpx # 确保这个库在requirements.txt中

"""
For more information on `huggingface_hub` Inference API support, please check the docs:
https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""

# 定义多个模型及其对应的InferenceClient
# !!重要!! 如果需要，请替换 'hf_YOUR_TOKEN_HERE' 为您的Hugging Face API Token。
# API Token可以在 https://huggingface.co/settings/tokens 生成。
# 使用API Token可以帮助解决一些访问限制问题，特别是对于热门模型。
# 同时，请确保您已经在Hugging Face网站上同意了Mistral 7B Instruct v0.2等模型的条款。
MODEL_CLIENTS = {
    "Zephyr 7B Beta": InferenceClient("HuggingFaceH4/zephyr-7b-beta"), # 默认使用API Token
    "Mistral 7B Instruct v0.2": InferenceClient("mistralai/Mistral-7B-Instruct-v0.2"), # 默认使用API Token
    # 如果您需要使用API Token，可以这样写：
    # "Zephyr 7B Beta": InferenceClient("HuggingFaceH4/zephyr-7b-beta", token="hf_YOUR_TOKEN_HERE"),
    # "Mistral 7B Instruct v0.2": InferenceClient("mistralai/Mistral-7B-Instruct-v0.2", token="hf_YOUR_TOKEN_HERE"),
    # 更多模型示例（请根据需要添加或删除）：
    # "Llama 2 7B Chat": InferenceClient("meta-llama/Llama-2-7b-chat-hf", token="hf_YOUR_TOKEN_HERE"), # Llama 2通常需要访问权限和API Token
    # "OpenHermes 2.5 Mistral 7B": InferenceClient("teknium/OpenHermes-2.5-Mistral-7B"),
}

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    # 新增参数，用于选择模型
    selected_model_name,
):
    # 根据选择的模型名称获取对应的client
    client = MODEL_CLIENTS.get(selected_model_name)
    if not client:
        # 如果模型名称无效，直接返回错误信息
        yield "错误：未找到选定的模型客户端。请检查模型名称是否正确或已添加到列表中。"
        return

    messages = [{"role": "system", "content": system_message}]

    # 构建完整的对话历史
    for val in history:
        if val[0]:  # 用户消息
            messages.append({"role": "user", "content": val[0]})
        if val[1]:  # 助手消息
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message}) # 添加当前用户消息

    response = ""
    try:
        # 使用选定的client进行推理
        # client.chat_completion() 默认是一个生成器，用于流式传输
        for message_chunk in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True, # 启用流式传输
            temperature=temperature,
            top_p=top_p,
        ):
            # 确保 chunk 和 content 存在，以防API响应格式异常
            if message_chunk.choices and message_chunk.choices[0].delta and message_chunk.choices[0].delta.content is not None:
                token = message_chunk.choices[0].delta.content
                response += token
                yield response # 逐步返回生成的文本
            else:
                # 可能是流的末尾，或者是一个空的内容块
                pass

    # 错误处理：捕获可能出现的各种异常
    except HfHubHTTPError as e:
        error_message = ""
        if e.response.status_code == 402:
            error_message = "抱歉，此模型服务需要付费访问或您的Hugging Face账户额度已用尽。请检查您的Hugging Face账户设置或Space日志。"
        elif e.response.status_code == 429:
            error_message = "抱歉，请求过于频繁，触发了速率限制。请稍后再试，或考虑使用API Token提升额度。"
        elif e.response.status_code == 401 or e.response.status_code == 403:
            error_message = "抱歉，模型访问权限不足或API Token无效/缺失。请确保您已在Hugging Face上登录，接受模型条款，并正确配置API Token。"
        elif e.response.status_code == 503:
            error_message = "模型服务当前不可用，可能正在加载或维护中。请稍后再试。"
        else:
            error_message = f"模型服务出现HTTP错误 ({e.response.status_code})：{e.response.text}。请检查Hugging Face Space日志。"
        print(f"HfHubHTTPError: {e}") # 打印到控制台以供调试
        yield error_message # 将错误信息显示给用户

    except InferenceTimeoutError as e:
        error_message = "模型响应超时，可能是请求过于复杂或服务器繁忙。请尝试减少'Max new tokens'或稍后再试。"
        print(f"InferenceTimeoutError: {e}")
        yield error_message

    except httpx.HTTPStatusError as e:
        # 这是httpx库抛出的HTTP错误，可能发生在HfHubHTTPError之外
        error_message = f"与Hugging Face服务通信时发生HTTP错误 ({e.response.status_code})：{e.response.text}。"
        print(f"HTTPStatusError: {e}")
        yield error_message

    except Exception as e:
        # 捕获所有其他未预期的错误
        error_message = f"发生未知错误：{type(e).__name__} - {e}。请查看Hugging Face Space日志了解更多详情。"
        print(f"General Error: {e}")
        yield error_message


"""
For information on how to customize the ChatInterface, peruse the gradio docs:
https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="你是一个友好的AI助手，尽力提供帮助。", label="系统消息 (System message)"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="最大生成token数 (Max new tokens)"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="随机性 (Temperature)"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (核心采样)",
        ),
        # 新增一个Dropdown用于选择模型
        gr.Dropdown(
            list(MODEL_CLIENTS.keys()), # 选项为MODEL_CLIENTS的键（模型名称）
            value=list(MODEL_CLIENTS.keys())[0], # 默认选中第一个模型
            label="选择语言模型 (Select Model)",
            interactive=True, # 允许用户更改
        ),
    ],
    title="多模型AI聊天助手", # 给界面添加一个标题
    description="选择一个语言模型，开始与AI对话。您可以调整参数或切换模型进行比较。", # 添加描述
    submit_btn="发送",
    stop_btn="停止",
    clear_btn="清空对话",
)


if __name__ == "__main__":
    demo.launch()