#!/usr/bin/env python3
"""
🧪 Test script untuk FunCaptcha Solver API
Test authentication, endpoints, dan response formats
"""

import os
import base64
import json
import requests
import time
from typing import Optional
import argparse

# Configuration
DEFAULT_LOCAL_URL = "http://localhost:7860"
DEFAULT_HF_URL = "https://your-space-name.hf.space"
DEFAULT_API_KEY = "test-key-123"

class FunCaptchaAPITester:
    def __init__(self, base_url: str, api_key: str):
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key
        self.session = requests.Session()
        
    def test_health_check(self) -> bool:
        """Test health check endpoint (no auth required)"""
        print("🔍 Testing health check endpoint...")
        
        try:
            response = self.session.get(f"{self.base_url}/health", timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                status = data.get('status', 'unknown')
                ml_backend = data.get('ml_backend', 'none')
                onnx_available = data.get('onnx_runtime_available', False)
                pytorch_available = data.get('pytorch_available', False)
                tensorflow_available = data.get('tensorflow_available', False)
                warnings = data.get('warnings', [])
                
                print(f"✅ Health check passed")
                print(f"   Status: {status}")
                print(f"   ML Backend: {ml_backend}")
                print(f"   ONNX Runtime: {'✅ Available' if onnx_available else '❌ Not Available'}")
                print(f"   PyTorch: {'✅ Available' if pytorch_available else '❌ Not Available'}")
                print(f"   TensorFlow: {'✅ Available' if tensorflow_available else '❌ Not Available'}")
                print(f"   Models loaded: {data.get('models_loaded', 0)}")
                print(f"   Cache entries: {data.get('cache_entries', 0)}")
                
                if warnings:
                    print(f"   ⚠️ Warnings:")
                    for warning in warnings:
                        print(f"      - {warning}")
                
                return True
            else:
                print(f"❌ Health check failed: {response.status_code}")
                print(f"   Response: {response.text}")
                return False
                
        except Exception as e:
            print(f"❌ Health check error: {e}")
            return False
    
    def test_root_endpoint(self) -> bool:
        """Test root endpoint (no auth required)"""
        print("🔍 Testing root endpoint...")
        
        try:
            response = self.session.get(f"{self.base_url}/", timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                print(f"✅ Root endpoint accessible")
                print(f"   Service: {data.get('service')}")
                print(f"   Version: {data.get('version')}")
                return True
            else:
                print(f"❌ Root endpoint failed: {response.status_code}")
                return False
                
        except Exception as e:
            print(f"❌ Root endpoint error: {e}")
            return False
    
    def test_authentication(self) -> bool:
        """Test API authentication"""
        print("🔍 Testing API authentication...")
        
        # Test without API key
        try:
            response = self.session.post(f"{self.base_url}/solve", timeout=10, json={
                "challenge_type": "pick_the",
                "image_b64": "data:image/png;base64,test",
                "target_label": "test"
            })
            
            if response.status_code == 401:
                print("✅ Authentication correctly rejects requests without API key")
            else:
                print(f"⚠️  Expected 401 for missing auth, got {response.status_code}")
                
        except Exception as e:
            print(f"❌ Auth test error: {e}")
            return False
        
        # Test with wrong API key
        try:
            headers = {"Authorization": "Bearer wrong-key"}
            response = self.session.post(f"{self.base_url}/solve", 
                                       headers=headers, timeout=10, json={
                "challenge_type": "pick_the", 
                "image_b64": "data:image/png;base64,test",
                "target_label": "test"
            })
            
            if response.status_code == 401:
                print("✅ Authentication correctly rejects wrong API key")
                return True
            else:
                print(f"⚠️  Expected 401 for wrong key, got {response.status_code}")
                return False
                
        except Exception as e:
            print(f"❌ Wrong auth test error: {e}")
            return False
    
    def create_test_image_b64(self) -> str:
        """Create a simple test image in base64 format"""
        # Simple 1x1 pixel PNG in base64
        # This is a minimal valid PNG image
        png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
        return f"data:image/png;base64,{png_b64}"
    
    def test_solve_endpoint_pick_the(self) -> bool:
        """Test solve endpoint dengan pick_the challenge"""
        print("🔍 Testing solve endpoint (pick_the)...")
        
        try:
            headers = {"Authorization": f"Bearer {self.api_key}"}
            test_image = self.create_test_image_b64()
            
            payload = {
                "challenge_type": "pick_the",
                "image_b64": test_image,
                "target_label": "ice cream"
            }
            
            start_time = time.time()
            response = self.session.post(f"{self.base_url}/solve", 
                                       headers=headers, 
                                       json=payload, 
                                       timeout=30)
            response_time = time.time() - start_time
            
            print(f"   Response time: {response_time:.2f}s")
            print(f"   Status code: {response.status_code}")
            
            if response.status_code == 200:
                data = response.json()
                status = data.get('status')
                print("✅ Solve endpoint accessible with valid auth")
                print(f"   Status: {status}")
                print(f"   Processing time: {data.get('processing_time', 0):.3f}s")
                
                if status == 'error':
                    print(f"   Error message: {data.get('message', 'Unknown error')}")
                    if 'No ML backend available' in data.get('message', '') or 'ONNX Runtime not available' in data.get('message', ''):
                        print("   ℹ️ This is expected jika ML backend tidak tersedia")
                        return True  # Still consider this as success - API working properly
                
                if 'box' in data:
                    print(f"   Box coordinates: {data['box']}")
                if 'confidence' in data:
                    print(f"   Confidence: {data['confidence']:.3f}")
                    
                return True
            else:
                print(f"❌ Solve endpoint failed: {response.status_code}")
                print(f"   Response: {response.text}")
                return False
                
        except Exception as e:
            print(f"❌ Solve endpoint error: {e}")
            return False
    
    def test_solve_endpoint_upright(self) -> bool:
        """Test solve endpoint dengan upright challenge"""
        print("🔍 Testing solve endpoint (upright)...")
        
        try:
            headers = {"Authorization": f"Bearer {self.api_key}"}
            test_image = self.create_test_image_b64()
            
            payload = {
                "challenge_type": "upright",
                "image_b64": test_image
            }
            
            start_time = time.time()
            response = self.session.post(f"{self.base_url}/solve", 
                                       headers=headers, 
                                       json=payload, 
                                       timeout=30)
            response_time = time.time() - start_time
            
            print(f"   Response time: {response_time:.2f}s")
            print(f"   Status code: {response.status_code}")
            
            if response.status_code == 200:
                data = response.json()
                status = data.get('status')
                print("✅ Upright solve endpoint works")
                print(f"   Status: {status}")
                print(f"   Processing time: {data.get('processing_time', 0):.3f}s")
                
                if status == 'error':
                    print(f"   Error message: {data.get('message', 'Unknown error')}")
                    if 'No ML backend available' in data.get('message', '') or 'ONNX Runtime not available' in data.get('message', ''):
                        print("   ℹ️ This is expected jika ML backend tidak tersedia")
                        return True  # Still consider this as success - API working properly
                
                if 'button_index' in data:
                    print(f"   Button index: {data['button_index']}")
                if 'confidence' in data:
                    print(f"   Confidence: {data['confidence']:.3f}")
                    
                return True
            else:
                print(f"❌ Upright solve failed: {response.status_code}")
                print(f"   Response: {response.text}")
                return False
                
        except Exception as e:
            print(f"❌ Upright solve error: {e}")
            return False
    
    def test_invalid_requests(self) -> bool:
        """Test invalid request handling"""
        print("🔍 Testing invalid request handling...")
        
        headers = {"Authorization": f"Bearer {self.api_key}"}
        
        # Test 1: Invalid challenge type
        try:
            response = self.session.post(f"{self.base_url}/solve",
                                       headers=headers,
                                       json={
                                           "challenge_type": "invalid_type",
                                           "image_b64": self.create_test_image_b64()
                                       },
                                       timeout=10)
            
            if response.status_code == 400:
                print("✅ Correctly rejects invalid challenge type")
            else:
                print(f"⚠️  Expected 400 for invalid challenge type, got {response.status_code}")
        
        except Exception as e:
            print(f"❌ Invalid challenge type test error: {e}")
            return False
        
        # Test 2: Missing target_label for pick_the
        try:
            response = self.session.post(f"{self.base_url}/solve",
                                       headers=headers,
                                       json={
                                           "challenge_type": "pick_the", 
                                           "image_b64": self.create_test_image_b64()
                                           # Missing target_label
                                       },
                                       timeout=10)
            
            if response.status_code == 400:
                print("✅ Correctly rejects pick_the without target_label")
            else:
                print(f"⚠️  Expected 400 for missing target_label, got {response.status_code}")
        
        except Exception as e:
            print(f"❌ Missing target_label test error: {e}")
            return False
        
        return True
    
    def run_all_tests(self) -> bool:
        """Run all test cases"""
        print("🚀 Starting FunCaptcha API tests...")
        print(f"🌐 Base URL: {self.base_url}")
        print(f"🔑 API Key: {self.api_key[:8]}...")
        print("=" * 60)
        
        tests = [
            ("Health Check", self.test_health_check),
            ("Root Endpoint", self.test_root_endpoint),
            ("Authentication", self.test_authentication),
            ("Solve Pick-The", self.test_solve_endpoint_pick_the),
            ("Solve Upright", self.test_solve_endpoint_upright),
            ("Invalid Requests", self.test_invalid_requests)
        ]
        
        results = []
        for test_name, test_func in tests:
            print(f"\\n📋 Running test: {test_name}")
            print("-" * 40)
            try:
                result = test_func()
                results.append((test_name, result))
                if result:
                    print(f"✅ {test_name} PASSED")
                else:
                    print(f"❌ {test_name} FAILED") 
            except Exception as e:
                print(f"💥 {test_name} CRASHED: {e}")
                results.append((test_name, False))
            
            time.sleep(1)  # Brief pause between tests
        
        # Summary
        print("\\n" + "=" * 60)
        print("📊 TEST SUMMARY")
        print("=" * 60)
        
        passed = sum(1 for _, result in results if result)
        total = len(results)
        
        for test_name, result in results:
            status = "✅ PASS" if result else "❌ FAIL"
            print(f"{test_name:<20} {status}")
        
        print(f"\\n🎯 Overall: {passed}/{total} tests passed")
        
        if passed == total:
            print("🎉 All tests PASSED! API is working correctly.")
            return True
        else:
            print(f"⚠️  {total - passed} tests FAILED. Check configuration and deployment.")
            return False

def main():
    parser = argparse.ArgumentParser(description="Test FunCaptcha Solver API")
    parser.add_argument("--url", default=DEFAULT_LOCAL_URL, 
                       help=f"Base URL for API (default: {DEFAULT_LOCAL_URL})")
    parser.add_argument("--api-key", default=DEFAULT_API_KEY,
                       help=f"API key for authentication (default: {DEFAULT_API_KEY})")
    parser.add_argument("--hf", action="store_true",
                       help="Use default Hugging Face URL")
    
    args = parser.parse_args()
    
    # Use HF URL if --hf flag provided
    if args.hf:
        base_url = DEFAULT_HF_URL
        print("🤗 Using Hugging Face Spaces URL")
        print("⚠️  Make sure to update DEFAULT_HF_URL in script dengan actual space URL!")
    else:
        base_url = args.url
    
    # Check for API key in environment
    api_key = os.getenv("FUNCAPTCHA_API_KEY", args.api_key)
    
    if api_key == DEFAULT_API_KEY and args.hf:
        print("⚠️  WARNING: Using default API key dengan HF Spaces!")
        print("   Set FUNCAPTCHA_API_KEY environment variable atau use --api-key")
    
    # Run tests
    tester = FunCaptchaAPITester(base_url, api_key)
    success = tester.run_all_tests()
    
    # Exit with appropriate code
    exit(0 if success else 1)

if __name__ == "__main__":
    main()