import os
import uuid
import shutil
import subprocess
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse
import uvicorn

app = FastAPI(title="MinerU Extraction API")

@app.post("/extract")
def extract_document(file: UploadFile = File(...)):
    # 1. 创建独立的任务文件夹，防止多个请求同时发生时文件冲突
    task_id = str(uuid.uuid4())[:8]
    work_dir = os.path.abspath(f"./api_workspace/{task_id}")
    out_dir = os.path.join(work_dir, "OUT")
    os.makedirs(work_dir, exist_ok=True)
    
    input_file_path = os.path.join(work_dir, file.filename)
    
    try:
        # 2. 接收并保存客户端上传的文件
        print(f"[Task {task_id}] 收到文件: {file.filename}，开始处理...")
        with open(input_file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        # 3. 调用你刚刚跑通的 MinerU 命令行
        # 相当于在终端执行 mineru -p xxx -o xxx
        subprocess.run(["mineru", "-p", input_file_path, "-o", out_dir], check=True)
        
        # 4. 将 MinerU 的输出结果打包成 ZIP 文件
        print(f"[Task {task_id}] 解析完成，正在打包结果...")
        zip_base_path = os.path.join(work_dir, "result")
        shutil.make_archive(zip_base_path, 'zip', out_dir)
        zip_file_path = f"{zip_base_path}.zip"
        
        # 5. 将 ZIP 文件作为响应返回给客户端
        return FileResponse(
            zip_file_path,
            media_type="application/zip",
            filename=f"parsed_{file.filename}.zip"
        )
        
    except Exception as e:
        return {"status": "error", "message": str(e)}

if __name__ == "__main__":
    # 绑定到 0.0.0.0 代表监听本机所有 IP (包括 192.168.4.6)
    # 端口设置为你要求的 5000
    print("🚀 MinerU API 启动成功！监听地址: http://192.168.4.6:5000")
    uvicorn.run(app, host="0.0.0.0", port=5000)