#!/usr/bin/env python3
"""
Test custom task status support in optimized shot and asset queries.
This test verifies that the optimized queries include both default and custom task statuses.
"""

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 create_test_project(token):
    """Create a test project for custom status testing"""
    headers = {'Authorization': f'Bearer {token}'}
    
    project_data = {
        'name': f'Custom Status Test Project {int(time.time())}',
        'code_name': f'CSTP{int(time.time())}',
        'client_name': 'Test Client',
        'project_type': 'tv',
        'description': 'Test project for custom status optimization'
    }
    
    response = requests.post(f'{BASE_URL}/projects/', json=project_data, headers=headers)
    if response.status_code == 201:
        project = response.json()
        print(f"✅ Created test project: {project['name']} (ID: {project['id']})")
        return project['id']
    else:
        print(f"❌ Failed to create test project: {response.status_code}")
        return None

def create_custom_task_statuses(token, project_id):
    """Create custom task statuses for the project"""
    headers = {'Authorization': f'Bearer {token}'}
    
    custom_statuses = [
        {'name': 'Ready for Review', 'color': '#8B5CF6'},
        {'name': 'On Hold', 'color': '#FFA500'},
        {'name': 'Blocked', 'color': '#DC2626'}
    ]
    
    created_statuses = []
    for status_data in custom_statuses:
        response = requests.post(
            f'{BASE_URL}/projects/{project_id}/task-statuses',
            json=status_data,
            headers=headers
        )
        if response.status_code == 201:
            status = response.json()['status']
            created_statuses.append(status)
            print(f"✅ Created custom status: {status['name']} ({status['id']})")
        else:
            print(f"❌ Failed to create custom status {status_data['name']}: {response.status_code}")
    
    return created_statuses

def create_test_episode(token, project_id):
    """Create a test episode"""
    headers = {'Authorization': f'Bearer {token}'}
    
    episode_data = {
        'name': 'Test Episode',
        'episode_number': 1,
        'description': 'Test episode for custom status testing'
    }
    
    response = requests.post(f'{BASE_URL}/episodes/?project_id={project_id}', json=episode_data, headers=headers)
    if response.status_code == 201:
        episode = response.json()
        print(f"✅ Created test episode: {episode['name']} (ID: {episode['id']})")
        return episode['id']
    else:
        print(f"❌ Failed to create test episode: {response.status_code}")
        return None

def create_test_shot(token, episode_id):
    """Create a test shot"""
    headers = {'Authorization': f'Bearer {token}'}
    
    shot_data = {
        'name': 'TEST_001',
        'description': 'Test shot for custom status testing',
        'frame_start': 1001,
        'frame_end': 1100
    }
    
    response = requests.post(f'{BASE_URL}/shots/?episode_id={episode_id}', json=shot_data, headers=headers)
    if response.status_code == 201:
        shot = response.json()
        print(f"✅ Created test shot: {shot['name']} (ID: {shot['id']})")
        return shot['id']
    else:
        print(f"❌ Failed to create test shot: {response.status_code}")
        return None

def create_test_asset(token, project_id):
    """Create a test asset"""
    headers = {'Authorization': f'Bearer {token}'}
    
    asset_data = {
        'name': 'TestCharacter',
        'category': 'characters',
        'description': 'Test asset for custom status testing'
    }
    
    response = requests.post(f'{BASE_URL}/assets/?project_id={project_id}', json=asset_data, headers=headers)
    if response.status_code == 201:
        asset = response.json()
        print(f"✅ Created test asset: {asset['name']} (ID: {asset['id']})")
        return asset['id']
    else:
        print(f"❌ Failed to create test asset: {response.status_code}")
        return None

