import gradio as gr
import requests
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# ========== 从环境变量读取密钥 ==========
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
CEREBRAS_API_KEY = os.getenv("CEREBRAS_API_KEY")

# ========== 模型定义 ==========
LOCAL_MODELS = {
    "DistilGPT2 (最快)": "distilgpt2",
    "DialoGPT-small (轻量)": "microsoft/DialoGPT-small",
    "DialoGPT-medium (均衡)": "microsoft/DialoGPT-medium",
}

OPENROUTER_MODELS = {
    "Llama 3.3 70B (免费)": "meta-llama/llama-3.3-70b-instruct:free",
    "Nemotron 3 Super (免费)": "nvidia/nemotron-3-super:free",
    "Qwen3 Next 80B (免费)": "qwen/qwen3-next-80b:free",
}

CEREBRAS_MODELS = {
    "Llama 3.1 8B (Cerebras)": "llama3.1-8b",
}

local_model_cache = {}

def load_local_model(model_key):
    if model_key in local_model_cache:
        return local_model_cache[model_key]
    model_id = LOCAL_MODELS[model_key]
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    local_model_cache[model_key] = (model, tokenizer)
    return model, tokenizer

def is_chinese(text):
    return any('\u4e00' <= c <= '\u9fff' for c in text)

def local_predict(message, model_key):
    if is_chinese(message):
        return "⚠️ This model only supports English input. 请使用英文。"
    model, tokenizer = load_local_model(model_key)
    inputs = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
    outputs = model.generate(
        inputs,
        max_new_tokens=100,
        do_sample=True,
        temperature=0.7,
        repetition_penalty=1.2,
        pad_token_id=tokenizer.eos_token_id
    )
    reply = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
    return reply.strip() or "I'm not sure how to respond."

def openrouter_predict(message, model_name):
    if not OPENROUTER_API_KEY:
        return "❌ OpenRouter API key 未设置，请在 Space Secrets 中添加 OPENROUTER_API_KEY"
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "HTTP-Referer": "https://your-space.hf.space",
        "X-Title": "AI ChatBot",
        "Content-Type": "application/json"
    }
    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": message}],
        "max_tokens": 500,
        "temperature": 0.7
    }
    try:
        resp = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload, timeout=60)
        if resp.status_code == 200:
            return resp.json()["choices"][0]["message"]["content"]
        elif resp.status_code == 402:
            return "⚠️ 免费额度已用完或需启用付费模型。"
        else:
            return f"❌ API 错误 {resp.status_code}"
    except Exception as e:
        return f"❌ 请求失败: {e}"

def cerebras_predict(message, model_name):
    if not CEREBRAS_API_KEY:
        return "❌ Cerebras API key 未设置，请在 Space Secrets 中添加 CEREBRAS_API_KEY"
    headers = {
        "Authorization": f"Bearer {CEREBRAS_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": message}],
        "max_tokens": 500,
        "temperature": 0.7
    }
    try:
        resp = requests.post("https://api.cerebras.ai/v1/chat/completions", headers=headers, json=payload, timeout=60)
        if resp.status_code == 200:
            return resp.json()["choices"][0]["message"]["content"]
        else:
            return f"❌ Cerebras 错误 {resp.status_code}"
    except Exception as e:
        return f"❌ 请求失败: {e}"

def predict(message, history, backend, model_choice):
    if not message.strip():
        return ""
    if backend == "本地 CPU 模型 (限英文)":
        return local_predict(message, model_choice)
    elif backend == "OpenRouter (不限语言)":
        return openrouter_predict(message, model_choice)
    elif backend == "Cerebras (不限语言)":
        return cerebras_predict(message, model_choice)
    else:
        return "未知后端"

with gr.Blocks(title="多后端 AI 聊天机器人", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧠 多后端 AI 聊天机器人")
    gr.Markdown("**本地 CPU 模型**仅限英文；**OpenRouter / Cerebras** 不限语言，且质量更高。")

    with gr.Row():
        backend_radio = gr.Radio(
            choices=["本地 CPU 模型 (限英文)", "OpenRouter (不限语言)", "Cerebras (不限语言)"],
            label="选择后端",
            value="本地 CPU 模型 (限英文)"
        )
        model_dropdown = gr.Dropdown(
            choices=list(LOCAL_MODELS.keys()),
            label="模型选择",
            value=list(LOCAL_MODELS.keys())[2]
        )

    def update_model_dropdown(backend):
        if backend == "本地 CPU 模型 (限英文)":
            return gr.update(choices=list(LOCAL_MODELS.keys()), value=list(LOCAL_MODELS.keys())[2], visible=True)
        elif backend == "OpenRouter (不限语言)":
            return gr.update(choices=list(OPENROUTER_MODELS.keys()), value=list(OPENROUTER_MODELS.keys())[0], visible=True)
        else:
            return gr.update(choices=list(CEREBRAS_MODELS.keys()), value=list(CEREBRAS_MODELS.keys())[0], visible=True)

    backend_radio.change(fn=update_model_dropdown, inputs=backend_radio, outputs=model_dropdown)

    chatbot = gr.ChatInterface(
        fn=predict,
        additional_inputs=[backend_radio, model_dropdown],
        title=None,
        description="多轮对话，自动记忆上下文。"
    )

demo.launch(server_port=7860)