"""
Test script for Dynamic Task Type Support (Task 15)

This test verifies that:
1. New task types are automatically included in aggregated queries
2. Custom task types appear in optimized responses
3. Projects that add new task types after optimization work correctly

Requirements: 3.4
"""

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

BASE_URL = "http://localhost:8000"


def login() -> str:
    """Login and return access token"""
    response = requests.post(f"{BASE_URL}/auth/login", json={
        "email": "admin@vfx.com",
        "password": "admin123"
    })
    if response.status_code == 200:
        return response.json()["access_token"]
    else:
        raise Exception(f"Login failed: {response.status_code} - {response.text}")


def create_test_project(token: str) -> int:
    """Create a test project and return its ID"""
    headers = {"Authorization": f"Bearer {token}"}
    
    project_data = {
        "name": f"Dynamic Task Type Test Project {int(time.time())}",
        "code_name": f"dtt_test_{int(time.time())}",
        "client_name": "Test Client",
        "project_type": "tv",
        "description": "Test project for dynamic task type support"
    }
    
    response = requests.post(f"{BASE_URL}/projects/", headers=headers, json=project_data)
    if response.status_code == 201:
        return response.json()["id"]
    else:
        raise Exception(f"Failed to create project: {response.status_code} - {response.text}")


def create_test_episode(token: str, project_id: int) -> int:
    """Create a test episode and return its ID"""
    headers = {"Authorization": f"Bearer {token}"}
    
    episode_data = {
        "name": "Test Episode",
        "episode_number": 1,
        "description": "Test episode for dynamic task type support"
    }
    
    response = requests.post(f"{BASE_URL}/episodes/projects/{project_id}/episodes", headers=headers, json=episode_data)
    if response.status_code == 201:
        return response.json()["id"]
    else:
        raise Exception(f"Failed to create episode: {response.status_code} - {response.text}")


def create_test_shot(token: str, episode_id: int, shot_name: str) -> int:
    """Create a test shot and return its ID"""
    headers = {"Authorization": f"Bearer {token}"}
    
    shot_data = {
        "name": shot_name,
        "description": f"Test shot {shot_name}",
        "frame_start": 1001,
        "frame_end": 1100
    }
    
    response = requests.post(f"{BASE_URL}/shots/?episode_id={episode_id}", headers=headers, json=shot_data)
    if response.status_code == 201:
        return response.json()["id"]
    else:
        raise Exception(f"Failed to create shot: {response.status_code} - {response.text}")


def create_test_asset(token: str, project_id: int, asset_name: str) -> int:
    """Create a test asset and return its ID"""
    headers = {"Authorization": f"Bearer {token}"}
    
    asset_data = {
        "name": asset_name,
        "category": "characters",
        "description": f"Test asset {asset_name}"
    }
    
    response = requests.post(f"{BASE_URL}/assets/?project_id={project_id}", headers=headers, json=asset_data)
    if response.status_code == 201:
        return response.json()["id"]
    else:
        raise Exception(f"Failed to create asset: {response.status_code} - {response.text}")


def add_custom_task_type(token: str, project_id: int, task_type: str, category: str) -> bool:
    """Add a custom task type to the project"""
    headers = {"Authorization": f"Bearer {token}"}
    
    data = {
        "task_type": task_type,
        "category": category
    }
    
    response = requests.post(f"{BASE_URL}/projects/{project_id}/custom-task-types", headers=headers, json=data)
    return response.status_code in [201, 409]  # 201 = created, 409 = already exists


def get_shots_with_task_status(token: str, project_id: int) -> List[Dict[str, Any]]:
    """Get shots with their task status information"""
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.get(f"{BASE_URL}/shots/?project_id={project_id}", headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"Failed to get shots: {response.status_code} - {response.text}")


def get_assets_with_task_status(token: str, project_id: int) -> List[Dict[str, Any]]:
    """Get assets with their task status information"""
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.get(f"{BASE_URL}/assets/?project_id={project_id}", headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"Failed to get assets: {response.status_code} - {response.text}")


def create_task_for_shot(token: str, shot_id: int, task_type: str) -> bool:
    """Create a task for a shot with the specified task type"""
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.post(f"{BASE_URL}/shots/{shot_id}/tasks?task_type={task_type}", headers=headers)
    return response.status_code in [201, 200]  # 201 = created, 200 = already exists


def create_task_for_asset(token: str, asset_id: int, task_type: str) -> bool:
    """Create a task for an asset with the specified task type"""
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.post(f"{BASE_URL}/assets/{asset_id}/tasks?task_type={task_type}", headers=headers)
    return response.status_code in [201, 200]  # 201 = created, 200 = already exists


