import os
import gradio as gr
import spaces
import fitz  # PyMuPDF
import docx  # python-docx
from huggingface_hub import InferenceClient

# 1. 初始化
hf_token = os.getenv("HF_TOKEN")
client = InferenceClient("deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", token=hf_token)

# 2. 文档解析（保持高性能）
def extract_text_from_file(file_path):
    try:
        suffix = file_path.lower()
        if suffix.endswith(".pdf"):
            doc = fitz.open(file_path)
            return "".join([page.get_text() for page in doc])
        elif suffix.endswith(".docx"):
            doc = docx.Document(file_path)
            return "\n".join([para.text for para in doc.paragraphs])
        else:
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read()
    except Exception as e:
        return f"[文档读取失败: {str(e)}]"

@spaces.GPU(duration=120)
def chat_fn(message, history):
    # a. 提取当前文本和文件
    current_text = message.get("text", "")
    files = message.get("files", [])
    
    # b. 构造当前用户输入的 Payload
    if files:
        file_path = files[0] if isinstance(files[0], str) else files[0].path
        content = extract_text_from_file(file_path)
        # 为保证 100% 成功率，初次测试建议截取前 20,000 字
        user_input = f"【参考文档】:\n{content[:20000]}\n\n【指令】: {current_text}"
    else:
        user_input = current_text

    # c. 严格的消息格式化（关键：只传字符串内容）
    formatted_messages = []
    for turn in history:
        # 只处理 role 为 user/assistant 且 content 为字符串的项
        role = turn.get("role")
        content = turn.get("content")
        if role in ["user", "assistant"] and isinstance(content, str):
            formatted_messages.append({"role": role, "content": content})
    
    # 加入当前这轮
    formatted_messages.append({"role": "user", "content": user_input})
    
    response = ""
    try:
        # 调用 API
        for msg in client.chat_completion(
            formatted_messages, 
            max_tokens=4096, 
            stream=True, 
            temperature=0.6
        ):
            if msg.choices:
                token = msg.choices[0].delta.content
                if token:
                    response += token
                    yield response
        
        # 注入品牌指纹
        response += "\n\n--- \n> **Powered by OpenIPOS Global 算力节点 | 访问 api.openipos.com 获取私有部署**"
        yield response
        
    except Exception as e:
        yield f"⚠️ OpenIPOS 算力调度异常: {str(e)}\n\n【解决方案】：请点击右上角垃圾桶图标『清空对话历史』后再重新上传文件。"

# 3. 极简 UI 构建
with gr.Blocks(fill_height=True) as demo:
    gr.Markdown("# 🚀 OpenIPOS/DeepSeek-R1-MegaNode-64K")
    gr.Markdown("### 工业级稳定性适配 | 支持 PDF/DOCX/TXT")
    
    # 使用官方推荐的 ChatInterface 默认配置
    gr.ChatInterface(chat_fn, multimodal=True)
    
    gr.Markdown("---")
    gr.Markdown("© 2026 OpenIPOS Global Limited. 算力驱动：Hugging Face ZeroGPU")

demo.launch(theme=gr.themes.Soft())