import gradio as gr
import requests
import json
import base64
from io import BytesIO
from PIL import Image
import os

API_BASE = "https://www.gpt4novel.com/api/xiaoshuoai/ext/v1"

def get_models(api_key):
    try:
        r = requests.get(f"{API_BASE}/models", headers={"Authorization": f"Bearer {api_key}"})
        models = [m["id"] for m in r.json()["data"]]
        return sorted(models, reverse=True)
    except:
        return ["nalang-xl-0826-10k"]

def extract_card_from_png(file_obj):
    if not file_obj.name.endswith(".png"):
        return None
    img = Image.open(file_obj.name)
    data = img.info.get("chara")
    if data:
        decoded = base64.b64decode(data).decode()
        return json.loads(decoded)
    return None

def chat(message, history, model, api_key, system_prompt):
    messages = [{"role": "system", "content": system_prompt or "你是一个有帮助的AI助手。"}]
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        if h[1]: messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    resp = requests.post(
        f"{API_BASE}/chat/completions",
        json={"model": model, "messages": messages, "stream": True, "temperature": 0.7, "max_tokens": 800},
        headers={"Authorization": f"Bearer {api_key}"},
        stream=True,
        timeout=60
    )
    resp.raise_for_status()
    for line in resp.iter_lines():
        if line and line.startswith(b"data: "):
            try:
                data = json.loads(line[6:].decode())
                if token := data["choices"][0]["delta"].get("content"):
                    yield token
            except:
                continue

with gr.Blocks() as demo:
    gr.HTML(open("index.html").read())
    chatbot = gr.Chatbot(height="70vh")
    with gr.Row():
        msg = gr.Textbox(scale=8, placeholder="发送消息...", container=False)
        send = gr.Button("发送", scale=1)
    model = gr.Dropdown(choices=[], label="模型")
    api_key = gr.Textbox(label="API Key", type="password")
    card_file = gr.File(label="导入角色卡 (.json 或 .png)", file_types=[".json", ".png"])

    system_prompt = gr.Textbox(label="系统提示词（角色卡自动填充）", lines=4, value="你是一个有帮助的AI助手。")

    def load_card(file):
        if not file: return "", []
        if str(file).endswith(".png"):
            card = extract_card_from_png(file)
        else:
            card = json.load(open(file.name, encoding="utf-8"))
        if card:
            prompt = card.get("data", card.get("system_prompt", card.get("description", "")))
            name = card.get("name", card.get("data", {}).get("name", "角色"))
            return prompt, [[None, f"已加载角色卡：{name}"]]
        return "", []

    card_file.change(load_card, card_file, [system_prompt, chatbot])

    def connect_key(key):
        models = get_models(key)
        return gr.Dropdown(choices=models, value=models[0] if models else None)

    api_key.submit(connect_key, api_key, model)

    send.click(chat, [msg, chatbot, model, api_key, system_prompt], chatbot).then(lambda: "", None, msg)
    msg.submit(chat, [msg, chatbot, model, api_key, system_prompt], chatbot).then(lambda: "", None, msg)

demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)