import os
import gradio as gr
import google.generativeai as genai

# --- 代理设置 ---
# 如果您所在的地区无法直接访问Gemini API，请取消下面两行的注释，
# 并将 'http://your-proxy-host:port' 替换为您的代理服务器地址和端口。
# 例如: 'http://127.0.0.1:7890'
# os.environ['HTTPS_PROXY'] = 'http://your-proxy-host:port'
# os.environ['HTTP_PROXY'] = 'http://your-proxy-host:port'
# ------------------

# 从环境变量中获取API密钥
try:
    api_key = os.environ["apikey"]
except KeyError:
    raise RuntimeError("apikey not set in environment. Please get a key from https://aistudio.google.com/app/apikey and set it as an environment variable.")

genai.configure(api_key=api_key)
model = genai.GenerativeModel('gemini-2.5-flash-preview-05-20')

# 预设字典，key为shortcut，value为实际prompt
PRESETS = {
    "why": "请帮我解释用户输入的名词或短语，尽量通俗简短。1. 请使用纯文本格式，不要有任何markdown文本标记。公式尽量使用unicode编码 2.不要输出额外问候语和与结果无关的内容。",
    "coder": "请你扮演一名高级Python开发者，回答时只给出代码和简要注释，不要多余解释。",
    "summarizer": "请你扮演一名专业的内容摘要员，将输入内容压缩为50字以内的中文摘要。",
    # 可继续添加更多预设
}

with gr.Blocks() as demo:
    preset_state = gr.State("")

    with gr.Row():
        chatbot = gr.Chatbot(label="Gemini Flash Chatbot", scale=4)
        with gr.Column(scale=1):
            last_bot_response = gr.Textbox(visible=False, label="Last response")
            copy_btn = gr.Button("Copy Last Response")

    msg = gr.Textbox(label="What do you want to ask?")
    clear = gr.ClearButton([msg, chatbot])

    def user(user_message, history):
        return gr.update(value="", interactive=False), history + [[user_message, None]]

    def bot(history, preset_key):
        if not history or not history[-1][0]:
            yield history, ""
            return

        preset = PRESETS.get(preset_key, None)
        gemini_history = []
        if preset:
            gemini_history.append({'role': 'system', 'parts': [preset]})
        for user_msg, model_msg in history[:-1]:
            gemini_history.append({'role': 'user', 'parts': [user_msg]})
            if model_msg:
                gemini_history.append({'role': 'model', 'parts': [model_msg]})
        current_question = history[-1][0]
        gemini_history.append({'role': 'user', 'parts': [current_question]})

        try:
            response = model.generate_content(gemini_history, stream=True)
            answer = ""
            for chunk in response:
                if hasattr(chunk, 'text') and chunk.text:
                    answer += chunk.text
                    history[-1][1] = answer
                    yield history, answer
        except Exception as e:
            history[-1][1] = f"Error: {e}"
            yield history, history[-1][1]

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, preset_state], [chatbot, last_bot_response]
    )
    
    copy_btn.click(
        None,
        last_bot_response,
        None,
        js="(text) => {navigator.clipboard.writeText(text); return null;}"
    )

    def load_from_url(request: gr.Request):
        preset_key = request.query_params.get("p", "")
        q = request.query_params.get("q")
        
        final_history = None
        last_answer = ""

        if q:
            initial_history = [[q, None]]
            final_history, last_answer = bot(initial_history, preset_key)
        
        return {
            preset_state: preset_key,
            chatbot: gr.update(value=final_history) if final_history else None,
            last_bot_response: last_answer,
            msg: gr.update(value="")
        }

    demo.load(load_from_url, None, [preset_state, chatbot, last_bot_response, msg])

if __name__ == "__main__":
    demo.launch(debug=True, server_name="0.0.0.0", server_port=8000)