#!/usr/bin/env python3
"""Script to fix test files that use old set_broadcast_callback pattern."""

import re
import sys
from pathlib import Path


def fix_file(filepath: Path) -> bool:
    """Fix a single test file.
    
    Args:
        filepath: Path to the test file
        
    Returns:
        True if file was modified, False otherwise
    """
    content = filepath.read_text()
    original = content
    
    # Pattern 1: Replace set_broadcast_callback calls
    # Old: service.set_broadcast_callback(mock_broadcast)
    # New: progress_service.subscribe("progress_updated", mock_event_handler)
    
    # Pattern 2: Fix download_service fixture to return tuple
    if "async def download_service(" in content and "yield service" in content:
        content = re.sub(
            r'(async def download_service\([^)]+\):.*?)(yield service)',
            r'\1yield service, progress_service',
            content,
            flags=re.DOTALL
        )
    
    #Pattern 3: Unpack download_service in tests
    if "def test_" in content or "async def test_" in content:
        # Find tests that use download_service but don't unpack it
        content = re.sub(
            r'(async def test_[^\(]+\([^)]*download_service[^)]*\):.*?""".*?""")\s*broadcasts',
            r'\1\n        download_svc, progress_svc = download_service\n        broadcasts',
            content,
            flags=re.DOTALL,
            count=1  # Only first occurrence in each test
        )
    
    # Pattern 4: Replace set_broadcast_callback with subscribe
    content = re.sub(
        r'(\w+)\.set_broadcast_callback\((\w+)\)',
        r'progress_service.subscribe("progress_updated", \2)',
        content
    )
    
    # Pattern 5: Fix event handler signatures
    # Old: async def mock_broadcast(message_type: str, room: str, data: dict):
    # New: async def mock_event_handler(event):
    content = re.sub(
        r'async def (mock_broadcast\w*)\([^)]+\):(\s+"""[^"]*""")?(\s+)broadcasts\.append',
        r'async def mock_event_handler(event):\2\3broadcasts.append',
        content
    )
    
    # Pattern 6: Fix broadcast append calls
    # Old: broadcasts.append({"type": message_type, "data": data})
    # New: broadcasts.append({"type": event.event_type, "data": event.progress.to_dict()})
    content = re.sub(
        r'broadcasts\.append\(\{[^}]*"type":\s*message_type[^}]*\}\)',
        'broadcasts.append({"type": event.event_type, "data": event.progress.to_dict()})',
        content
    )
    
    # Pattern 7: Update download_service usage in tests to use unpacked version
    content = re.sub(
        r'await download_service\.add_to_queue\(',
        r'await download_svc.add_to_queue(',
        content
    )
    content = re.sub(
        r'await download_service\.start',
        r'await download_svc.start',
        content
    )
    content = re.sub(
        r'await download_service\.stop',
        r'await download_svc.stop',
        content
    )
    content = re.sub(
        r'await download_service\.get_queue_status\(',
        r'await download_svc.get_queue_status(',
        content
    )
    content = re.sub(
        r'await download_service\.remove_from_queue\(',
        r'await download_svc.remove_from_queue(',
        content
    )
    content = re.sub(
        r'await download_service\.clear_completed\(',
        r'await download_svc.clear_completed(',
        content
    )
    
    if content != original:
        filepath.write_text(content)
        print(f"✓ Fixed {filepath}")
        return True
    else:
        print(f"  Skipped {filepath} (no changes needed)")
        return False


def main():
    """Main function to fix all test files."""
    test_dir = Path(__file__).parent / "tests"
    
    # Find all test files that might need fixing
    test_files = list(test_dir.rglob("test_*.py"))
    
    print(f"Found {len(test_files)} test files")
    print("Fixing test files...")
    
    fixed_count = 0
    for test_file in test_files:
        if fix_file(test_file):
            fixed_count += 1
    
    print(f"\nFixed {fixed_count}/{len(test_files)} files")
    return 0 if fixed_count > 0 else 1


if __name__ == "__main__":
    sys.exit(main())