import gradio as gr
import openai
import os
import time

# --- 配置 ---
ARK_API_KEY = os.environ.get("ARK_API_KEY")
if not ARK_API_KEY:
    raise ValueError("ARK_API_KEY not found. Please set it in Space Secrets.")

client = openai.OpenAI(api_key=ARK_API_KEY, base_url="https://ark.cn-beijing.volces.com/api/v3")

# --- 可选模型列表 ---
# 键是用户看到的显示名称，值是实际的API模型ID
AVAILABLE_MODELS = {
    "DeepSeek R1 (250120)": "deepseek-r1-250120",
    "DeepSeek v3 (250324)": "deepseek-v3-250324" # 示例，你需要有这个模型的权限
    # 你可以添加更多你账户下可用的模型
    # "模型名称A": "模型ID_A",
    # "模型名称B": "模型ID_B",
}
DEFAULT_MODEL_DISPLAY_NAME = list(AVAILABLE_MODELS.keys())[0] # 默认选择第一个模型

# --- 自定义 CSS ---
custom_css = """
body {
    background-color: #e9ecef; /* 稍微调暗一点背景，让前景更突出 */
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}

.gradio-container {
    border-radius: 20px !important; /* 稍微增大圆角 */
    box-shadow: 0 8px 24px rgba(0,0,0,0.15) !important; /* 增强阴影 */
    background-color: #ffffff;
    /* border: 1px solid #dee2e6 !important; */ /* （可选）添加一个细边框 */
}

.message-bubble-user {
    background-color: #007bff !important;
    color: white !important;
    border-radius: 20px 20px 5px 20px !important; /* 调整圆角使其更有特色 */
    align-self: flex-end;
    max-width: 75%; /* 稍微增加最大宽度 */
    box-shadow: 3px 3px 8px rgba(0,0,0,0.2) !important; /* 添加阴影使其浮动 */
    margin-bottom: 10px !important; /* 增加一点底部间距 */
}

.message-bubble-bot {
    background-color: #f8f9fa !important; /* 使用更亮的背景增加对比 */
    color: #212529 !important; /* 使用更深的文字颜色增加对比 */
    border-radius: 20px 20px 20px 5px !important; /* 调整圆角 */
    align-self: flex-start;
    max-width: 75%; /* 稍微增加最大宽度 */
    border: 1px solid #ced4da !important; /* 添加边框使其更清晰 */
    box-shadow: -3px 3px 8px rgba(0,0,0,0.1) !important; /* 添加阴影 */
    margin-bottom: 10px !important; /* 增加一点底部间距 */
}

.message-bubble-user .message-text, .message-bubble-bot .message-text {
    font-size: 1.05rem !important; /* 稍微增大字体 */
    padding: 12px 18px !important; /* 增加内边距 */
    line-height: 1.6 !important; /* 增加行高，提高可读性 */
}

.avatar-container img {
    border-radius: 50% !important;
    width: 48px !important; /* 增大头像尺寸 */
    height: 48px !important; /* 增大头像尺寸 */
    object-fit: cover;
    margin: 8px !important; /* 增加外边距 */
    box-shadow: 0 2px 6px rgba(0,0,0,0.25) !important; /* 增强阴影 */
    border: 2px solid #ffffff !important; /* 给头像加一个白色边框，使其从背景中突出 */
}

.textbox textarea {
    border-radius: 25px !important; /* 增大圆角 */
    border: 2px solid #343a40 !important; /* 加粗边框 */
    background-color: #495057 !important;
    color: #f8f9fa !important;
    padding: 12px 18px !important; /* 增加内边距 */
    font-size: 1.05rem !important; /* 稍微增大字体 */
    box-shadow: inset 0 1px 3px rgba(0,0,0,0.2) !important; /* 添加内阴影增加深度感 */
}

.textbox textarea:focus {
    border-color: #007bff !important;
    background-color: #5a6268 !important;
    box-shadow: 0 0 0 0.25rem rgba(0, 123, 255, 0.35) !important; /* 增强聚焦阴影 */
    /* transform: scale(1.01); */ /* （可选）聚焦时轻微放大 */
}

.submit-button {
    background-color: #007bff !important;
    color: white !important;
    border-radius: 25px !important; /* 增大圆角 */
    font-weight: bold !important;
    font-size: 1.1rem !important; /* 增大字体 */
    padding: 12px 25px !important; /* 增加内边距 */
    transition: all 0.2s ease-in-out; /* 平滑所有过渡 */
    box-shadow: 0 3px 6px rgba(0, 123, 255, 0.3) !important; /* 添加阴影 */
    border: none !important; /* 移除默认边框（如果存在） */
}

.submit-button:hover {
    background-color: #0056b3 !important;
    box-shadow: 0 5px 10px rgba(0, 86, 179, 0.4) !important; /* 悬停时增强阴影 */
    transform: translateY(-2px); /* 悬停时轻微上移 */
}

.gr-prose h1, .gr-prose h2, .gr-prose h3 {
    color: #212529; /* 更深的颜色 */
    text-align: center;
    margin-bottom: 0.75em !important; /* 增加标题下的间距 */
    text-shadow: 1px 1px 2px rgba(0,0,0,0.1); /* 给标题添加轻微文字阴影 */
}
.gr-prose h1 {
    font-size: 2.8em !important; /* 增大 H1 */
    font-weight: 700 !important;
}
.gr-prose h2 {
    font-size: 2.2em !important; /* 增大 H2 */
    font-weight: 600 !important;
}
.gr-prose h3 {
    font-size: 1.8em !important; /* 增大 H3 */
    font-weight: 600 !important;
}

.gr-prose p {
    color: #343a40; /* 更深的颜色 */
    text-align: center;
    font-size: 1.15em !important; /* 增大段落字体 */
    line-height: 1.7 !important; /* 增加行高 */
    margin-bottom: 1.5em !important; /* 增加段落下间距 */
}

/* 为下拉菜单添加一些边距和视觉强调 */
.gradio-dropdown {
    margin-top: 10px !important; /* 增加上方外边距 */
    margin-bottom: 25px !important; /* 显著增加下方外边距，使其与其他元素分隔开 */
    /* 如果想让下拉框本身更突出，可能需要针对其内部元素如 .gr-select select 或类似的选择器 */
    /* 例如（需要根据实际Gradio版本检查选择器）： */
    /*
    select {
        border-width: 2px !important;
        border-color: #007bff !important;
        font-weight: bold !important;
    }
    */
}
"""

