#!/usr/bin/env python
#coding=utf-8

import requests
from config import MINERU_API_URL, MINERU_TOKEN

def test_create_task():
    """
    测试创建解析任务
    """
    print("测试创建解析任务...")
    
    api_url = f"{MINERU_API_URL}/extract/task"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {MINERU_TOKEN}"
    }
    # 使用官方示例 URL
    data = {
        "url": "https://cdn-mineru.openxlab.org.cn/demo/example.pdf",
        "model_version": "vlm"
    }
    
    try:
        response = requests.post(api_url, headers=headers, json=data)
        print(f"状态码: {response.status_code}")
        print(f"响应内容: {response.json()}")
        
        if response.status_code == 200:
            result = response.json()
            if result.get("code") == 0:
                task_id = result.get("data", {}).get("task_id")
                print(f"任务创建成功，任务ID: {task_id}")
                return task_id
            else:
                print(f"任务创建失败: {result.get('msg')}")
                return None
        else:
            print(f"请求失败，状态码: {response.status_code}")
            return None
            
    except Exception as e:
        print(f"测试失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def test_get_task_status(task_id):
    """
    测试获取任务状态
    """
    if not task_id:
        print("任务ID为空，跳过测试")
        return None
    
    print(f"测试获取任务状态...")
    
    api_url = f"{MINERU_API_URL}/extract/task/{task_id}"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {MINERU_TOKEN}"
    }
    
    try:
        response = requests.get(api_url, headers=headers)
        print(f"状态码: {response.status_code}")
        print(f"响应内容: {response.json()}")
        
        if response.status_code == 200:
            result = response.json()
            if result.get("code") == 0:
                status = result.get("data", {})
                print(f"任务状态获取成功，状态: {status.get('state')}")
                return status
            else:
                print(f"获取状态失败: {result.get('msg')}")
                return None
        else:
            print(f"请求失败，状态码: {response.status_code}")
            return None
            
    except Exception as e:
        print(f"测试失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def main():
    """
    主函数
    """
    print("开始测试 Mineru API...")
    
    # 测试创建任务
    task_id = test_create_task()
    
    # 测试获取任务状态
    if task_id:
        test_get_task_status(task_id)
    
    print("API 测试完成")


if __name__ == "__main__":
    main()