#!/usr/bin/env python3
"""
Test script to verify custom status support across multiple projects with different configurations.

This test ensures that:
1. Projects with no custom statuses work correctly
2. Projects with different custom status sets work independently
3. Mixed project queries handle custom statuses correctly
4. Default status fallback works when custom statuses are missing
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import json
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from database import Base
from models.user import User, UserRole
from models.project import Project, ProjectType, ProjectStatus, ProjectMember
from models.task import Task
from models.asset import Asset, AssetCategory, AssetStatus
from models.episode import Episode, EpisodeStatus
from models.shot import Shot, ShotStatus

# Create test database
TEST_DB = "test_multi_project_custom.db"
if os.path.exists(TEST_DB):
    try:
        os.remove(TEST_DB)
    except PermissionError:
        pass  # File in use, continue anyway

engine = create_engine(f"sqlite:///{TEST_DB}")
Base.metadata.create_all(engine)
SessionLocal = sessionmaker(bind=engine)

def setup_multi_project_data():
    """Set up test data with multiple projects having different custom status configurations."""
    print("Setting up multi-project test data...")
    
    db = SessionLocal()
    try:
        # Create test user
        user = User(
            email="coordinator@test.com",
            password_hash="test",
            first_name="Test",
            last_name="Coordinator",
            role=UserRole.COORDINATOR,
            is_approved=True
        )
        db.add(user)
        db.commit()
        
        # Project 1: No custom statuses (uses only default statuses)
        project1 = Project(
            name="Default Status Project",
            code_name="DSP",
            client_name="Client 1",
            project_type=ProjectType.CINEMA,
            status=ProjectStatus.IN_PROGRESS,
            custom_task_statuses=json.dumps([])  # Empty custom statuses
        )
        db.add(project1)
        
        # Project 2: Custom statuses with different names
        project2_custom = [
            {
                "id": "review_needed",
                "name": "Review Needed",
                "color": "#FFA500",
                "order": 0,
                "is_default": False
            },
            {
                "id": "client_approved",
                "name": "Client Approved",
                "color": "#32CD32",
                "order": 1,
                "is_default": True
            }
        ]
        
        project2 = Project(
            name="Custom Status Project",
            code_name="CSP",
            client_name="Client 2",
            project_type=ProjectType.TV,
            status=ProjectStatus.IN_PROGRESS,
            custom_task_statuses=json.dumps(project2_custom)
        )
        db.add(project2)
        
        # Project 3: Different custom statuses
        project3_custom = [
            {
                "id": "waiting_assets",
                "name": "Waiting for Assets",
                "color": "#FFD700",
                "order": 0,
                "is_default": False
            },
            {
                "id": "final_review",
                "name": "Final Review",
                "color": "#9370DB",
                "order": 1,
                "is_default": False
            }
        ]
        
        project3 = Project(
            name="Different Custom Project",
            code_name="DCP",
            client_name="Client 3",
            project_type=ProjectType.CINEMA,
            status=ProjectStatus.IN_PROGRESS,
            custom_task_statuses=json.dumps(project3_custom)
        )
        db.add(project3)
        
        db.commit()
        
        # Add user as member of all projects
        for project in [project1, project2, project3]:
            member = ProjectMember(
                user_id=user.id,
                project_id=project.id
            )
            db.add(member)
        
        db.commit()
        
        # Create episodes for each project
        episodes = []
        for i, project in enumerate([project1, project2, project3], 1):
            episode = Episode(
                project_id=project.id,
                name=f"Episode {i}",
                episode_number=i,
                status=EpisodeStatus.IN_PROGRESS
            )
            db.add(episode)
            episodes.append(episode)
        
        db.commit()
        
        # Create shots with different status patterns for each project
        
        # Project 1: Only default statuses
        shot1 = Shot(
            project_id=project1.id,
            episode_id=episodes[0].id,
            name="default_shot_001",
            frame_start=1001,
            frame_end=1100,
            status=ShotStatus.IN_PROGRESS
        )
        db.add(shot1)
        db.flush()
        
        # Tasks with only default statuses
        for task_type, status in [("layout", "not_started"), ("animation", "in_progress"), ("lighting", "approved")]:
            task = Task(
                project_id=project1.id,
                episode_id=episodes[0].id,
                shot_id=shot1.id,
                name=f"{shot1.name}_{task_type}",
                task_type=task_type,
                status=status
            )
            db.add(task)
        
        # Project 2: Mix of default and custom statuses
        shot2 = Shot(
            project_id=project2.id,
            episode_id=episodes[1].id,
            name="custom_shot_001",
            frame_start=1001,
            frame_end=1100,
            status=ShotStatus.IN_PROGRESS
        )
        db.add(shot2)
        db.flush()
        
        # Tasks with mix of default and custom statuses
        for task_type, status in [("layout", "review_needed"), ("animation", "submitted"), ("lighting", "client_approved")]:
            task = Task(
                project_id=project2.id,
                episode_id=episodes[1].id,
                shot_id=shot2.id,
                name=f"{shot2.name}_{task_type}",
                task_type=task_type,
                status=status
            )
            db.add(task)
        
        # Project 3: Different custom statuses
        shot3 = Shot(
            project_id=project3.id,
            episode_id=episodes[2].id,
            name="different_shot_001",
            frame_start=1001,
            frame_end=1100,
            status=ShotStatus.IN_PROGRESS
        )
        db.add(shot3)
        db.flush()
        
        # Tasks with different custom statuses
        for task_type, status in [("layout", "waiting_assets"), ("animation", "final_review"), ("lighting", "retake")]:
            task = Task(
                project_id=project3.id,
                episode_id=episodes[2].id,
                shot_id=shot3.id,
                name=f"{shot3.name}_{task_type}",
                task_type=task_type,
                status=status
            )
            db.add(task)
        
        # Create assets for each project with different status patterns
        
        # Project 1 asset: Only default statuses
        asset1 = Asset(
            project_id=project1.id,
            name="default_character",
            category=AssetCategory.CHARACTERS,
            status=AssetStatus.IN_PROGRESS
        )
        db.add(asset1)
        db.flush()
        
        for task_type, status in [("modeling", "not_started"), ("surfacing", "in_progress")]:
            task = Task(
                project_id=project1.id,
                asset_id=asset1.id,
                name=f"{asset1.name}_{task_type}",
                task_type=task_type,
                status=status
            )
            db.add(task)
        
        # Project 2 asset: Custom statuses
        asset2 = Asset(
            project_id=project2.id,
            name="custom_prop",
            category=AssetCategory.PROPS,
            status=AssetStatus.IN_PROGRESS
        )
        db.add(asset2)
        db.flush()
        
        for task_type, status in [("modeling", "review_needed"), ("surfacing", "client_approved")]:
            task = Task(
                project_id=project2.id,
                asset_id=asset2.id,
                name=f"{asset2.name}_{task_type}",
                task_type=task_type,
                status=status
            )
            db.add(task)
        
        # Project 3 asset: Different custom statuses
        asset3 = Asset(
            project_id=project3.id,
            name="different_set",
            category=AssetCategory.SETS,
            status=AssetStatus.IN_PROGRESS
        )
        db.add(asset3)
        db.flush()
        
        for task_type, status in [("modeling", "waiting_assets"), ("surfacing", "final_review")]:
            task = Task(
                project_id=project3.id,
                asset_id=asset3.id,
                name=f"{asset3.name}_{task_type}",
                task_type=task_type,
                status=status
            )
            db.add(task)
        
        db.commit()
        print("✅ Multi-project test data setup complete")
        return user, [project1, project2, project3], episodes
        
    except Exception as e:
        db.rollback()
        print(f"❌ Failed to setup multi-project test data: {e}")
        raise
    finally:
        db.close()


def test_project_isolation():
    """Test that custom statuses are properly isolated between projects."""
    print("\n=== Test 1: Project Custom Status Isolation ===")
    
    db = SessionLocal()
    try:
        user = db.query(User).first()
        projects = db.query(Project).all()
        
        from routers.shots import list_shots
        import asyncio
        
        # Test each project individually
        for project in projects:
            episode = db.query(Episode).filter(Episode.project_id == project.id).first()
            
            async def get_project_shots():
                return await list_shots(
                    episode_id=episode.id,
                    project_id=project.id,
                    task_status_filter=None,
                    sort_by=None,
                    sort_direction="asc",
                    skip=0,
                    limit=100,
                    db=db,
                    current_user=user
                )
            
            shots = asyncio.run(get_project_shots())
            
            print(f"\nProject: {project.name} (ID: {project.id})")
            print(f"  Custom statuses: {json.loads(project.custom_task_statuses or '[]')}")
            
            for shot in shots:
                print(f"  Shot: {shot.name}")
                print(f"    Task statuses: {shot.task_status}")
                
                # Verify that only appropriate statuses are present
                for task_type, status in shot.task_status.items():
                    if status.startswith('custom_') or status in ['review_needed', 'client_approved', 'waiting_assets', 'final_review']:
                        # This is a custom status, verify it belongs to this project
                        project_custom_statuses = json.loads(project.custom_task_statuses or '[]')
                        custom_status_ids = [cs['id'] for cs in project_custom_statuses]
                        
                        if status not in custom_status_ids and status not in ['not_started', 'in_progress', 'submitted', 'approved', 'retake']:
                            # This might be a status from another project - this would be an error
                            print(f"    ⚠️  Status '{status}' not found in project custom statuses: {custom_status_ids}")
                        else:
                            print(f"    ✅ Status '{status}' correctly belongs to this project")
        
        print("✅ Test 1 PASSED: Project custom status isolation works correctly")
        return True
        
    except Exception as e:
        print(f"❌ Test 1 FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        db.close()


def test_mixed_project_queries():
    """Test queries that span multiple projects with different custom statuses."""
    print("\n=== Test 2: Mixed Project Queries ===")
    
    db = SessionLocal()
    try:
        user = db.query(User).first()
        
        from routers.shots import list_shots
        from routers.assets import list_assets
        import asyncio
        
        # Test shots across all projects (no project filter)
        async def get_all_shots():
            return await list_shots(
                episode_id=None,
                project_id=None,  # No project filter - should get all accessible shots
                task_status_filter=None,
                sort_by=None,
                sort_direction="asc",
                skip=0,
                limit=100,
                db=db,
                current_user=user
            )
        
        all_shots = asyncio.run(get_all_shots())
        
        print(f"✅ Retrieved {len(all_shots)} shots across all projects")
        
        # Verify that each shot has appropriate statuses for its project
        for shot in all_shots:
            print(f"\nShot: {shot.name}")
            print(f"  Task statuses: {shot.task_status}")
            
            # Get the project for this shot to verify statuses
            shot_project = db.query(Project).filter(Project.id == shot.project_id).first()
            project_custom_statuses = json.loads(shot_project.custom_task_statuses or '[]')
            custom_status_ids = [cs['id'] for cs in project_custom_statuses]
            
            for task_type, status in shot.task_status.items():
                if status not in ['not_started', 'in_progress', 'submitted', 'approved', 'retake']:
                    # This is a custom status, verify it belongs to the shot's project
                    assert status in custom_status_ids, f"Custom status '{status}' not found in project {shot_project.name} custom statuses"
                    print(f"    ✅ Custom status '{status}' correctly belongs to project {shot_project.name}")
        
        # Test assets across all projects
        async def get_all_assets():
            return await list_assets(
                project_id=None,  # No project filter
                category=None,
                task_status_filter=None,
                sort_by=None,
                sort_direction="asc",
                skip=0,
                limit=100,
                db=db,
                current_user=user
            )
        
        all_assets = asyncio.run(get_all_assets())
        
        print(f"\n✅ Retrieved {len(all_assets)} assets across all projects")
        
        # Verify that each asset has appropriate statuses for its project
        for asset in all_assets:
            print(f"\nAsset: {asset.name}")
            print(f"  Task statuses: {asset.task_status}")
            
            # Get the project for this asset to verify statuses
            asset_project = db.query(Project).filter(Project.id == asset.project_id).first()
            project_custom_statuses = json.loads(asset_project.custom_task_statuses or '[]')
            custom_status_ids = [cs['id'] for cs in project_custom_statuses]
            
            for task_type, status in asset.task_status.items():
                if status not in ['not_started', 'in_progress', 'submitted', 'approved', 'retake']:
                    # This is a custom status, verify it belongs to the asset's project
                    assert status in custom_status_ids, f"Custom status '{status}' not found in project {asset_project.name} custom statuses"
                    print(f"    ✅ Custom status '{status}' correctly belongs to project {asset_project.name}")
        
        print("✅ Test 2 PASSED: Mixed project queries handle custom statuses correctly")
        return True
        
    except Exception as e:
        print(f"❌ Test 2 FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        db.close()


def test_custom_status_sorting_across_projects():
    """Test that custom status sorting works correctly across different projects."""
    print("\n=== Test 3: Custom Status Sorting Across Projects ===")
    
    db = SessionLocal()
    try:
        user = db.query(User).first()
        
        from routers.shots import list_shots
        import asyncio
        
        # Test sorting by task status across all projects
        async def get_sorted_shots():
            return await list_shots(
                episode_id=None,
                project_id=None,
                task_status_filter=None,
                sort_by="layout_status",  # Sort by layout task status
                sort_direction="asc",
                skip=0,
                limit=100,
                db=db,
                current_user=user
            )
        
        sorted_shots = asyncio.run(get_sorted_shots())
        
        print(f"✅ Retrieved {len(sorted_shots)} shots sorted by layout status")
        
        # Verify sorting works with mixed custom statuses from different projects
        for shot in sorted_shots:
            layout_status = shot.task_status.get('layout', 'not_started')
            print(f"  Shot {shot.name}: layout = {layout_status}")
            
            # Get the project to access custom statuses for sorting verification
            shot_project = db.query(Project).filter(Project.id == shot.project_id).first()
            project_custom_statuses = json.loads(shot_project.custom_task_statuses or '[]')
            
            # Import the sorting function to verify order
            from routers.shots import get_status_sort_order
            sort_order = get_status_sort_order(layout_status, project_custom_statuses)
            print(f"    Project: {shot_project.name}, Sort order: {sort_order}")
        
        print("✅ Test 3 PASSED: Custom status sorting across projects works correctly")
        return True
        
    except Exception as e:
        print(f"❌ Test 3 FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        db.close()


def main():
    """Run all multi-project custom status tests."""
    print("🚀 Starting Multi-Project Custom Status Tests")
    print("=" * 60)
    
    try:
        # Setup test data
        setup_multi_project_data()
        
        # Run tests
        tests = [
            test_project_isolation,
            test_mixed_project_queries,
            test_custom_status_sorting_across_projects
        ]
        
        passed = 0
        failed = 0
        
        for test in tests:
            try:
                if test():
                    passed += 1
                else:
                    failed += 1
            except Exception as e:
                print(f"❌ Test failed with exception: {e}")
                failed += 1
        
        print("\n" + "=" * 60)
        print(f"📊 Test Results: {passed} passed, {failed} failed")
        
        if failed == 0:
            print("🎉 All multi-project tests passed! Custom status support works across projects.")
            return True
        else:
            print("❌ Some tests failed. Multi-project custom status support needs fixes.")
            return False
            
    except Exception as e:
        print(f"❌ Test setup failed: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        # Cleanup
        try:
            if os.path.exists(TEST_DB):
                os.remove(TEST_DB)
        except PermissionError:
            pass  # File in use, ignore


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