#!/usr/bin/env python3
"""
Comprehensive test for custom task status support in optimized queries.
Tests both shots and assets with various custom status scenarios.
"""

import requests
import json
import time

BASE_URL = 'http://localhost:8000'

def login():
    """Login and get access token"""
    response = requests.post(f'{BASE_URL}/auth/login', json={'email': 'admin@vfx.com', 'password': 'admin123'})
    if response.status_code != 200:
        print("❌ Login failed")
        return None
    return response.json()['access_token']

def test_custom_status_filtering(token, project_id):
    """Test filtering by custom task status"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Custom Status Filtering ===")
    
    # Test shot filtering with custom status
    response = requests.get(f'{BASE_URL}/shots/?project_id={project_id}&task_status_filter=layout:custom_9441a740', headers=headers)
    if response.status_code == 200:
        shots = response.json()
        filtered_shots = [s for s in shots if s.get('task_status', {}).get('layout') == 'custom_9441a740']
        print(f"✅ Shot filtering with custom status: {len(filtered_shots)} shots found")
    else:
        print(f"❌ Shot filtering failed: {response.status_code}")
    
    # Test asset filtering with custom status
    response = requests.get(f'{BASE_URL}/assets/?project_id={project_id}&task_status_filter=modeling:custom_9441a740', headers=headers)
    if response.status_code == 200:
        assets = response.json()
        filtered_assets = [a for a in assets if a.get('task_status', {}).get('modeling') == 'custom_9441a740']
        print(f"✅ Asset filtering with custom status: {len(filtered_assets)} assets found")
    else:
        print(f"❌ Asset filtering failed: {response.status_code}")

def test_custom_status_sorting(token, project_id):
    """Test sorting by custom task status"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Custom Status Sorting ===")
    
    # Test shot sorting by task status
    response = requests.get(f'{BASE_URL}/shots/?project_id={project_id}&sort_by=layout_status&sort_direction=asc', headers=headers)
    if response.status_code == 200:
        shots = response.json()
        if len(shots) > 1:
            # Check if sorting is working (custom statuses should come after system statuses)
            statuses = [s.get('task_status', {}).get('layout', 'not_started') for s in shots]
            print(f"✅ Shot sorting by layout status: {statuses}")
        else:
            print("ℹ️ Only one shot, cannot verify sorting order")
    else:
        print(f"❌ Shot sorting failed: {response.status_code}")
    
    # Test asset sorting by task status
    response = requests.get(f'{BASE_URL}/assets/?project_id={project_id}&sort_by=modeling_status&sort_direction=asc', headers=headers)
    if response.status_code == 200:
        assets = response.json()
        if len(assets) > 1:
            statuses = [a.get('task_status', {}).get('modeling', 'not_started') for a in assets]
            print(f"✅ Asset sorting by modeling status: {statuses}")
        else:
            print("ℹ️ Only one asset, cannot verify sorting order")
    else:
        print(f"❌ Asset sorting failed: {response.status_code}")

