import gradio as gr
import torch
import os
from PIL import Image
import tempfile
import subprocess
import sys

# 安装MonkeyOCR
def install_monkeyocr():
    """安装MonkeyOCR依赖"""
    try:
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", 
            "git+https://github.com/Yuliang-Liu/MonkeyOCR.git"
        ])
        return True
    except:
        return False

# 在启动时安装
if not install_monkeyocr():
    print("Failed to install MonkeyOCR")

try:
    from monkeyocr import MonkeyOCR
except ImportError:
    print("MonkeyOCR not available, using placeholder")

class MonkeyOCRDemo:
    def __init__(self):
        self.model = None
        self.load_model()
    
    def load_model(self):
        """加载MonkeyOCR模型"""
        try:
            # 这里需要根据实际的MonkeyOCR API调整
            self.model = MonkeyOCR.from_pretrained("echo840/MonkeyOCR")
            print("MonkeyOCR模型加载成功")
        except Exception as e:
            print(f"模型加载失败: {e}")
            self.model = None
    
    def parse_document(self, file_input, prompt="解析这个文档"):
        """解析文档"""
        if self.model is None:
            return "模型未加载，请稍后再试", "模型加载失败"
        
        try:
            # 处理上传的文件
            if file_input is None:
                return "请上传一个文件", "错误：没有文件"
            
            # 根据文件类型处理
            file_path = file_input.name if hasattr(file_input, 'name') else file_input
            
            # 调用MonkeyOCR进行解析
            # 这里需要根据实际的API调整
            result = self.model.parse(file_path)
            
            # 返回解析结果
            markdown_result = result.get('markdown', '解析失败')
            status = "解析成功"
            
            return markdown_result, status
            
        except Exception as e:
            error_msg = f"解析过程中出现错误: {str(e)}"
            return error_msg, "错误"

# 创建Demo实例
demo_instance = MonkeyOCRDemo()

def parse_with_monkeyocr(file_input, prompt="解析文档内容"):
    """Gradio包装函数"""
    return demo_instance.parse_document(file_input, prompt)

# 创建Gradio界面
def create_interface():
    with gr.Blocks(title="MonkeyOCR Document Parser", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🐵 MonkeyOCR 文档解析器
        
        MonkeyOCR采用Structure-Recognition-Relation (SRR)三元组范式，能够高效解析中英文文档。
        
        **支持格式**: PDF, PNG, JPG, JPEG
        **功能**: 文本识别、公式解析、表格提取、结构化输出
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="上传文档",
                    file_types=[".pdf", ".png", ".jpg", ".jpeg"],
                    type="filepath"
                )
                
                prompt_input = gr.Textbox(
                    label="提示词（可选）",
                    placeholder="请描述您希望如何解析这个文档...",
                    value="解析这个文档的所有内容，包括文本、表格和公式",
                    lines=2
                )
                
                parse_btn = gr.Button("🚀 开始解析", variant="primary", size="lg")
                
                gr.Markdown("""
                ### 使用说明
                1. 上传PDF文档或图片
                2. 输入解析提示词（可选）
                3. 点击"开始解析"按钮
                4. 等待解析完成，查看结果
                """)
            
            with gr.Column(scale=2):
                status_output = gr.Textbox(
                    label="解析状态",
                    value="等待上传文件...",
                    interactive=False
                )
                
                result_output = gr.Textbox(
                    label="解析结果（Markdown格式）",
                    lines=20,
                    max_lines=30,
                    show_copy_button=True
                )
        
        # 示例文件
        gr.Examples(
            examples=[
                ["./examples/sample.pdf", "解析这个学术论文"],
                ["./examples/table.png", "提取表格数据"],
                ["./examples/formula.jpg", "识别数学公式"],
            ],
            inputs=[file_input, prompt_input],
            outputs=[result_output, status_output],
            fn=parse_with_monkeyocr,
            cache_examples=False
        )
        
        # 绑定事件
        parse_btn.click(
            fn=parse_with_monkeyocr,
            inputs=[file_input, prompt_input],
            outputs=[result_output, status_output]
        )
        
        # 文件上传时的处理
        file_input.change(
            fn=lambda x: "文件已上传，请点击解析按钮" if x else "等待上传文件...",
            inputs=[file_input],
            outputs=[status_output]
        )
    
    return demo

# 启动应用
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )