#!/usr/bin/env python3
"""
HCFS API v2 Test Client

Comprehensive test client for validating the production API functionality.
"""

import asyncio
import json
import time
from typing import List, Dict, Any
import httpx
import websocket
import threading

class HCFSAPIClient:
    """Test client for HCFS API v2."""
    
    def __init__(self, base_url: str = "http://localhost:8000", api_key: str = "dev-key-123"):
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key
        self.headers = {
            "X-API-Key": api_key,
            "Content-Type": "application/json"
        }
    
    async def test_health_check(self) -> Dict[str, Any]:
        """Test health check endpoint."""
        print("🔍 Testing health check...")
        
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{self.base_url}/health")
            
            print(f"   Status: {response.status_code}")
            if response.status_code == 200:
                data = response.json()
                print(f"   System Status: {data['status']}")
                print(f"   Components: {len(data['components'])}")
                return data
            else:
                print(f"   Error: {response.text}")
                return {}
    
    async def test_context_crud(self) -> Dict[str, Any]:
        """Test context CRUD operations."""
        print("\n📋 Testing Context CRUD operations...")
        
        results = {"create": False, "read": False, "update": False, "delete": False}
        
        async with httpx.AsyncClient() as client:
            # Test Create
            create_data = {
                "path": "/test/api_test",
                "content": "This is a test context for API validation",
                "summary": "API test context",
                "author": "test_client"
            }
            
            response = await client.post(
                f"{self.base_url}/api/v1/contexts",
                json=create_data,
                headers=self.headers
            )
            
            if response.status_code == 200:
                created_context = response.json()["data"]
                context_id = created_context["id"]
                print(f"   ✅ Create: Context {context_id} created")
                results["create"] = True
                
                # Test Read
                response = await client.get(
                    f"{self.base_url}/api/v1/contexts/{context_id}",
                    headers=self.headers
                )
                
                if response.status_code == 200:
                    read_context = response.json()["data"]
                    print(f"   ✅ Read: Context {context_id} retrieved")
                    results["read"] = True
                    
                    # Test Update
                    update_data = {
                        "content": "Updated test context content",
                        "summary": "Updated summary"
                    }
                    
                    response = await client.put(
                        f"{self.base_url}/api/v1/contexts/{context_id}",
                        json=update_data,
                        headers=self.headers
                    )
                    
                    if response.status_code == 200:
                        print(f"   ✅ Update: Context {context_id} updated")
                        results["update"] = True
                    else:
                        print(f"   ❌ Update failed: {response.status_code}")
                    
                    # Test Delete
                    response = await client.delete(
                        f"{self.base_url}/api/v1/contexts/{context_id}",
                        headers=self.headers
                    )
                    
                    if response.status_code == 200:
                        print(f"   ✅ Delete: Context {context_id} deleted")
                        results["delete"] = True
                    else:
                        print(f"   ❌ Delete failed: {response.status_code}")
                        
                else:
                    print(f"   ❌ Read failed: {response.status_code}")
            else:
                print(f"   ❌ Create failed: {response.status_code} - {response.text}")
        
        return results
    
    async def test_search_functionality(self) -> Dict[str, Any]:
        """Test search functionality."""
        print("\n🔍 Testing Search functionality...")
        
        results = {"semantic": False, "hybrid": False, "keyword": False}
        
        # First, create some test contexts
        test_contexts = [
            {
                "path": "/ml/algorithms",
                "content": "Machine learning algorithms and neural networks for data analysis",
                "summary": "ML algorithms overview",
                "author": "test_client"
            },
            {
                "path": "/web/development",
                "content": "Web development using FastAPI and modern frameworks",
                "summary": "Web dev guide",
                "author": "test_client"
            },
            {
                "path": "/database/systems",
                "content": "Database management systems and SQL optimization techniques",
                "summary": "Database guide",
                "author": "test_client"
            }
        ]
        
        context_ids = []
        
        async with httpx.AsyncClient() as client:
            # Create test contexts
            for context_data in test_contexts:
                response = await client.post(
                    f"{self.base_url}/api/v1/contexts",
                    json=context_data,
                    headers=self.headers
                )
                if response.status_code == 200:
                    context_ids.append(response.json()["data"]["id"])
            
            print(f"   Created {len(context_ids)} test contexts")
            
            # Wait a moment for embeddings to be generated
            await asyncio.sleep(2)
            
            # Test Semantic Search
            search_data = {
                "query": "machine learning neural networks",
                "search_type": "semantic",
                "top_k": 5
            }
            
            response = await client.post(
                f"{self.base_url}/api/v1/search",
                json=search_data,
                headers=self.headers
            )
            
            if response.status_code == 200:
                search_results = response.json()
                print(f"   ✅ Semantic Search: {search_results['total_results']} results in {search_results['search_time_ms']:.2f}ms")
                results["semantic"] = True
            else:
                print(f"   ❌ Semantic Search failed: {response.status_code}")
            
            # Test Hybrid Search
            search_data["search_type"] = "hybrid"
            search_data["semantic_weight"] = 0.7
            
            response = await client.post(
                f"{self.base_url}/api/v1/search",
                json=search_data,
                headers=self.headers
            )
            
            if response.status_code == 200:
                search_results = response.json()
                print(f"   ✅ Hybrid Search: {search_results['total_results']} results in {search_results['search_time_ms']:.2f}ms")
                results["hybrid"] = True
            else:
                print(f"   ❌ Hybrid Search failed: {response.status_code}")
            
            # Test Keyword Search
            search_data["search_type"] = "keyword"
            
            response = await client.post(
                f"{self.base_url}/api/v1/search",
                json=search_data,
                headers=self.headers
            )
            
            if response.status_code == 200:
                search_results = response.json()
                print(f"   ✅ Keyword Search: {search_results['total_results']} results")
                results["keyword"] = True
            else:
                print(f"   ❌ Keyword Search failed: {response.status_code}")
            
            # Cleanup test contexts
            for context_id in context_ids:
                await client.delete(
                    f"{self.base_url}/api/v1/contexts/{context_id}",
                    headers=self.headers
                )
        
        return results
    
    async def test_batch_operations(self) -> Dict[str, Any]:
        """Test batch operations."""
        print("\n📦 Testing Batch operations...")
        
        batch_contexts = [
            {
                "path": f"/batch/test_{i}",
                "content": f"Batch test context {i} with sample content",
                "summary": f"Batch context {i}",
                "author": "batch_client"
            }
            for i in range(5)
        ]
        
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.base_url}/api/v1/contexts/batch",
                json={"contexts": batch_contexts},
                headers=self.headers
            )
            
            if response.status_code == 200:
                batch_result = response.json()["data"]
                print(f"   ✅ Batch Create: {batch_result['success_count']}/{batch_result['total_items']} succeeded")
                
                # Cleanup
                for context_id in batch_result["created_ids"]:
                    await client.delete(
                        f"{self.base_url}/api/v1/contexts/{context_id}",
                        headers=self.headers
                    )
                
                return {"batch_create": True, "success_rate": batch_result['success_count'] / batch_result['total_items']}
            else:
                print(f"   ❌ Batch Create failed: {response.status_code}")
                return {"batch_create": False, "success_rate": 0.0}
    
    async def test_pagination(self) -> Dict[str, Any]:
        """Test pagination functionality."""
        print("\n📄 Testing Pagination...")
        
        # Create multiple contexts for pagination testing
        contexts = [
            {
                "path": f"/pagination/test_{i}",
                "content": f"Pagination test context {i}",
                "summary": f"Context {i}",
                "author": "pagination_client"
            }
            for i in range(15)
        ]
        
        context_ids = []
        
        async with httpx.AsyncClient() as client:
            # Create contexts
            for context_data in contexts:
                response = await client.post(
                    f"{self.base_url}/api/v1/contexts",
                    json=context_data,
                    headers=self.headers
                )
                if response.status_code == 200:
                    context_ids.append(response.json()["data"]["id"])
            
            # Test pagination
            response = await client.get(
                f"{self.base_url}/api/v1/contexts?page=1&page_size=5&path_prefix=/pagination",
                headers=self.headers
            )
            
            if response.status_code == 200:
                page_data = response.json()
                pagination_info = page_data["pagination"]
                
                print(f"   ✅ Page 1: {len(page_data['data'])} items")
                print(f"   Total: {pagination_info['total_items']}, Pages: {pagination_info['total_pages']}")
                print(f"   Has Next: {pagination_info['has_next']}, Has Previous: {pagination_info['has_previous']}")
                
                # Cleanup
                for context_id in context_ids:
                    await client.delete(
                        f"{self.base_url}/api/v1/contexts/{context_id}",
                        headers=self.headers
                    )
                
                return {
                    "pagination_working": True,
                    "total_items": pagination_info['total_items'],
                    "items_per_page": len(page_data['data'])
                }
            else:
                print(f"   ❌ Pagination failed: {response.status_code}")
                return {"pagination_working": False}
    
    async def test_statistics_endpoint(self) -> Dict[str, Any]:
        """Test statistics endpoint."""
        print("\n📊 Testing Statistics endpoint...")
        
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{self.base_url}/api/v1/stats",
                headers=self.headers
            )
            
            if response.status_code == 200:
                stats = response.json()
                print(f"   ✅ Statistics retrieved")
                print(f"   Total Contexts: {stats['context_stats']['total_contexts']}")
                print(f"   Active Connections: {stats['system_stats']['active_connections']}")
                print(f"   Cache Hit Rate: {stats['system_stats']['cache_hit_rate']:.2%}")
                return {"stats_available": True, "data": stats}
            else:
                print(f"   ❌ Statistics failed: {response.status_code}")
                return {"stats_available": False}
    
    def test_websocket_connection(self) -> Dict[str, Any]:
        """Test WebSocket connection."""
        print("\n🔌 Testing WebSocket connection...")
        
        try:
            ws_url = self.base_url.replace("http", "ws") + "/ws"
            
            def on_message(ws, message):
                print(f"   📨 WebSocket message: {message}")
            
            def on_error(ws, error):
                print(f"   ❌ WebSocket error: {error}")
            
            def on_close(ws, close_status_code, close_msg):
                print(f"   🔐 WebSocket closed")
            
            def on_open(ws):
                print(f"   ✅ WebSocket connected")
                # Send subscription request
                subscription = {
                    "type": "subscribe",
                    "data": {
                        "path_prefix": "/test",
                        "event_types": ["created", "updated", "deleted"]
                    }
                }
                ws.send(json.dumps(subscription))
                
                # Close after a moment
                threading.Timer(2.0, ws.close).start()
            
            ws = websocket.WebSocketApp(
                ws_url,
                on_open=on_open,
                on_message=on_message,
                on_error=on_error,
                on_close=on_close
            )
            
            ws.run_forever(ping_interval=30, ping_timeout=10)
            return {"websocket_working": True}
            
        except Exception as e:
            print(f"   ❌ WebSocket test failed: {e}")
            return {"websocket_working": False}
    
    async def test_error_handling(self) -> Dict[str, Any]:
        """Test error handling."""
        print("\n🚨 Testing Error handling...")
        
        results = {}
        
        async with httpx.AsyncClient() as client:
            # Test 404 - Non-existent context
            response = await client.get(
                f"{self.base_url}/api/v1/contexts/999999",
                headers=self.headers
            )
            
            if response.status_code == 404:
                print("   ✅ 404 handling works")
                results["404_handling"] = True
            else:
                print(f"   ❌ Expected 404, got {response.status_code}")
                results["404_handling"] = False
            
            # Test 422 - Invalid data
            invalid_data = {
                "path": "",  # Invalid empty path
                "content": "",  # Invalid empty content
            }
            
            response = await client.post(
                f"{self.base_url}/api/v1/contexts",
                json=invalid_data,
                headers=self.headers
            )
            
            if response.status_code == 422:
                print("   ✅ Validation error handling works")
                results["validation_handling"] = True
            else:
                print(f"   ❌ Expected 422, got {response.status_code}")
                results["validation_handling"] = False
        
        return results
    
    async def run_comprehensive_test(self) -> Dict[str, Any]:
        """Run all tests comprehensively."""
        print("🧪 HCFS API v2 Comprehensive Test Suite")
        print("=" * 50)
        
        start_time = time.time()
        all_results = {}
        
        # Run all tests
        all_results["health"] = await self.test_health_check()
        all_results["crud"] = await self.test_context_crud()
        all_results["search"] = await self.test_search_functionality()
        all_results["batch"] = await self.test_batch_operations()
        all_results["pagination"] = await self.test_pagination()
        all_results["statistics"] = await self.test_statistics_endpoint()
        all_results["errors"] = await self.test_error_handling()
        
        # WebSocket test (runs synchronously)
        print("\n🔌 Testing WebSocket (this may take a moment)...")
        all_results["websocket"] = self.test_websocket_connection()
        
        total_time = time.time() - start_time
        
        # Generate summary
        print(f"\n📋 TEST SUMMARY")
        print("=" * 30)
        
        total_tests = 0
        passed_tests = 0
        
        for category, results in all_results.items():
            if isinstance(results, dict):
                category_tests = len([v for v in results.values() if isinstance(v, bool)])
                category_passed = len([v for v in results.values() if v is True])
                total_tests += category_tests
                passed_tests += category_passed
                
                if category_tests > 0:
                    success_rate = (category_passed / category_tests) * 100
                    print(f"   {category.upper()}: {category_passed}/{category_tests} ({success_rate:.1f}%)")
        
        overall_success_rate = (passed_tests / total_tests) * 100 if total_tests > 0 else 0
        
        print(f"\n🎯 OVERALL RESULTS:")
        print(f"   Tests Passed: {passed_tests}/{total_tests}")
        print(f"   Success Rate: {overall_success_rate:.1f}%")
        print(f"   Total Time: {total_time:.2f}s")
        
        if overall_success_rate >= 80:
            print(f"   Status: ✅ API IS PRODUCTION READY!")
        elif overall_success_rate >= 60:
            print(f"   Status: ⚠️  API needs some improvements")
        else:
            print(f"   Status: ❌ API has significant issues")
        
        return {
            "summary": {
                "total_tests": total_tests,
                "passed_tests": passed_tests,
                "success_rate": overall_success_rate,
                "total_time": total_time
            },
            "detailed_results": all_results
        }


async def main():
    """Main function to run API tests."""
    import argparse
    
    parser = argparse.ArgumentParser(description="HCFS API v2 Test Client")
    parser.add_argument("--url", default="http://localhost:8000", help="API base URL")
    parser.add_argument("--api-key", default="dev-key-123", help="API key for authentication")
    parser.add_argument("--test", choices=["all", "health", "crud", "search", "batch", "websocket"], 
                       default="all", help="Specific test to run")
    
    args = parser.parse_args()
    
    client = HCFSAPIClient(base_url=args.url, api_key=args.api_key)
    
    if args.test == "all":
        await client.run_comprehensive_test()
    elif args.test == "health":
        await client.test_health_check()
    elif args.test == "crud":
        await client.test_context_crud()
    elif args.test == "search":
        await client.test_search_functionality()
    elif args.test == "batch":
        await client.test_batch_operations()
    elif args.test == "websocket":
        client.test_websocket_connection()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n🛑 Test interrupted by user")
    except Exception as e:
        print(f"\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()