#!/usr/bin/env python3
"""
Quick API test for SAM3 endpoint
Usage: python test_api.py
"""
import requests
import base64
import sys

ENDPOINT_URL = "https://p6irm2x7y9mwp4l4.us-east-1.aws.endpoints.huggingface.cloud"

def test_health():
    """Test health endpoint"""
    print("Testing /health endpoint...")
    response = requests.get(f"{ENDPOINT_URL}/health")

    if response.status_code == 200:
        data = response.json()
        print(f"✅ Health check passed")
        print(f"   Model: {data['model']}")
        print(f"   GPU: {'Available' if data['gpu_available'] else 'Not available'}")
        print(f"   VRAM: {data['vram']['free_gb']:.1f}GB free / {data['vram']['total_gb']:.1f}GB total")
        return True
    else:
        print(f"❌ Health check failed: {response.status_code}")
        return False

def test_inference():
    """Test inference with sample image"""
    print("\nTesting inference endpoint...")

    # Load test image
    import os
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(os.path.dirname(script_dir))
    test_image_path = os.path.join(project_root, "assets", "test_images", "test.jpg")

    try:
        with open(test_image_path, "rb") as f:
            image_b64 = base64.b64encode(f.read()).decode()
    except FileNotFoundError:
        print(f"❌ Test image not found at: {test_image_path}")
        return False

    # Make request
    response = requests.post(
        ENDPOINT_URL,
        json={
            "inputs": image_b64,
            "parameters": {
                "classes": ["pothole", "asphalt"]
            }
        },
        timeout=30
    )

    if response.status_code == 200:
        results = response.json()
        print(f"✅ Inference successful ({response.elapsed.total_seconds():.2f}s)")
        print(f"   Generated {len(results)} masks:")
        for result in results:
            mask_size = len(base64.b64decode(result['mask']))
            print(f"   - {result['label']}: {mask_size:,} bytes (score: {result['score']:.2f})")
        return True
    else:
        print(f"❌ Inference failed: {response.status_code}")
        print(f"   Response: {response.text}")
        return False

def main():
    print("=" * 60)
    print("SAM3 API Test")
    print("=" * 60)
    print(f"Endpoint: {ENDPOINT_URL}\n")

    health_ok = test_health()
    inference_ok = test_inference()

    print("\n" + "=" * 60)
    if health_ok and inference_ok:
        print("✅ All tests passed!")
        sys.exit(0)
    else:
        print("❌ Some tests failed")
        sys.exit(1)

if __name__ == "__main__":
    main()