import requests
import os
from PIL import Image
import io

def test_api(image_path, api_url="http://localhost:8000"):
    """测试纹身移除API"""
    
    # 检查输入图片是否存在
    if not os.path.exists(image_path):
        print(f"错误：找不到输入图片 {image_path}")
        return
    
    try:
        # 准备文件
        with open(image_path, 'rb') as f:
            files = {'file': (os.path.basename(image_path), f, 'image/jpeg')}
            
            # 发送请求
            print("正在发送请求到API...")
            response = requests.post(f"{api_url}/remove-tattoo", files=files)
            
            if response.status_code == 200:
                # 保存结果
                output_path = os.path.splitext(image_path)[0] + "_api_removed.png"
                image = Image.open(io.BytesIO(response.content))
                image.save(output_path)
                print(f"处理完成！结果已保存至：{output_path}")
            else:
                print(f"API请求失败，状态码：{response.status_code}")
                print(f"错误信息：{response.text}")
                
    except requests.exceptions.ConnectionError:
        print("连接错误：无法连接到API服务器")
        print("请确保API服务器正在运行（python -m uvicorn api.app:app --host 0.0.0.0 --port 8000）")
    except Exception as e:
        print(f"测试过程中出现错误：{str(e)}")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) != 2:
        print("使用方法：python test_api.py <图片路径>")
        print("例如：python test_api.py test_image.jpg")
        sys.exit(1)
    
    test_api(sys.argv[1])