#!/usr/bin/env python3

"""
Test script to validate the asset router optimization implementation.
This script tests the optimized asset endpoints to ensure they work correctly
and provide the expected performance improvements.
"""

import requests
import json
import time
from typing import Dict, List

# Configuration
BASE_URL = "http://localhost:8000"
API_BASE = f"{BASE_URL}/api"

def get_auth_headers() -> Dict[str, str]:
    """Get authentication headers for API requests."""
    # Login to get access token
    login_data = {
        "username": "admin",
        "password": "admin123"
    }
    
    response = requests.post(f"{API_BASE}/auth/login", json=login_data)
    if response.status_code != 200:
        raise Exception(f"Login failed: {response.text}")
    
    token_data = response.json()
    access_token = token_data["access_token"]
    
    return {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json"
    }

def test_list_assets_optimization():
    """Test the optimized list_assets endpoint."""
    print("Testing optimized list_assets endpoint...")
    
    headers = get_auth_headers()
    
    # Test basic asset listing
    start_time = time.time()
    response = requests.get(f"{API_BASE}/assets/", headers=headers)
    end_time = time.time()
    
    print(f"List assets response time: {end_time - start_time:.3f}s")
    
    if response.status_code != 200:
        print(f"❌ List assets failed: {response.status_code} - {response.text}")
        return False
    
    assets = response.json()
    print(f"✅ Retrieved {len(assets)} assets")
    
    # Verify response structure
    if assets:
        asset = assets[0]
        required_fields = ['id', 'name', 'category', 'task_status', 'task_details', 'task_count']
        missing_fields = [field for field in required_fields if field not in asset]
        
        if missing_fields:
            print(f"❌ Missing fields in asset response: {missing_fields}")
            return False
        
        print(f"✅ Asset response structure is correct")
        print(f"   - Task status: {asset.get('task_status', {})}")
        print(f"   - Task details count: {len(asset.get('task_details', []))}")
        print(f"   - Task count: {asset.get('task_count', 0)}")
    
    return True

def test_get_asset_optimization():
    """Test the optimized get_asset endpoint."""
    print("\nTesting optimized get_asset endpoint...")
    
    headers = get_auth_headers()
    
    # First get list of assets to find an asset ID
    response = requests.get(f"{API_BASE}/assets/", headers=headers)
    if response.status_code != 200:
        print("❌ Could not get assets list for testing")
        return False
    
    assets = response.json()
    if not assets:
        print("❌ No assets found for testing")
        return False
    
    asset_id = assets[0]['id']
    
    # Test single asset retrieval
    start_time = time.time()
    response = requests.get(f"{API_BASE}/assets/{asset_id}", headers=headers)
    end_time = time.time()
    
    print(f"Get asset response time: {end_time - start_time:.3f}s")
    
    if response.status_code != 200:
        print(f"❌ Get asset failed: {response.status_code} - {response.text}")
        return False
    
    asset = response.json()
    print(f"✅ Retrieved asset: {asset['name']}")
    
    # Verify response structure
    required_fields = ['id', 'name', 'category', 'task_count']
    missing_fields = [field for field in required_fields if field not in asset]
    
    if missing_fields:
        print(f"❌ Missing fields in asset response: {missing_fields}")
        return False
    
    print(f"✅ Asset response structure is correct")
    print(f"   - Task count: {asset.get('task_count', 0)}")
    
    return True

def test_asset_task_status_endpoint():
    """Test the asset task status endpoint."""
    print("\nTesting asset task status endpoint...")
    
    headers = get_auth_headers()
    
    # First get list of assets to find an asset ID
    response = requests.get(f"{API_BASE}/assets/", headers=headers)
    if response.status_code != 200:
        print("❌ Could not get assets list for testing")
        return False
    
    assets = response.json()
    if not assets:
        print("❌ No assets found for testing")
        return False
    
    asset_id = assets[0]['id']
    
    # Test asset task status retrieval
    start_time = time.time()
    response = requests.get(f"{API_BASE}/assets/{asset_id}/task-status", headers=headers)
    end_time = time.time()
    
    print(f"Get asset task status response time: {end_time - start_time:.3f}s")
    
    if response.status_code != 200:
        print(f"❌ Get asset task status failed: {response.status_code} - {response.text}")
        return False
    
    task_details = response.json()
    print(f"✅ Retrieved {len(task_details)} task details")
    
    # Verify response structure
    if task_details:
        task = task_details[0]
        required_fields = ['task_type', 'status', 'task_id']
        missing_fields = [field for field in required_fields if field not in task]
        
        if missing_fields:
            print(f"❌ Missing fields in task detail response: {missing_fields}")
            return False
        
        print(f"✅ Task detail response structure is correct")
        print(f"   - Task type: {task.get('task_type')}")
        print(f"   - Status: {task.get('status')}")
    
    return True

def test_asset_filtering():
    """Test asset filtering functionality."""
    print("\nTesting asset filtering...")
    
    headers = get_auth_headers()
    
    # Test filtering by category
    response = requests.get(f"{API_BASE}/assets/?category=characters", headers=headers)
    if response.status_code == 200:
        assets = response.json()
        print(f"✅ Category filter works: {len(assets)} character assets")
    else:
        print(f"⚠️  Category filter test skipped: {response.status_code}")
    
    # Test task status filtering
    response = requests.get(f"{API_BASE}/assets/?task_status_filter=modeling:not_started", headers=headers)
    if response.status_code == 200:
        assets = response.json()
        print(f"✅ Task status filter works: {len(assets)} assets with modeling not_started")
    else:
        print(f"⚠️  Task status filter test skipped: {response.status_code}")
    
    return True

def main():
    """Run all optimization tests."""
    print("🚀 Starting Asset Router Optimization Tests")
    print("=" * 50)
    
    try:
        # Test all endpoints
        tests = [
            test_list_assets_optimization,
            test_get_asset_optimization,
            test_asset_task_status_endpoint,
            test_asset_filtering
        ]
        
        results = []
        for test in tests:
            try:
                result = test()
                results.append(result)
            except Exception as e:
                print(f"❌ Test failed with exception: {e}")
                results.append(False)
        
        # Summary
        print("\n" + "=" * 50)
        print("📊 Test Results Summary")
        print("=" * 50)
        
        passed = sum(results)
        total = len(results)
        
        print(f"✅ Passed: {passed}/{total}")
        if passed == total:
            print("🎉 All asset router optimization tests passed!")
        else:
            print(f"❌ {total - passed} tests failed")
        
        return passed == total
        
    except Exception as e:
        print(f"❌ Test suite failed: {e}")
        return False

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)