def test_mixed_status_scenarios(token, project_id):
    """Test scenarios with mixed system and custom statuses"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Mixed Status Scenarios ===")
    
    # Get all task statuses for the project
    response = requests.get(f'{BASE_URL}/projects/{project_id}/task-statuses', headers=headers)
    if response.status_code == 200:
        all_statuses = response.json()
        system_statuses = [s['id'] for s in all_statuses['system_statuses']]
        custom_statuses = [s['id'] for s in all_statuses['statuses']]
        
        print(f"System statuses: {system_statuses}")
        print(f"Custom statuses: {custom_statuses}")
        
        # Verify that optimized queries include both types
        response = requests.get(f'{BASE_URL}/shots/?project_id={project_id}', headers=headers)
        if response.status_code == 200:
            shots = response.json()
            if shots:
                shot = shots[0]
                task_status = shot.get('task_status', {})
                
                # Check for both system and custom statuses
                found_system = any(status in system_statuses for status in task_status.values())
                found_custom = any(status in custom_statuses for status in task_status.values())
                
                print(f"✅ Shot contains system statuses: {found_system}")
                print(f"✅ Shot contains custom statuses: {found_custom}")
        
        response = requests.get(f'{BASE_URL}/assets/?project_id={project_id}', headers=headers)
        if response.status_code == 200:
            assets = response.json()
            if assets:
                asset = assets[0]
                task_status = asset.get('task_status', {})
                
                # Check for both system and custom statuses
                found_system = any(status in system_statuses for status in task_status.values())
                found_custom = any(status in custom_statuses for status in task_status.values())
                
                print(f"✅ Asset contains system statuses: {found_system}")
                print(f"✅ Asset contains custom statuses: {found_custom}")
    else:
        print(f"❌ Failed to get project task statuses: {response.status_code}")

def test_performance_with_custom_statuses(token, project_id):
    """Test that performance is maintained with custom statuses"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Performance with Custom Statuses ===")
    
    # Time the shot query
    start_time = time.time()
    response = requests.get(f'{BASE_URL}/shots/?project_id={project_id}', headers=headers)
    shot_time = time.time() - start_time
    
    if response.status_code == 200:
        shots = response.json()
        print(f"✅ Shot query with custom statuses: {len(shots)} shots in {shot_time:.3f}s")
    else:
        print(f"❌ Shot query failed: {response.status_code}")
    
    # Time the asset query
    start_time = time.time()
    response = requests.get(f'{BASE_URL}/assets/?project_id={project_id}', headers=headers)
    asset_time = time.time() - start_time
    
    if response.status_code == 200:
        assets = response.json()
        print(f"✅ Asset query with custom statuses: {len(assets)} assets in {asset_time:.3f}s")
    else:
        print(f"❌ Asset query failed: {response.status_code}")
    
    # Both queries should be fast (under 500ms as per requirements)
    if shot_time < 0.5 and asset_time < 0.5:
        print("✅ Performance requirements met (< 500ms)")
    else:
        print(f"⚠️ Performance concern: shot={shot_time:.3f}s, asset={asset_time:.3f}s")

def test_data_consistency(token, project_id):
    """Test data consistency between task_status and task_details"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Data Consistency ===")
    
    # Test shots
    response = requests.get(f'{BASE_URL}/shots/?project_id={project_id}', headers=headers)
    if response.status_code == 200:
        shots = response.json()
        for shot in shots:
            task_status = shot.get('task_status', {})
            task_details = shot.get('task_details', [])
            
            # Create a map from task_details for comparison
            details_status = {td['task_type']: td['status'] for td in task_details}
            
            # Check consistency
            inconsistencies = []
            for task_type, status in task_status.items():
                if task_type in details_status and details_status[task_type] != status:
                    inconsistencies.append(f"{task_type}: {status} vs {details_status[task_type]}")
            
            if inconsistencies:
                print(f"❌ Shot {shot['name']} has inconsistencies: {inconsistencies}")
            else:
                print(f"✅ Shot {shot['name']} data is consistent")
    
    # Test assets
    response = requests.get(f'{BASE_URL}/assets/?project_id={project_id}', headers=headers)
    if response.status_code == 200:
        assets = response.json()
        for asset in assets:
            task_status = asset.get('task_status', {})
            task_details = asset.get('task_details', [])
            
            # Create a map from task_details for comparison
            details_status = {td['task_type']: td['status'] for td in task_details}
            
            # Check consistency
            inconsistencies = []
            for task_type, status in task_status.items():
                if task_type in details_status and details_status[task_type] != status:
                    inconsistencies.append(f"{task_type}: {status} vs {details_status[task_type]}")
            
            if inconsistencies:
                print(f"❌ Asset {asset['name']} has inconsistencies: {inconsistencies}")
            else:
                print(f"✅ Asset {asset['name']} data is consistent")

def main():
    print("=" * 80)
    print("Comprehensive Custom Task Status Support Test")
    print("=" * 80)
    
    # Login
    token = login()
    if not token:
        return
    
    # Use the existing test project (ID 10) that has custom statuses
    project_id = 10
    
    # Run all tests
    test_custom_status_filtering(token, project_id)
    test_custom_status_sorting(token, project_id)
    test_mixed_status_scenarios(token, project_id)
    test_performance_with_custom_statuses(token, project_id)
    test_data_consistency(token, project_id)
    
    print("\n" + "=" * 80)
    print("🎉 Comprehensive Custom Status Support Test Complete!")
    print("=" * 80)

if __name__ == "__main__":
    main()