"""
Heavy tests for data processing tasks using distributed compute
"""
import pytest
import time
import json
import hashlib
import threading
from distributed_compute import Coordinator, Worker
from distributed_compute.task import Task


def process_large_dataset(data_chunk):
    """Process a chunk of data with heavy computation"""
    result = []
    for item in data_chunk:
        # Simulate heavy data transformation
        processed = {
            'id': item['id'],
            'value': sum([item['value'] ** 2 for _ in range(100)]) / 100,
            'hash': hashlib.sha256(str(item).encode()).hexdigest(),
            'timestamp': time.time()
        }
        result.append(processed)
    return result


def aggregate_data(data_list):
    """Aggregate processed data"""
    total = sum(item['value'] for sublist in data_list for item in sublist)
    count = sum(len(sublist) for sublist in data_list)
    return {'total': total, 'count': count, 'average': total / count if count > 0 else 0}


@pytest.mark.timeout(60)
def test_large_dataset_processing():
    """Test processing large dataset across multiple workers"""
    # Create large dataset
    dataset = [{'id': i, 'value': i * 2 + 1} for i in range(1000)]
    chunk_size = 100
    chunks = [dataset[i:i+chunk_size] for i in range(0, len(dataset), chunk_size)]
    
    # Start coordinator
    coordinator = Coordinator(port=5557, verbose=False)
    coordinator.start_server()
    time.sleep(0.5)
    
    # Start workers in threads
    workers = []
    for i in range(3):
        worker = Worker('localhost', 5557, name=f'worker_{i}')
        workers.append(worker)
        thread = threading.Thread(target=worker.start, daemon=True)
        thread.start()
    
    time.sleep(2)  # Wait for workers to connect
    
    try:
        # Distribute computation
        results = coordinator.map(process_large_dataset, chunks, timeout=50)
        
        # Verify results
        assert len(results) == len(chunks)
        assert all(len(result) == chunk_size for result in results)
        
        # Aggregate results
        aggregated = aggregate_data(results)
        assert aggregated['count'] == len(dataset)
        assert aggregated['total'] > 0
    finally:
        coordinator.stop_server()
        for worker in workers:
            worker.stop()


def transform_json_data(json_string):
    """Heavy JSON transformation"""
    data = json.loads(json_string)
    # Simulate complex transformations
    transformed = {}
    for key, value in data.items():
        if isinstance(value, (int, float)):
            transformed[f"{key}_squared"] = value ** 2
            transformed[f"{key}_cubed"] = value ** 3
        else:
            transformed[key] = str(value).upper()
    return transformed


@pytest.mark.timeout(60)
def test_json_data_transformation():
    """Test JSON data transformation across workers"""
    # Create JSON data
    json_data = [
        json.dumps({'id': i, 'value': i * 10, 'name': f'item_{i}'})
        for i in range(500)
    ]
    
    coordinator = Coordinator(port=5558, verbose=False)
    coordinator.start_server()
    time.sleep(0.5)
    
    workers = []
    for i in range(3):
        worker = Worker('localhost', 5558, name=f'worker_{i}')
        workers.append(worker)
        thread = threading.Thread(target=worker.start, daemon=True)
        thread.start()
    
    time.sleep(2)
    
    try:
        results = coordinator.map(transform_json_data, json_data, timeout=50)
        
        assert len(results) == len(json_data)
        assert all('id_squared' in result for result in results)
        assert all('value_cubed' in result for result in results)
    finally:
        coordinator.stop_server()
        for worker in workers:
            worker.stop()


def filter_and_sort_data(data_range):
    """Filter and sort data with heavy computation"""
    start, end = data_range
    data = list(range(start, end))
    
    # Heavy filtering
    filtered = [x for x in data if sum(int(d) for d in str(x)) % 3 == 0]
    
    # Heavy sorting with custom key
    sorted_data = sorted(filtered, key=lambda x: sum(int(d) ** 2 for d in str(x)), reverse=True)
    
    return sorted_data[:10]  # Return top 10


@pytest.mark.timeout(60)
def test_filter_and_sort():
    """Test filtering and sorting large datasets"""
    # Create data ranges
    ranges = [(i * 1000, (i + 1) * 1000) for i in range(20)]
    
    coordinator = Coordinator(port=5559, verbose=False)
    coordinator.start_server()
    time.sleep(0.5)
    
    workers = []
    for i in range(4):
        worker = Worker('localhost', 5559, name=f'worker_{i}')
        workers.append(worker)
        thread = threading.Thread(target=worker.start, daemon=True)
        thread.start()
    
    time.sleep(2)
    
    try:
        results = coordinator.map(filter_and_sort_data, ranges, timeout=50)
        
        assert len(results) == len(ranges)
        assert all(isinstance(result, list) for result in results)
        assert all(len(result) <= 10 for result in results)
    finally:
        coordinator.stop_server()
        for worker in workers:
            worker.stop()
