import gradio as gr
import requests
import os
import json
import fitz  # PyMuPDF
from pathlib import Path

# --- 核心配置 ---
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
MODEL_ID = "google/gemini-2.0-flash-001"  # 性能与稳定性的平衡点

# --- 你的专属 HTML 声明 (完整保留) ---
INFO_HTML = """
<div style="text-align: left; border-left: 4px solid #2196F3; padding-left: 15px; margin-bottom: 20px;">
    <h3>MG TaxAI | 跨境财税合规实验室 (Beta)</h3>
    <p>本系统依托 <b>MG 核心智库</b> 构建，旨在实现解析结果实时溯源至各国官方税收协定与法律文本。目前系统正处于<b>知识库全量装载阶段</b>，已优先上线核心业务国家的官方协定库。</p>
    <p>我们正持续同步全球各主要经济体的国别投资税收指南及多税种年度税收报告。受限于测试版的数据填充进度，相关解析结果仅供专业参考。MG团队正加速完善每一条咨询建议的合规证据链，以确保交付专家级的数字化合规支持。</p>
    <hr style="border: 0; border-top: 1px solid #eee; margin: 10px 0;">
    <p style="font-size: 0.85em; color: #666;">
        <b>⚠️ AI 免责声明：</b><br>
        本系统生成的内容由人工智能根据现有库文件分析得出，不构成正式的法律或税务建议。在使用本系统结果进行任何商业决策前，请务必咨询 MG Consult 专业团队。
    </p>
</div>
"""

# --- 深度知识库检索引擎 (RAG) ---
def get_knowledge_context(query):
    """
    扫描本地 Treaties 和 InvestmentGuide 文件夹中的相关 PDF 内容。
    这里实现一个基础的关键词匹配检索，确保 AI 能“读”到书。
    """
    context_chunks = []
    base_dirs = ["Treaties", "InvestmentGuide"]
    
    # 简单的关键词提取（可以根据需要改进）
    keywords = [word for word in query.split() if len(word) > 1]
    
    for folder in base_dirs:
        path = Path(folder)
        if not path.exists(): continue
        
        # 扫描文件夹下的 PDF
        for pdf_file in path.rglob("*.pdf"):
            # 如果文件名包含关键词，优先读取
            if any(kw.lower() in pdf_file.name.lower() for kw in keywords):
                try:
                    with fitz.open(pdf_file) as doc:
                        # 提取前 2 页的核心内容（防止 Token 溢出）
                        text = "".join([page.get_text() for page in doc[:2]])
                        context_chunks.append(f"来自文件 [{pdf_file.name}]:\n{text}")
                except:
                    continue
    
    return "\n\n".join(context_chunks)[:4000] # 限制长度，保证 API 响应速度

# --- API 专家级调用逻辑 ---
def ask_ai(message, history):
    if not OPENROUTER_API_KEY:
        return "⚠️ 未检测到 API Key，请检查环境变量设置。"

    # 1. 获取本地知识库背景
    local_context = get_knowledge_context(message)

    # 2. 构造专家级指令 (System Prompt)
    # 这里决定了回复的质量和长度
    system_instruction = """
    你是一位资深的 MG Consulting 国际税务AI。
    你的回答必须具备以下特征：
    1. 深度：不仅给出结论，还要解释背后的税法逻辑。
    2. 广度：考虑多税种（所得税、增值税、预提税等）及双边协定影响。
    3. 严谨：优先引用提供的背景知识。如果背景中没有，请基于 2025-2026 年最新财税知识库回答。
    4. 结构：使用 Markdown 标题、列表和加粗，使回复易于阅读且专业。
    """

    # 3. 组装对话历史
    messages = [{"role": "system", "content": system_instruction}]
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    # 4. 组装当前请求
    current_input = f"【参考知识库内容】：\n{local_context}\n\n【用户当前咨询】：\n{message}"
    messages.append({"role": "user", "content": current_input})

    # 5. 发送请求
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"}
    payload = {
        "model": MODEL_ID,
        "messages": messages,
        "temperature": 0.2, # 降低随机性，提升严谨度
        "top_p": 0.9
    }

    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=90)
        if response.status_code == 200:
            return response.json()['choices'][0]['message']['content']
        return f"❌ 接口响应异常 ({response.status_code}): {response.text}"
    except Exception as e:
        return f"💥 系统连接超时或失败: {str(e)}"

# --- 界面构建 ---
with gr.Blocks(title="MG TaxAI Lab", fill_height=True, css=".gradio-container {background-color: #f9f9f9}") as demo:
    gr.HTML(INFO_HTML)
    
    chatbot = gr.ChatInterface(
        fn=ask_ai,
        fill_height=True,
        retry_btn="🔄 重新生成",
        undo_btn="↩️ 撤回上条",
        clear_btn="🗑️ 清空对话",
    )

if __name__ == "__main__":
    demo.launch()