#!/usr/bin/env python3
"""
Test script for v0.1.6 - Large payload handling
Tests chunked transmission and compression with heavy workloads
"""

import numpy as np
import time
from distributed_compute import Coordinator, Worker
import threading

def heavy_computation(x):
    """Task that returns a large result (simulating heavy workload)."""
    # Create a large result - 10MB array
    large_array = np.random.rand(1280, 1000)  # ~10MB
    
    # Do some computation
    result = np.sum(large_array) + x * x
    
    return result

def test_large_payload():
    """Test with tasks that return large payloads."""
    print("=" * 60)
    print("Testing v0.1.6: Large Payload Handling")
    print("=" * 60)
    
    # Start coordinator
    coordinator = Coordinator(host="0.0.0.0", port=5556, verbose=True)
    coordinator.start_server()
    print("✓ Coordinator started on port 5556")
    
    # Start a local worker
    worker = Worker(
        coordinator_host="localhost",
        coordinator_port=5556,
        max_concurrent_tasks=2,
        name="test-worker-1"
    )
    
    worker_thread = threading.Thread(target=worker.start, kwargs={"block": True}, daemon=True)
    worker_thread.start()
    print("✓ Worker started")
    
    # Wait for worker to connect
    time.sleep(2)
    
    # Check worker connected
    stats = coordinator.get_stats()
    if stats["workers"] == 0:
        print("✗ ERROR: No workers connected")
        coordinator.stop_server()
        return False
    
    print(f"✓ {stats['workers']} worker(s) connected")
    
    # Run test with large payloads
    print("\nRunning test with 10 tasks (each returning ~10MB)...")
    start_time = time.time()
    
    try:
        results = coordinator.map(
            heavy_computation,
            range(10),
            timeout=60.0
        )
        
        elapsed = time.time() - start_time
        print(f"✓ All tasks completed in {elapsed:.2f}s")
        print(f"✓ Received {len(results)} results")
        
        # Verify results are reasonable
        for i, result in enumerate(results):
            if result is None:
                print(f"✗ ERROR: Task {i} returned None")
                coordinator.stop_server()
                return False
        
        print("✓ All results valid")
        
        # Show final stats
        final_stats = coordinator.get_stats()
        print(f"\nFinal stats:")
        print(f"  - Tasks completed: {final_stats['tasks_completed']}")
        print(f"  - Workers: {final_stats['workers']}")
        
        coordinator.stop_server()
        worker.stop()
        
        print("\n" + "=" * 60)
        print("✓ TEST PASSED: Large payloads handled successfully!")
        print("=" * 60)
        return True
        
    except Exception as e:
        print(f"\n✗ TEST FAILED: {e}")
        coordinator.stop_server()
        worker.stop()
        return False

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