"""
Test custom task type integration with asset and shot creation
"""
import sys
import os
sys.path.insert(0, os.path.dirname(__file__))

from database import SessionLocal, engine
from models.project import Project
from models.asset import Asset, AssetCategory
from models.episode import Episode
from models.shot import Shot
from models.task import Task
from models.user import User, UserRole
from sqlalchemy.orm.attributes import flag_modified
import json

def test_custom_task_types_integration():
    """Test that custom task types work with asset and shot creation"""
    db = SessionLocal()
    
    try:
        print("\n=== Testing Custom Task Types Integration ===\n")
        
        # 1. Get or create a test project with custom task types
        print("1. Getting/creating test project with custom task types...")
        project = db.query(Project).first()
        
        if not project:
            print("   No existing project found, creating new one...")
            project = Project(
                name="Custom Task Test Project",
                status="planning"
            )
            db.add(project)
            db.commit()
            db.refresh(project)
        
        # Add custom task types to the project
        project.custom_asset_task_types = ["grooming", "lookdev"]
        project.custom_shot_task_types = ["previz", "matchmove"]
        flag_modified(project, 'custom_asset_task_types')
        flag_modified(project, 'custom_shot_task_types')
        db.commit()
        db.refresh(project)
        
        print(f"   ✓ Using project: {project.name} (ID: {project.id})")
        print(f"   ✓ Custom asset types: {project.custom_asset_task_types}")
        print(f"   ✓ Custom shot types: {project.custom_shot_task_types}")
        
        # 2. Create an asset with custom task type
        print("\n2. Creating asset with custom task type...")
        asset = Asset(
            project_id=project.id,
            name="Test Character",
            category=AssetCategory.CHARACTERS,
            description="Test character with custom grooming task"
        )
        db.add(asset)
        db.commit()
        db.refresh(asset)
        print(f"   ✓ Created asset: {asset.name} (ID: {asset.id})")
        
        # Create standard and custom tasks for the asset
        standard_tasks = ["modeling", "surfacing", "rigging"]
        custom_tasks = ["grooming", "lookdev"]
        
        for task_type in standard_tasks + custom_tasks:
            task = Task(
                project_id=project.id,
                asset_id=asset.id,
                task_type=task_type,
                name=f"{asset.name} - {task_type.title()}",
                description=f"{task_type.title()} task for {asset.name}"
            )
            db.add(task)
        
        db.commit()
        print(f"   ✓ Created {len(standard_tasks)} standard tasks")
        print(f"   ✓ Created {len(custom_tasks)} custom tasks")
        
        # 3. Verify tasks were created
        print("\n3. Verifying asset tasks...")
        asset_tasks = db.query(Task).filter(Task.asset_id == asset.id).all()
        print(f"   ✓ Total tasks: {len(asset_tasks)}")
        for task in asset_tasks:
            task_label = "CUSTOM" if task.task_type in custom_tasks else "STANDARD"
            print(f"   - {task.name} ({task_label})")
        
        # 4. Create an episode and shot with custom task type
        print("\n4. Creating episode and shot with custom task type...")
        episode = Episode(
            project_id=project.id,
            name="Episode 01",
            episode_number=1
        )
        db.add(episode)
        db.commit()
        db.refresh(episode)
        print(f"   ✓ Created episode: {episode.name} (ID: {episode.id})")
        
        shot = Shot(
            episode_id=episode.id,
            name="SH010",
            frame_start=1001,
            frame_end=1100
        )
        db.add(shot)
        db.commit()
        db.refresh(shot)
        print(f"   ✓ Created shot: {shot.name} (ID: {shot.id})")
        
        # Create standard and custom tasks for the shot
        standard_shot_tasks = ["layout", "animation", "lighting", "compositing"]
        custom_shot_tasks = ["previz", "matchmove"]
        
        for task_type in standard_shot_tasks + custom_shot_tasks:
            task = Task(
                project_id=project.id,
                episode_id=episode.id,
                shot_id=shot.id,
                task_type=task_type,
                name=f"{shot.name}_{task_type}",
                description=f"{task_type.title()} task for shot {shot.name}"
            )
            db.add(task)
        
        db.commit()
        print(f"   ✓ Created {len(standard_shot_tasks)} standard tasks")
        print(f"   ✓ Created {len(custom_shot_tasks)} custom tasks")
        
        # 5. Verify shot tasks were created
        print("\n5. Verifying shot tasks...")
        shot_tasks = db.query(Task).filter(Task.shot_id == shot.id).all()
        print(f"   ✓ Total tasks: {len(shot_tasks)}")
        for task in shot_tasks:
            task_label = "CUSTOM" if task.task_type in custom_shot_tasks else "STANDARD"
            print(f"   - {task.name} ({task_label})")
        
        # 6. Test querying all task types for the project
        print("\n6. Testing task type queries...")
        from routers.assets import get_all_asset_task_types
        from routers.shots import get_all_shot_task_types
        
        all_asset_types = get_all_asset_task_types(project.id, db)
        all_shot_types = get_all_shot_task_types(project.id, db)
        
        print(f"   ✓ All asset task types: {all_asset_types}")
        print(f"   ✓ All shot task types: {all_shot_types}")
        
        # Verify custom types are included
        assert "grooming" in all_asset_types, "Custom asset type 'grooming' not found"
        assert "lookdev" in all_asset_types, "Custom asset type 'lookdev' not found"
        assert "previz" in all_shot_types, "Custom shot type 'previz' not found"
        assert "matchmove" in all_shot_types, "Custom shot type 'matchmove' not found"
        print("   ✓ All custom types are included in queries")
        
        print("\n=== All Tests Passed! ===\n")
        
        # Cleanup
        print("Cleaning up test data...")
        db.query(Task).filter(Task.project_id == project.id).delete()
        db.query(Shot).filter(Shot.episode_id == episode.id).delete()
        db.query(Episode).filter(Episode.project_id == project.id).delete()
        db.query(Asset).filter(Asset.project_id == project.id).delete()
        db.query(Project).filter(Project.id == project.id).delete()
        db.commit()
        print("✓ Cleanup complete")
        
    except Exception as e:
        print(f"\n❌ Error: {str(e)}")
        import traceback
        traceback.print_exc()
        db.rollback()
    finally:
        db.close()

if __name__ == "__main__":
    test_custom_task_types_integration()