#!/usr/bin/env python3
"""
Test script for the new endpoint-per-model API architecture
"""

import requests
import json
import time

def test_endpoint_api(base_url="https://aurasystems-spanish-embeddings-api.hf.space"):
    """Test the new endpoint-based API"""
    
    print(f"Testing Endpoint-Based API at {base_url}")
    print("=" * 60)
    
    # Test root endpoint
    try:
        response = requests.get(f"{base_url}/")
        print(f"✓ Root endpoint: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Version: {data.get('version', 'N/A')}")
            print(f"  Startup model: {data.get('startup_model', 'N/A')}")
            print(f"  Available endpoints: {list(data.get('available_endpoints', {}).keys())}")
        else:
            print(f"  Error: {response.text}")
            return False
    except Exception as e:
        print(f"✗ Root endpoint failed: {e}")
        return False
    
    # Test health endpoint
    try:
        response = requests.get(f"{base_url}/health")
        print(f"✓ Health endpoint: {response.status_code}")
        if response.status_code == 200:
            health_data = response.json()
            print(f"  Startup model loaded: {health_data.get('startup_model_loaded', False)}")
            print(f"  Available models: {health_data.get('available_models', [])}")
            print(f"  Models count: {health_data.get('models_count', 0)}")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ Health endpoint failed: {e}")
    
    print("\n" + "=" * 60)
    print("TESTING MODEL ENDPOINTS")
    print("=" * 60)
    
    # Test jina-v3 endpoint (startup model)
    try:
        payload = {
            "texts": ["Hello world", "Bonjour le monde", "Hola mundo"],
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed/jina-v3", json=payload)
        print(f"✓ Jina-v3 endpoint: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Model: {data.get('model_used', 'N/A')}")
            print(f"  Embeddings: {data.get('num_texts', 0)} texts → {data.get('dimensions', 0)} dimensions")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ Jina-v3 endpoint failed: {e}")
    
    # Test roberta-ca endpoint (on-demand)
    try:
        payload = {
            "texts": ["Bon dia", "Com estàs?", "Catalunya és meravellosa"],
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed/roberta-ca", json=payload)
        print(f"✓ RoBERTa-ca endpoint: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Model: {data.get('model_used', 'N/A')}")
            print(f"  Embeddings: {data.get('num_texts', 0)} texts → {data.get('dimensions', 0)} dimensions")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ RoBERTa-ca endpoint failed: {e}")
    
    # Test jina endpoint (on-demand)
    try:
        payload = {
            "texts": ["Texto en español", "Text in English"],
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed/jina", json=payload)
        print(f"✓ Jina endpoint: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Model: {data.get('model_used', 'N/A')}")
            print(f"  Embeddings: {data.get('num_texts', 0)} texts → {data.get('dimensions', 0)} dimensions")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ Jina endpoint failed: {e}")
    
    # Test robertalex endpoint (Spanish legal)
    try:
        payload = {
            "texts": ["Artículo primero de la constitución", "El contrato será válido"],
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed/robertalex", json=payload)
        print(f"✓ RoBERTalex endpoint: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Model: {data.get('model_used', 'N/A')}")
            print(f"  Embeddings: {data.get('num_texts', 0)} texts → {data.get('dimensions', 0)} dimensions")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ RoBERTalex endpoint failed: {e}")
    
    # Test legal-bert endpoint (English legal)
    try:
        payload = {
            "texts": ["This agreement is legally binding", "The contract shall be valid"],
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed/legal-bert", json=payload)
        print(f"✓ Legal-BERT endpoint: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Model: {data.get('model_used', 'N/A')}")
            print(f"  Embeddings: {data.get('num_texts', 0)} texts → {data.get('dimensions', 0)} dimensions")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ Legal-BERT endpoint failed: {e}")
    
    print("\n" + "=" * 60)
    print("FINAL HEALTH CHECK")
    print("=" * 60)
    
    # Final health check to see all loaded models
    try:
        response = requests.get(f"{base_url}/health")
        if response.status_code == 200:
            health_data = response.json()
            print(f"✓ Final status: {health_data.get('status', 'unknown')}")
            print(f"  Available models: {health_data.get('available_models', [])}")
            print(f"  Total models loaded: {health_data.get('models_count', 0)}/5")
            
            endpoints = health_data.get('endpoints', {})
            for model, status in endpoints.items():
                print(f"  {model}: {status}")
    except Exception as e:
        print(f"✗ Final health check failed: {e}")
    
    return True

if __name__ == "__main__":
    test_endpoint_api()