def create_task_with_custom_status(token, shot_id=None, asset_id=None, task_type='layout', custom_status_id='ready_for_review'):
    """Create a task and update it to use a custom status"""
    headers = {'Authorization': f'Bearer {token}'}
    
    # Create the task
    if shot_id:
        response = requests.post(f'{BASE_URL}/shots/{shot_id}/tasks?task_type={task_type}', headers=headers)
        entity_type = 'shot'
        entity_id = shot_id
    else:
        response = requests.post(f'{BASE_URL}/assets/{asset_id}/tasks?task_type={task_type}', headers=headers)
        entity_type = 'asset'
        entity_id = asset_id
    
    task_id = None
    if response.status_code == 201:
        task_info = response.json()
        task_id = task_info['task_id']
        print(f"✅ Created {entity_type} task: {task_type} (ID: {task_id})")
    elif response.status_code == 400 and "already exists" in response.text:
        # Task already exists, get the existing task info
        try:
            task_info = response.json()
            task_id = task_info.get('task_id')
            if task_id:
                print(f"ℹ️ {entity_type.title()} task already exists: {task_type} (ID: {task_id})")
            else:
                print(f"ℹ️ {entity_type.title()} task already exists: {task_type} (no ID returned)")
                return None
        except:
            print(f"ℹ️ {entity_type.title()} task already exists: {task_type} (could not parse response)")
            return None
    else:
        print(f"❌ Failed to create {entity_type} task: {response.status_code} - {response.text}")
        return None
    
    if not task_id:
        return None
    
    # Update task status to custom status
    update_data = {'status': custom_status_id}
    response = requests.put(f'{BASE_URL}/tasks/{task_id}', json=update_data, headers=headers)
    if response.status_code == 200:
        print(f"✅ Updated task {task_id} to custom status: {custom_status_id}")
        return task_id
    else:
        print(f"❌ Failed to update task status: {response.status_code} - {response.text}")
        return task_id

def test_shot_optimization_with_custom_statuses(token, project_id, episode_id):
    """Test that optimized shot queries include custom task statuses"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Shot Optimization with Custom Statuses ===")
    
    # Get shots with optimized query
    response = requests.get(f'{BASE_URL}/shots/?project_id={project_id}', headers=headers)
    if response.status_code != 200:
        print(f"❌ Failed to get shots: {response.status_code}")
        return False
    
    shots = response.json()
    if not shots:
        print("❌ No shots found")
        return False
    
    shot = shots[0]
    task_status = shot.get('task_status', {})
    task_details = shot.get('task_details', [])
    
    print(f"Shot: {shot['name']}")
    print(f"Task status keys: {list(task_status.keys())}")
    print(f"Task details count: {len(task_details)}")
    
    # Check for custom statuses in task_status
    custom_status_found = False
    for task_type, status in task_status.items():
        if status not in ['not_started', 'in_progress', 'submitted', 'approved', 'retake']:
            custom_status_found = True
            print(f"✅ Found custom status in task_status: {task_type} = {status}")
    
    # Check for custom statuses in task_details
    custom_status_in_details = False
    for task_detail in task_details:
        status = task_detail.get('status')
        if status not in ['not_started', 'in_progress', 'submitted', 'approved', 'retake']:
            custom_status_in_details = True
            print(f"✅ Found custom status in task_details: {task_detail['task_type']} = {status}")
    
    # Verify all required fields are present in task_details
    required_fields = ['task_type', 'status', 'task_id', 'assigned_user_id']
    all_fields_present = True
    for task_detail in task_details:
        missing_fields = [f for f in required_fields if f not in task_detail]
        if missing_fields:
            print(f"❌ Missing fields in task details: {missing_fields}")
            all_fields_present = False
    
    if all_fields_present:
        print("✅ All required fields present in task_details")
    
    return custom_status_found or custom_status_in_details

def test_asset_optimization_with_custom_statuses(token, project_id):
    """Test that optimized asset queries include custom task statuses"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Asset Optimization with Custom Statuses ===")
    
    # Get assets with optimized query
    response = requests.get(f'{BASE_URL}/assets/?project_id={project_id}', headers=headers)
    if response.status_code != 200:
        print(f"❌ Failed to get assets: {response.status_code}")
        return False
    
    assets = response.json()
    if not assets:
        print("❌ No assets found")
        return False
    
    asset = assets[0]
    task_status = asset.get('task_status', {})
    task_details = asset.get('task_details', [])
    
    print(f"Asset: {asset['name']}")
    print(f"Task status keys: {list(task_status.keys())}")
    print(f"Task details count: {len(task_details)}")
    
    # Check for custom statuses in task_status
    custom_status_found = False
    for task_type, status in task_status.items():
        if status not in ['not_started', 'in_progress', 'submitted', 'approved', 'retake']:
            custom_status_found = True
            print(f"✅ Found custom status in task_status: {task_type} = {status}")
    
    # Check for custom statuses in task_details
    custom_status_in_details = False
    for task_detail in task_details:
        status = task_detail.get('status')
        if status not in ['not_started', 'in_progress', 'submitted', 'approved', 'retake']:
            custom_status_in_details = True
            print(f"✅ Found custom status in task_details: {task_detail['task_type']} = {status}")
    
    # Verify all required fields are present in task_details
    required_fields = ['task_type', 'status', 'task_id', 'assigned_user_id']
    all_fields_present = True
    for task_detail in task_details:
        missing_fields = [f for f in required_fields if f not in task_detail]
        if missing_fields:
            print(f"❌ Missing fields in task details: {missing_fields}")
            all_fields_present = False
    
    if all_fields_present:
        print("✅ All required fields present in task_details")
    
    return custom_status_found or custom_status_in_details