# --- 聊天逻辑 (修改以接收模型ID) ---
def call_openai_with_history(message: str, history: list, selected_model_display_name: str):
    """
    调用 OpenAI API，并处理聊天历史。
    message: 当前用户输入
    history: Gradio 的聊天历史
    selected_model_display_name: 用户从下拉菜单选择的模型显示名称
    """
    actual_model_id = AVAILABLE_MODELS.get(selected_model_display_name, AVAILABLE_MODELS[DEFAULT_MODEL_DISPLAY_NAME])
    print(f"使用模型: {actual_model_id}") # 调试信息

    messages_for_api = [{"role": "system", "content": "你是一个乐于助人、风趣幽默的AI助手。"}]

    for user_msg, bot_msg in history:
        # Gradio 的 history 格式是 [[user_msg1, bot_msg1], ...]
        # 当 chatbot type='messages' 时，history 格式是 [{'role':'user', 'content':...}, {'role':'assistant', 'content':...}]
        # 这里需要根据 chatbot 的 type 来调整 history 的处理方式
        # 由于 chatbot_component.type = 'messages', history 已经是 OpenAI 期望的格式了
        # 但是 ChatInterface 传递给 fn 的 history 仍然是 [[user, bot], ...] 格式
        # 所以我们仍然需要转换
        if isinstance(user_msg, str): # 确保是 [[user, bot], ...] 格式
            messages_for_api.append({"role": "user", "content": user_msg})
        if bot_msg and isinstance(bot_msg, str):
             messages_for_api.append({"role": "assistant", "content": bot_msg})

    messages_for_api.append({"role": "user", "content": message})

    try:
        response = client.chat.completions.create(
            model=actual_model_id, # 使用选择的模型ID
            messages=messages_for_api,
            temperature=0.7,
            max_tokens=1024, # 可以适当增加 max_tokens
        )
        bot_response = response.choices[0].message.content
        return bot_response

    except openai.APIError as e:
        error_message = f"OpenAI API Error: {e}"
        print(error_message)
        if hasattr(e, 'message'):
            return f"抱歉，调用AI服务时出错: {e.message}"
        return f"抱歉，调用AI服务时出错: {str(e)}"
    except Exception as e:
        error_message = f"An unexpected error occurred: {e}"
        print(error_message)
        return "抱歉，发生了一个未知错误。"


# --- Gradio 界面 ---
with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("# 🤖 Starry的智能AI聊天伙伴")
    gr.Markdown("与基于大模型的AI助手进行友好对话吧！")

    # 添加模型选择下拉菜单
    model_selector = gr.Dropdown(
        choices=list(AVAILABLE_MODELS.keys()), # 用户看到的选项是显示名称
        value=DEFAULT_MODEL_DISPLAY_NAME,     # 默认值
        label="选择AI模型",
        info="更改模型后，请刷新窗口，新的对话将使用所选模型。",
        elem_classes="gradio-dropdown" # 应用自定义CSS类 (可选)
    )

    chatbot_component = gr.Chatbot(
        label="聊天窗口",
        type='messages', # 确保使用 'messages' 类型，以便更好地处理头像和角色
        avatar_images=("images/user_avatar.png", "images/bot_avatar.png"),
        height=500,
        show_label=False # 隐藏 Chatbot 自己的标签，因为上面已经有标题了
    )

    chat_interface = gr.ChatInterface(
        fn=call_openai_with_history,
        chatbot=chatbot_component,
        type='messages',
        additional_inputs=[model_selector], # 将下拉菜单作为额外输入传递给 fn
        submit_btn="➤ 发送",
    )

    # 如果你的 Gradio 版本较低，可能没有 avatar_images 直接在 Chatbot 构造函数里
    # 如果头像图片不显示，请确保 `images` 文件夹和里面的 `user_avatar.png` `bot_avatar.png`
    # 文件与你的 app.py 在同一级别，或者你提供了正确的相对/绝对路径。
    # 对于 Hugging Face Spaces，通常将 `images` 文件夹放在仓库根目录。

# --- 启动应用 ---
if __name__ == "__main__":
    # demo.queue() # 如果 API 调用时间较长，可以启用队列
    demo.launch()