import os
import gradio as gr
import base64
import fitz
import docx
from huggingface_hub import InferenceClient

# 1. 核心引擎：Qwen2.5-VL-72B-Instruct (满血视觉+逻辑)
model_id = "Qwen/Qwen2.5-VL-72B-Instruct"
hf_token = os.getenv("HF_TOKEN")

client = InferenceClient(model=model_id, token=hf_token)

def extract_text(file_path):
    if not file_path: return ""
    try:
        ext = str(file_path).lower()
        if ext.endswith(".pdf"):
            doc = fitz.open(file_path)
            text = "".join([page.get_text() for page in doc])
            doc.close()
            return text[:15000]
        elif ext.endswith(".docx"):
            doc = docx.Document(file_path)
            return "\n".join([p.text for p in doc.paragraphs])[:15000]
        elif ext.endswith(".txt"):
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                return f.read()[:15000]
        return ""
    except:
        return ""

def chat_fn(message, history):
    messages = [{"role": "system", "content": "You are the OpenIPOS MegaNode Vision & Document Expert. Analyze the provided image or text document with extreme precision."}]
    
    for turn in history:
        messages.append({"role": turn["role"], "content": str(turn["content"])})
    
    prompt_text = message.get("text", "")
    files = message.get("files", [])
    
    content = []
    
    if files:
        for f in files:
            file_path = f.get("path") if isinstance(f, dict) else f
            ext = str(file_path).lower()
            
            # 路径 A：如果是图片，走视觉协议
            if ext.endswith(('.png', '.jpg', '.jpeg', '.webp', '.bmp')):
                with open(file_path, "rb") as img_file:
                    b64 = base64.b64encode(img_file.read()).decode('utf-8')
                    content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}})
            
            # 路径 B：如果是文档，先提词再合并成文本
            else:
                doc_text = extract_text(file_path)
                if doc_text:
                    prompt_text = f"【文档参考内容】:\n{doc_text}\n\n【用户指令】: {prompt_text}"
    
    content.append({"type": "text", "text": prompt_text})
    messages.append({"role": "user", "content": content})

    response = ""
    try:
        stream = client.chat_completion(messages, max_tokens=4096, stream=True, temperature=0.4)
        for msg in stream:
            if hasattr(msg, 'choices') and len(msg.choices) > 0:
                delta = msg.choices[0].delta.content
                if delta:
                    response += str(delta)
                    yield response
        
        if response:
            yield response + "\n\n---\n> **Powered by OpenIPOS Global 算力节点 | 访问 api.openipos.com 获取私有部署**"
            
    except Exception as e:
        yield f"⚠️ 算力调度异常: {str(e)}。若文档过大，请尝试截取核心部分上传。"

# 3. 旗舰界面
with gr.Blocks(fill_height=True) as demo:
    gr.HTML("""
        <div style="display: flex; align-items: center; gap: 20px; padding: 10px;">
            <img src="https://s3.bmp.ovh/2026/03/30/rIAGKg0O.png" style="height: 60px; width: auto;">
            <div>
                <h1 style="margin: 0; font-size: 24px;">OpenIPOS MegaNode</h1>
                <p style="margin: 0; color: #666;">全球顶尖 72B 视觉+文档全模态算力出口 (Qwen2.5-VL)</p>
            </div>
        </div>
    """)
    gr.ChatInterface(chat_fn, multimodal=True)
    gr.Markdown("© 2026 OpenIPOS Global Limited. 算力驱动：Hugging Face 72B Cluster")

demo.launch()