#!/usr/bin/env python3
"""
Test script for MagicDoc API
"""

import requests
import json
import os

def test_health_check(base_url="http://localhost:8002"):
    """Test health check endpoint"""
    try:
        response = requests.get(f"{base_url}/health")
        print(f"Health check status: {response.status_code}")
        print(f"Response: {response.json()}")
        return response.status_code == 200
    except Exception as e:
        print(f"Health check failed: {e}")
        return False

def test_file_parse(base_url="http://localhost:8002", file_path=None):
    """Test file parse endpoint"""
    if not file_path or not os.path.exists(file_path):
        print(f"File not found: {file_path}")
        return False
    
    try:
        with open(file_path, 'rb') as f:
            files = {'files': (os.path.basename(file_path), f, 'application/octet-stream')}
            data = {
                'output_dir': './output',
                'lang_list': 'ch',
                'backend': 'pipeline',
                'parse_method': 'auto',
                'formula_enable': True,
                'table_enable': True,
                'return_md': True,
                'return_middle_json': False,
                'return_model_output': False,
                'return_content_list': False,
                'return_images': False,
                'start_page_id': 0,
                'end_page_id': 99999
            }
            
            response = requests.post(f"{base_url}/file_parse", files=files, data=data)
            print(f"File parse status: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                print(f"Success! Converted {len(result.get('markdown', ''))} characters")
                print(f"Time cost: {result.get('time_cost', 'N/A')}s")
                return True
            else:
                print(f"Error: {response.text}")
                return False
                
    except Exception as e:
        print(f"File parse failed: {e}")
        return False

def main():
    """Main test function"""
    print("Testing MagicDoc API...")
    
    # Test health check
    print("\n1. Testing health check...")
    if not test_health_check():
        print("Health check failed. Make sure the service is running.")
        return
    
    # Test file parse (if sample file exists)
    print("\n2. Testing file parse...")
    sample_files = [
        "../sample_doc/20220707_na_decision-2.docx",
        "../sample_doc/20220707_na_decision-2.pdf",
        "../sample_doc/short_doc.md"
    ]
    
    for sample_file in sample_files:
        if os.path.exists(sample_file):
            print(f"Testing with {sample_file}...")
            if test_file_parse(file_path=sample_file):
                print("File parse test passed!")
                break
        else:
            print(f"Sample file not found: {sample_file}")
    
    print("\nTest completed!")

if __name__ == "__main__":
    main()