"""
Test script for MedSAM HuggingFace Space
Run this after deploying your Space to verify it works
"""

import requests
import json
import base64
import numpy as np
from PIL import Image
from io import BytesIO
import sys

# UPDATE THIS after deploying your Space!
SPACE_URL = "https://YOUR_USERNAME-medsam-inference.hf.space/api/predict"


def test_space_with_image(image_path: str, x: int, y: int):
    """
    Test the MedSAM Space with an image
    
    Args:
        image_path: Path to test image
        x: X coordinate for segmentation point
        y: Y coordinate for segmentation point
    """
    print(f"🧪 Testing MedSAM Space: {SPACE_URL}")
    print(f"   Image: {image_path}")
    print(f"   Point: ({x}, {y})")
    print()
    
    try:
        # 1. Load and encode image
        print("📸 Loading image...")
        with open(image_path, "rb") as f:
            image_bytes = f.read()
        
        image = Image.open(BytesIO(image_bytes))
        print(f"   Size: {image.size}")
        print(f"   Mode: {image.mode}")
        
        # Encode as base64
        img_base64 = base64.b64encode(image_bytes).decode()
        print(f"   Base64 size: {len(img_base64)} chars")
        print()
        
        # 2. Prepare points JSON
        print("📍 Preparing points...")
        points_json = json.dumps({
            "coords": [[x, y]],
            "labels": [1],  # 1 = foreground
            "multimask_output": True
        })
        print(f"   Points JSON: {points_json}")
        print()
        
        # 3. Call API
        print("🚀 Calling Space API...")
        response = requests.post(
            SPACE_URL,
            json={
                "data": [
                    f"data:image/jpeg;base64,{img_base64}",
                    points_json
                ]
            },
            timeout=120
        )
        
        print(f"   Status code: {response.status_code}")
        
        if response.status_code != 200:
            print(f"❌ Error: {response.status_code}")
            print(f"   Response: {response.text}")
            return False
        
        print()
        
        # 4. Parse result
        print("📊 Parsing result...")
        result = response.json()
        
        # Gradio wraps output in data array
        if "data" not in result or len(result["data"]) == 0:
            print("❌ Error: Unexpected response format")
            print(f"   Response: {json.dumps(result, indent=2)}")
            return False
        
        output_json = result["data"][0]
        output = json.loads(output_json)
        
        if not output.get("success", False):
            print(f"❌ Error: {output.get('error', 'Unknown error')}")
            return False
        
        print("✅ Success!")
        print(f"   Number of masks: {output['num_masks']}")
        print(f"   Scores: {output['scores']}")
        print()
        
        # 5. Process masks
        print("🎭 Processing masks...")
        for i, (mask_data, score) in enumerate(zip(output['masks'], output['scores'])):
            mask_array = np.array(mask_data['mask_data'], dtype=bool)
            print(f"   Mask {i+1}:")
            print(f"      Shape: {mask_array.shape}")
            print(f"      Score: {score:.4f}")
            print(f"      Pixels: {np.sum(mask_array)} / {mask_array.size}")
            print(f"      Coverage: {100 * np.sum(mask_array) / mask_array.size:.2f}%")
        
        # 6. Get best mask
        best_idx = np.argmax(output['scores'])
        best_mask = np.array(output['masks'][best_idx]['mask_data'], dtype=bool)
        best_score = output['scores'][best_idx]
        
        print()
        print(f"🏆 Best mask: #{best_idx+1} (score: {best_score:.4f})")
        print()
        
        # 7. Save visualization
        print("💾 Saving visualization...")
        
        # Create visualization
        image_array = np.array(image)
        
        # Create colored mask overlay
        mask_overlay = np.zeros((*best_mask.shape, 3), dtype=np.uint8)
        mask_overlay[best_mask] = [255, 0, 0]  # Red
        
        # Blend with original image
        if len(image_array.shape) == 2:  # Grayscale
            image_array = np.stack([image_array] * 3, axis=-1)
        
        blended = image_array.copy()
        blended[best_mask] = (
            0.6 * image_array[best_mask] + 
            0.4 * mask_overlay[best_mask]
        ).astype(np.uint8)
        
        # Save
        output_path = "test_result_visualization.png"
        Image.fromarray(blended).save(output_path)
        print(f"   Saved: {output_path}")
        
        # Save mask only
        mask_path = "test_result_mask.png"
        Image.fromarray((best_mask * 255).astype(np.uint8)).save(mask_path)
        print(f"   Saved: {mask_path}")
        print()
        
        print("=" * 60)
        print("✅ TEST PASSED! Your Space is working correctly!")
        print("=" * 60)
        
        return True
        
    except requests.exceptions.Timeout:
        print("❌ Error: Request timeout (>120 seconds)")
        print("   The Space might be sleeping or overloaded")
        print("   Try again in 30 seconds")
        return False
        
    except requests.exceptions.RequestException as e:
        print(f"❌ Error: Request failed: {e}")
        return False
        
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False


def check_space_status(space_url: str):
    """Check if the Space is online"""
    print(f"🔍 Checking Space status: {space_url}")
    
    try:
        # Try to access the Space homepage
        homepage_url = space_url.replace("/api/predict", "")
        response = requests.get(homepage_url, timeout=10)
        
        if response.status_code == 200:
            print("✅ Space is online!")
            return True
        else:
            print(f"⚠️  Space returned status {response.status_code}")
            return False
            
    except requests.exceptions.RequestException as e:
        print(f"❌ Cannot reach Space: {e}")
        print("   Make sure you've deployed the Space and updated SPACE_URL")
        return False


if __name__ == "__main__":
    print("=" * 60)
    print("MedSAM HuggingFace Space Test")
    print("=" * 60)
    print()
    
    # Check if SPACE_URL is updated
    if "YOUR_USERNAME" in SPACE_URL:
        print("❌ Error: Please update SPACE_URL in this script!")
        print("   Replace YOUR_USERNAME with your HuggingFace username")
        print()
        print("   Example:")
        print('   SPACE_URL = "https://johndoe-medsam-inference.hf.space/api/predict"')
        sys.exit(1)
    
    # Check Space status
    check_space_status(SPACE_URL)
    print()
    
    # Get test image
    if len(sys.argv) < 2:
        print("Usage: python test_space.py <image_path> [x] [y]")
        print()
        print("Example:")
        print("   python test_space.py test_image.jpg 200 150")
        print()
        sys.exit(1)
    
    image_path = sys.argv[1]
    x = int(sys.argv[2]) if len(sys.argv) > 2 else 200
    y = int(sys.argv[3]) if len(sys.argv) > 3 else 150
    
    # Run test
    success = test_space_with_image(image_path, x, y)
    
    sys.exit(0 if success else 1)