#!/usr/bin/env python3
"""
Test script for worker components.

Run: python -m worker.test_worker
"""

import os
import sys
import asyncio
import logging

# Add parent to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(name)s | %(levelname)s | %(message)s')
logger = logging.getLogger("WorkerTest")


def test_gpu_config():
    """Test GPU auto-configuration."""
    logger.info("=" * 50)
    logger.info("Testing GPU Configuration")
    logger.info("=" * 50)
    
    from worker.gpu_config import auto_configure_gpu, detect_gpu
    
    # Detect GPU
    gpu_name, vram_gb = detect_gpu()
    logger.info(f"GPU: {gpu_name}")
    logger.info(f"VRAM: {vram_gb:.1f} GB")
    
    # Auto configure
    config = auto_configure_gpu()
    logger.info(f"Embedding batch: {config.embedding_batch_size}")
    logger.info(f"Music batch: {config.music_batch_size}")
    logger.info(f"VAD workers: {config.vad_workers}")
    logger.info(f"Chunk workers: {config.chunk_workers}")
    
    return config


def test_r2_connection():
    """Test R2 storage connection."""
    logger.info("=" * 50)
    logger.info("Testing R2 Connection")
    logger.info("=" * 50)
    
    from worker.storage import R2Client, R2Config
    
    # Use hardcoded test credentials
    config = R2Config(
        access_key_id="c3c9190ae7ff98b10271ea8db6940210",
        secret_access_key="eab9394d02b48a865634105b92c74751ec9a311c56884f7aead5d76476c6b576",
        account_id="cb908ed13329eb7b186e06ab51bda190",
        bucket="diarization-output",
    )
    
    client = R2Client(config)
    
    # Test connection
    try:
        response = client.client.list_buckets()
        logger.info(f"✅ R2 Connection successful!")
        logger.info(f"   Buckets: {[b['Name'] for b in response['Buckets']]}")
        
        # Ensure our bucket exists
        client.ensure_bucket()
        logger.info(f"   Target bucket '{config.bucket}' ready")
        
        return True
    except Exception as e:
        logger.error(f"❌ R2 Connection failed: {e}")
        return False


async def test_local_coordinator():
    """Test local coordinator with CSV."""
    logger.info("=" * 50)
    logger.info("Testing Local Coordinator")
    logger.info("=" * 50)
    
    from worker.coordinator import LocalCoordinator
    
    # Check if CSV exists
    csv_path = "data/english_podcasts.csv"
    if not os.path.exists(csv_path):
        csv_path = "/home/ubuntu/maya3data/data/english_podcasts.csv"
    
    if not os.path.exists(csv_path):
        logger.warning(f"CSV not found at {csv_path}")
        return False
    
    coordinator = LocalCoordinator(csv_path, progress_file="/tmp/test_progress.json")
    
    logger.info(f"Stats: {coordinator.stats}")
    
    # Get a few video IDs
    videos = []
    for _ in range(5):
        video = coordinator.get_next()
        if video:
            videos.append(video)
    
    logger.info(f"Sample videos: {videos}")
    
    return True


def test_pipeline_import():
    """Test that pipeline can be imported."""
    logger.info("=" * 50)
    logger.info("Testing Pipeline Import")
    logger.info("=" * 50)
    
    try:
        from src.config import Config
        from pipeline import process_single_video
        
        config = Config()
        logger.info(f"✅ Pipeline imported successfully")
        logger.info(f"   Output dir: {config.output_dir}")
        logger.info(f"   Embedding batch: {config.embedding_batch_size}")
        logger.info(f"   VAD workers: {config.vad_workers}")
        
        return True
    except ImportError as e:
        logger.error(f"❌ Pipeline import failed: {e}")
        return False


async def run_all_tests():
    """Run all tests."""
    logger.info("\n" + "=" * 70)
    logger.info("🧪 WORKER COMPONENT TESTS")
    logger.info("=" * 70 + "\n")
    
    results = {}
    
    # Test GPU config
    try:
        test_gpu_config()
        results['gpu_config'] = True
    except Exception as e:
        logger.error(f"GPU config test failed: {e}")
        results['gpu_config'] = False
    
    # Test R2
    try:
        results['r2'] = test_r2_connection()
    except Exception as e:
        logger.error(f"R2 test failed: {e}")
        results['r2'] = False
    
    # Test local coordinator
    try:
        results['coordinator'] = await test_local_coordinator()
    except Exception as e:
        logger.error(f"Coordinator test failed: {e}")
        results['coordinator'] = False
    
    # Test pipeline import
    try:
        results['pipeline'] = test_pipeline_import()
    except Exception as e:
        logger.error(f"Pipeline test failed: {e}")
        results['pipeline'] = False
    
    # Summary
    logger.info("\n" + "=" * 70)
    logger.info("📊 TEST RESULTS")
    logger.info("=" * 70)
    
    for test, passed in results.items():
        status = "✅ PASS" if passed else "❌ FAIL"
        logger.info(f"   {test:20s}: {status}")
    
    all_passed = all(results.values())
    logger.info("=" * 70)
    logger.info(f"Overall: {'✅ ALL TESTS PASSED' if all_passed else '❌ SOME TESTS FAILED'}")
    logger.info("=" * 70)
    
    return all_passed


if __name__ == "__main__":
    asyncio.run(run_all_tests())