"""
VSS VLM Path Test Script
Tests the /api/v1/query/media endpoint with sample images.

Usage:
    python scripts/test_vlm_path.py
    
Prerequisites:
    - VSS API running at localhost:8080
    - vLLM server (Cosmos) running on GPU 6
"""

import asyncio
import base64
import json
import time
import httpx
from pathlib import Path
from typing import Optional, Dict, Any

# === Configuration ===

VSS_API_URL = "http://localhost:8080/api/v1"
DATA_DIR = Path("/workspace/AI_Agents/VLM_Agent/VLM_8xH100_export/data/frames")

# Sample images by category
SAMPLE_IMAGES = {
    "fire_detection": [
        "2026-02-13-11-50-09-1.jpg",
        "2026-02-13-11-50-17.jpg",
    ],
    "people_detection": [
        "big-shopping-mall-sample-.png",
        "people-in-mall-hallway.png",
    ],
    "vehicle_detection": [
        "YT_Parking_LOT.png",
        "heavy-traffic-in-urban-city.png",
    ],
    "weapon_detection": [
        "cash-counter-attacker.jpg",
        "robbers-jewellery-store.jpg",
    ],
}

# Test queries for visual analysis
VLM_TEST_QUERIES = [
    "Describe what you see in this image.",
    "How many people are visible?",
    "Are there any vehicles in this scene?",
    "Is there anything unusual or concerning?",
    "Describe the scene and identify key objects.",
]


# === Utility Functions ===

def load_image_base64(image_path: Path) -> str:
    """Load image and convert to base64."""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def get_media_type(filename: str) -> str:
    """Get MIME type from filename."""
    ext = filename.lower().split(".")[-1]
    return {
        "jpg": "image/jpeg",
        "jpeg": "image/jpeg",
        "png": "image/png",
    }.get(ext, "image/jpeg")


# === API Client ===

class VSSTestClient:
    """Test client for VSS API."""
    
    def __init__(self, base_url: str = VSS_API_URL):
        self.base_url = base_url
        self.client = httpx.AsyncClient(timeout=120.0)
    
    async def health_check(self) -> Dict[str, Any]:
        """Check API health."""
        response = await self.client.get(f"{self.base_url}/health")
        return response.json()
    
    async def query_text(self, query: str, include_trace: bool = True) -> Dict[str, Any]:
        """Send text-only query."""
        payload = {
            "query": query,
            "options": {"include_trace": include_trace}
        }
        response = await self.client.post(f"{self.base_url}/query", json=payload)
        response.raise_for_status()
        return response.json()
    
    async def query_with_media(
        self,
        query: str,
        image_path: Path,
        include_trace: bool = True
    ) -> Dict[str, Any]:
        """Send query with image."""
        image_b64 = load_image_base64(image_path)
        filename = image_path.name
        
        payload = {
            "query": query,
            "media": {
                "type": "image",
                "data": image_b64,
                "filename": filename
            },
            "options": {"include_trace": include_trace}
        }
        
        response = await self.client.post(f"{self.base_url}/query/media", json=payload)
        response.raise_for_status()
        return response.json()
    
    async def close(self):
        await self.client.aclose()


# === Test Functions ===

async def test_health():
    """Test API health endpoint."""
    print("\n" + "=" * 60)
    print("TEST: Health Check")
    print("=" * 60)
    
    client = VSSTestClient()
    try:
        result = await client.health_check()
        print(f"Status: {result.get('status', 'unknown')}")
        
        components = result.get("components", {})
        for name, status in components.items():
            indicator = "✓" if status.get("healthy") else "✗"
            print(f"  {indicator} {name}: {status.get('status', 'unknown')}")
        
        return result.get("status") == "healthy"
    except Exception as e:
        print(f"✗ Health check failed: {e}")
        return False
    finally:
        await client.close()


async def test_vlm_direct():
    """Test VLM inference directly (bypass API)."""
    print("\n" + "=" * 60)
    print("TEST: Direct VLM Client")
    print("=" * 60)
    
    try:
        # Import from VSS package
        import sys
        sys.path.insert(0, "/workspace/AI_Agents/VSS_Orchestration")
        from vss.inference.vlm_client import VLMClient
        
        vlm = VLMClient()
        
        # Health check
        healthy = await vlm.health_check()
        print(f"VLM Server Health: {'✓ OK' if healthy else '✗ FAIL'}")
        
        if not healthy:
            return False
        
        # Test with sample image
        image_path = DATA_DIR / "vehicle_detection" / "YT_Parking_LOT.png"
        if not image_path.exists():
            print(f"✗ Sample image not found: {image_path}")
            return False
        
        print(f"Testing with: {image_path.name}")
        
        start = time.time()
        response = await vlm.describe_scene(
            images=[load_image_base64(image_path)],
            context="parking lot surveillance"
        )
        elapsed = time.time() - start
        
        print(f"Response ({elapsed:.2f}s):")
        print("-" * 40)
        print(response[:500] + "..." if len(response) > 500 else response)
        print("-" * 40)
        
        return True
        
    except ImportError as e:
        print(f"✗ Import error: {e}")
        print("  Make sure VSS package is in PYTHONPATH")
        return False
    except Exception as e:
        print(f"✗ VLM test failed: {e}")
        return False