def test_dynamic_task_type_support():
    """Main test function for dynamic task type support"""
    print("=" * 80)
    print("TESTING DYNAMIC TASK TYPE SUPPORT")
    print("=" * 80)
    
    try:
        # Step 1: Login
        print("\n1. Logging in...")
        token = login()
        print("✓ Login successful")
        
        # Step 2: Create test project
        print("\n2. Creating test project...")
        project_id = create_test_project(token)
        print(f"✓ Created project with ID: {project_id}")
        
        # Step 3: Create test episode
        print("\n3. Creating test episode...")
        episode_id = create_test_episode(token, project_id)
        print(f"✓ Created episode with ID: {episode_id}")
        
        # Step 4: Create test shots and assets
        print("\n4. Creating test shots and assets...")
        shot_id = create_test_shot(token, episode_id, "TEST_001")
        asset_id = create_test_asset(token, project_id, "TestCharacter")
        print(f"✓ Created shot ID: {shot_id}, asset ID: {asset_id}")
        
        # Step 5: Get initial shots and assets (should have standard task types only)
        print("\n5. Getting initial shots and assets with standard task types...")
        initial_shots = get_shots_with_task_status(token, project_id)
        initial_assets = get_assets_with_task_status(token, project_id)
        
        print(f"✓ Found {len(initial_shots)} shots, {len(initial_assets)} assets")
        
        if initial_shots:
            shot_task_types = list(initial_shots[0].get('task_status', {}).keys())
            print(f"  Shot task types: {shot_task_types}")
        
        if initial_assets:
            asset_task_types = list(initial_assets[0].get('task_status', {}).keys())
            print(f"  Asset task types: {asset_task_types}")
        
        # Step 6: Add custom task types
        print("\n6. Adding custom task types...")
        custom_shot_type = "fx_simulation"
        custom_asset_type = "grooming"
        
        shot_type_added = add_custom_task_type(token, project_id, custom_shot_type, "shot")
        asset_type_added = add_custom_task_type(token, project_id, custom_asset_type, "asset")
        
        print(f"✓ Custom shot task type '{custom_shot_type}' added: {shot_type_added}")
        print(f"✓ Custom asset task type '{custom_asset_type}' added: {asset_type_added}")
        
        # Step 7: Get shots and assets after adding custom task types
        print("\n7. Getting shots and assets after adding custom task types...")
        updated_shots = get_shots_with_task_status(token, project_id)
        updated_assets = get_assets_with_task_status(token, project_id)
        
        # Verify custom task types are included in aggregated responses
        shot_success = False
        asset_success = False
        
        if updated_shots:
            updated_shot_task_types = list(updated_shots[0].get('task_status', {}).keys())
            print(f"  Updated shot task types: {updated_shot_task_types}")
            shot_success = custom_shot_type in updated_shot_task_types
            print(f"  Custom shot task type '{custom_shot_type}' in response: {shot_success}")
        
        if updated_assets:
            updated_asset_task_types = list(updated_assets[0].get('task_status', {}).keys())
            print(f"  Updated asset task types: {updated_asset_task_types}")
            asset_success = custom_asset_type in updated_asset_task_types
            print(f"  Custom asset task type '{custom_asset_type}' in response: {asset_success}")
        
        # Step 8: Create tasks with custom task types
        print("\n8. Creating tasks with custom task types...")
        shot_task_created = create_task_for_shot(token, shot_id, custom_shot_type)
        asset_task_created = create_task_for_asset(token, asset_id, custom_asset_type)
        
        print(f"✓ Shot task with custom type created: {shot_task_created}")
        print(f"✓ Asset task with custom type created: {asset_task_created}")
        
        # Step 9: Verify tasks appear in aggregated data
        print("\n9. Verifying tasks appear in aggregated data...")
        final_shots = get_shots_with_task_status(token, project_id)
        final_assets = get_assets_with_task_status(token, project_id)
        
        shot_task_in_details = False
        asset_task_in_details = False
        
        if final_shots:
            shot_task_details = final_shots[0].get('task_details', [])
            shot_task_in_details = any(task.get('task_type') == custom_shot_type for task in shot_task_details)
            print(f"  Custom shot task in task_details: {shot_task_in_details}")
            
            # Check task status is updated
            shot_task_status = final_shots[0].get('task_status', {}).get(custom_shot_type)
            print(f"  Custom shot task status: {shot_task_status}")
        
        if final_assets:
            asset_task_details = final_assets[0].get('task_details', [])
            asset_task_in_details = any(task.get('task_type') == custom_asset_type for task in asset_task_details)
            print(f"  Custom asset task in task_details: {asset_task_in_details}")
            
            # Check task status is updated
            asset_task_status = final_assets[0].get('task_status', {}).get(custom_asset_type)
            print(f"  Custom asset task status: {asset_task_status}")
        
        # Step 10: Test results summary
        print("\n" + "=" * 80)
        print("TEST RESULTS SUMMARY")
        print("=" * 80)
        
        tests_passed = 0
        total_tests = 4
        
        print(f"\n1. Custom shot task type included in aggregated query: {'✓ PASS' if shot_success else '✗ FAIL'}")
        if shot_success:
            tests_passed += 1
        
        print(f"2. Custom asset task type included in aggregated query: {'✓ PASS' if asset_success else '✗ FAIL'}")
        if asset_success:
            tests_passed += 1
        
        print(f"3. Custom shot task appears in task_details: {'✓ PASS' if shot_task_in_details else '✗ FAIL'}")
        if shot_task_in_details:
            tests_passed += 1
        
        print(f"4. Custom asset task appears in task_details: {'✓ PASS' if asset_task_in_details else '✗ FAIL'}")
        if asset_task_in_details:
            tests_passed += 1
        
        print(f"\nOverall Result: {tests_passed}/{total_tests} tests passed")
        
        if tests_passed == total_tests:
            print("🎉 ALL TESTS PASSED - Dynamic task type support is working correctly!")
            return True
        else:
            print("❌ SOME TESTS FAILED - Dynamic task type support needs attention")
            return False
        
    except Exception as e:
        print(f"\n❌ TEST ERROR: {str(e)}")
        return False


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