def test_custom_task_types_support(token, project_id):
    """Test that custom task types are also supported in optimized queries"""
    headers = {'Authorization': f'Bearer {token}'}
    
    print("\n=== Testing Custom Task Types Support ===")
    
    # Add custom task types to the project
    project_update = {
        'custom_shot_task_types': ['previz', 'matchmove'],
        'custom_asset_task_types': ['texturing', 'grooming']
    }
    
    response = requests.put(f'{BASE_URL}/projects/{project_id}', json=project_update, headers=headers)
    if response.status_code == 200:
        print("✅ Added custom task types to project")
    else:
        print(f"❌ Failed to add custom task types: {response.status_code}")
        return False
    
    # Test shots with custom task 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', {})
            custom_types_found = [t for t in task_status.keys() if t in ['previz', 'matchmove']]
            if custom_types_found:
                print(f"✅ Found custom task types in shot task_status: {custom_types_found}")
            else:
                print("ℹ️ Custom task types initialized as 'not_started' (expected behavior)")
    
    # Test assets with custom task types
    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', {})
            custom_types_found = [t for t in task_status.keys() if t in ['texturing', 'grooming']]
            if custom_types_found:
                print(f"✅ Found custom task types in asset task_status: {custom_types_found}")
            else:
                print("ℹ️ Custom task types initialized as 'not_started' (expected behavior)")
    
    return True

def main():
    print("=" * 80)
    print("Testing Custom Task Status Support in Optimized Queries")
    print("=" * 80)
    
    # Login
    token = login()
    if not token:
        return
    
    # Create test project
    project_id = create_test_project(token)
    if not project_id:
        return
    
    # Create custom task statuses
    custom_statuses = create_custom_task_statuses(token, project_id)
    if not custom_statuses:
        print("❌ Failed to create custom statuses")
        return
    
    # Create test episode
    episode_id = create_test_episode(token, project_id)
    if not episode_id:
        return
    
    # Create test shot
    shot_id = create_test_shot(token, episode_id)
    if not shot_id:
        return
    
    # Create test asset
    asset_id = create_test_asset(token, project_id)
    if not asset_id:
        return
    
    # Create tasks with custom statuses
    custom_status_id = custom_statuses[0]['id']  # Use first custom status
    shot_task_id = create_task_with_custom_status(token, shot_id=shot_id, task_type='layout', custom_status_id=custom_status_id)
    asset_task_id = create_task_with_custom_status(token, asset_id=asset_id, task_type='modeling', custom_status_id=custom_status_id)
    
    # If we couldn't create asset task, try a different task type
    if not asset_task_id:
        print("Trying different task type for asset...")
        asset_task_id = create_task_with_custom_status(token, asset_id=asset_id, task_type='surfacing', custom_status_id=custom_status_id)
    
    # Test optimized queries
    shot_test_passed = test_shot_optimization_with_custom_statuses(token, project_id, episode_id)
    asset_test_passed = test_asset_optimization_with_custom_statuses(token, project_id)
    custom_types_test_passed = test_custom_task_types_support(token, project_id)
    
    # Summary
    print("\n" + "=" * 80)
    print("TEST SUMMARY")
    print("=" * 80)
    
    if shot_test_passed:
        print("✅ Shot optimization with custom statuses: PASSED")
    else:
        print("❌ Shot optimization with custom statuses: FAILED")
    
    if asset_test_passed:
        print("✅ Asset optimization with custom statuses: PASSED")
    else:
        print("❌ Asset optimization with custom statuses: FAILED")
    
    if custom_types_test_passed:
        print("✅ Custom task types support: PASSED")
    else:
        print("❌ Custom task types support: FAILED")
    
    overall_success = shot_test_passed and asset_test_passed and custom_types_test_passed
    
    if overall_success:
        print("\n🎉 ALL TESTS PASSED - Custom status support is working correctly!")
    else:
        print("\n❌ SOME TESTS FAILED - Custom status support needs attention")
    
    return overall_success

if __name__ == "__main__":
    main()