async def test_api_vlm_endpoint():
    """Test /api/v1/query/media endpoint."""
    print("\n" + "=" * 60)
    print("TEST: API VLM Endpoint (/query/media)")
    print("=" * 60)
    
    client = VSSTestClient()
    results = []
    
    try:
        # Test with different categories
        test_cases = [
            ("vehicle_detection", "YT_Parking_LOT.png", "How many vehicles are in this parking lot?"),
            ("people_detection", "people-in-mall-hallway.png", "Describe the people and their activities."),
            ("fire_detection", "2026-02-13-11-50-17.jpg", "Is there any fire or smoke visible?"),
        ]
        
        for category, filename, query in test_cases:
            image_path = DATA_DIR / category / filename
            
            if not image_path.exists():
                print(f"\n✗ Image not found: {image_path}")
                results.append({"test": f"{category}/{filename}", "passed": False, "error": "File not found"})
                continue
            
            print(f"\n--- {category}: {filename} ---")
            print(f"Query: {query}")
            
            start = time.time()
            try:
                response = await client.query_with_media(query, image_path, include_trace=True)
                elapsed = time.time() - start
                
                print(f"Status: {response.get('status')}")
                print(f"Time: {response.get('processing_time_ms', int(elapsed*1000))}ms")
                
                # Print trace
                trace = response.get("trace", [])
                if trace:
                    print("Trace:")
                    for step in trace:
                        print(f"  → {step}")
                
                # Print response preview
                answer = response.get("response", "")
                print(f"Response: {answer[:200]}..." if len(answer) > 200 else f"Response: {answer}")
                
                # Metadata
                metadata = response.get("metadata", {})
                if metadata:
                    print(f"Intent: {metadata.get('intent')}, Source: {metadata.get('response_source')}")
                
                results.append({
                    "test": f"{category}/{filename}",
                    "passed": response.get("status") == "success",
                    "time_ms": response.get("processing_time_ms", int(elapsed*1000)),
                })
                
            except httpx.HTTPStatusError as e:
                print(f"✗ HTTP Error: {e.response.status_code}")
                print(f"  Detail: {e.response.text[:200]}")
                results.append({"test": f"{category}/{filename}", "passed": False, "error": str(e)})
            except Exception as e:
                print(f"✗ Error: {e}")
                results.append({"test": f"{category}/{filename}", "passed": False, "error": str(e)})
        
        return results
        
    finally:
        await client.close()


async def test_routing_with_image():
    """Test that image upload correctly routes to VLM."""
    print("\n" + "=" * 60)
    print("TEST: Routing Verification (Image → VLM)")
    print("=" * 60)
    
    client = VSSTestClient()
    
    try:
        image_path = DATA_DIR / "vehicle_detection" / "heavy-traffic-in-urban-city.png"
        
        if not image_path.exists():
            print(f"✗ Image not found: {image_path}")
            return False
        
        # Query that would normally go to SQL, but with image should go to VLM
        query = "How many cars are there?"
        
        print(f"Query: {query}")
        print(f"Image: {image_path.name}")
        print("Expected: Route to VLM (image upload forces VLM path)")
        
        response = await client.query_with_media(query, image_path, include_trace=True)
        
        metadata = response.get("metadata", {})
        intent = metadata.get("intent")
        source = metadata.get("response_source")
        
        print(f"\nActual Intent: {intent}")
        print(f"Response Source: {source}")
        
        # Verify it went to VLM
        if source == "vlm":
            print("✓ Correctly routed to VLM")
            return True
        else:
            print(f"✗ Unexpected routing: {source}")
            return False
        
    except Exception as e:
        print(f"✗ Test failed: {e}")
        return False
    finally:
        await client.close()


# === Main ===

async def main():
    print("=" * 60)
    print("VSS VLM Path Test Suite")
    print("=" * 60)
    
    results = {}
    
    # Run tests
    results["health"] = await test_health()
    
    if not results["health"]:
        print("\n✗ API not healthy, skipping remaining tests")
        return 1
    
    results["vlm_direct"] = await test_vlm_direct()
    results["api_endpoint"] = await test_api_vlm_endpoint()
    results["routing"] = await test_routing_with_image()
    
    # Summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    
    for name, result in results.items():
        if isinstance(result, list):
            passed = sum(1 for r in result if r.get("passed"))
            total = len(result)
            print(f"  {name}: {passed}/{total} passed")
        else:
            status = "✓ PASS" if result else "✗ FAIL"
            print(f"  {name}: {status}")
    
    all_passed = all(
        r if isinstance(r, bool) else all(x.get("passed") for x in r)
        for r in results.values()
    )
    
    print()
    if all_passed:
        print("All tests passed! VLM path is operational.")
    else:
        print("Some tests failed. Check output above.")
    
    return 0 if all_passed else 1


if __name__ == "__main__":
    exit(asyncio.run(main()))