"""
Test massive scaling with SHA-256 calculations distributed across cores and threads
"""
from cpu_parallel_driver import CPUDriver
import numpy as np
import multiprocessing as mp
import hashlib
from concurrent.futures import ThreadPoolExecutor
import time

# Constants for realistic scaling
NUM_CORES = 8  # Start with a smaller core count for testing
THREADS_PER_CORE = 4  # Fewer threads per core for testing
VECTOR_SIZE = 256
CHUNK_SIZE = 100  # Smaller chunks for testing
BATCH_SIZE = 10000  # Total number of SHA-256 operations
BATCH_WINDOW = 1000  # Number of instructions to process at once

def sha256_vector(data, stats):
    """Calculate SHA-256 on vector data and check for blocks"""
    hasher = hashlib.sha256()
    hasher.update(data.tobytes())
    hash_bytes = hasher.digest()
    
    # Count this hash
    stats.increment_hashes(1)
    
    # Check if we found a block (hash starts with enough zeros)
    # Adjust difficulty by changing the number of zeros required
    if hash_bytes.startswith(b'\x00\x00'):  # Requires 2 zero bytes (moderate difficulty)
        stats.increment_blocks()
        
    return np.frombuffer(hash_bytes, dtype=np.float32)

def create_sha_program(size=BATCH_SIZE, stats=None):
    """Create program with SHA-256 calculations"""
    if stats is None:
        raise ValueError("Stats object is required")
        
    program = []
    print(f"\nCreating SHA-256 program with {size} instructions...")
    
    # Create parallel SHA operations in chunks
    for i in range(size):
        # Distribute across cores and threads
        core_id = i % NUM_CORES
        thread_id = (i // NUM_CORES) % THREADS_PER_CORE
        
        if i % 1000 == 0:  # Log progress every 1000 instructions
            print(f"Creating instruction {i}/{size}")
        
        instruction = {
            'type': 'SHA256',
            'core_id': core_id,
            'thread_id': thread_id,
            'vector_offset': i * VECTOR_SIZE,
            'dest': i % 1024,  # Circular buffer for results
            'size': VECTOR_SIZE
        }
        
        if stats:
            instruction['stats'] = stats
            
        program.append(instruction)
        
    return program

def print_stats(stats):
    """Print current mining statistics"""
    current_stats = stats.get_stats()
    
    print("\033[H\033[J")  # Clear screen
    print(f"Mining Statistics:")
    print(f"- Total Hashes: {current_stats['total_hashes']:,}")
    print(f"- Blocks Found: {current_stats['blocks_found']}")
    print(f"- Runtime: {current_stats['duration']:.1f} seconds")
    print(f"- Current Hashrate: {current_stats['hashrate']:,.2f} H/s")
    print(f"- Hashes per Block: {current_stats['total_hashes'] / (current_stats['blocks_found'] if current_stats['blocks_found'] > 0 else 1):,.1f}")
    
def calculate_performance_rating(hashrate):
    """Calculate a performance rating based on hashrate"""
    # Base ratings on typical hardware performance levels
    ratings = [
        (1e9, "Exceptional - Server Farm Level"),
        (500e6, "Excellent - Mining Rig Level"),
        (100e6, "Very Good - High-End GPU Level"),
        (50e6, "Good - Mid-Range GPU Level"),
        (10e6, "Fair - Low-End GPU Level"),
        (1e6, "Basic - CPU Mining Level"),
        (0, "Limited - Development Level")
    ]
    
    for threshold, rating in ratings:
        if hashrate >= threshold:
            return rating
    return ratings[-1][1]

def main():
    # Initialize mining_stats
    from mining_stats import MiningStats
    stats = MiningStats()
    start = time.time()
    
    print(f"Starting SHA-256 Mining Test (60s):")
    print(f"- Cores: {NUM_CORES}")
    print(f"- Threads per core: {THREADS_PER_CORE}")
    print(f"- Total threads: {NUM_CORES * THREADS_PER_CORE}")
    print(f"- Vector size: {VECTOR_SIZE}")
    print(f"- Target: Run for 60 seconds")

    # Create shared memory for vectors
    print("\nCreating shared memory...")
    shared_vectors = mp.RawArray('f', BATCH_SIZE * VECTOR_SIZE)
    vectors_np = np.frombuffer(shared_vectors, dtype=np.float32).reshape(BATCH_SIZE, VECTOR_SIZE)
    vectors_np[:] = np.random.rand(BATCH_SIZE, VECTOR_SIZE)

    # Initialize CPU driver
    print("Initializing CPU driver...")
    driver = CPUDriver(
        num_cores=NUM_CORES, 
        shared_memory=shared_vectors,
        stats=stats
    )
    print("Initialized CPU driver")
    
    # Set runtime limit
    END_TIME = start + 60  # Run for 60 seconds
    
    # Create SHA-256 test program
    print("\nGenerating test program...")
    program = create_sha_program(BATCH_SIZE, stats=stats)
    
    # Run parallel execution
    print("\nStarting parallel SHA-256 calculations...")
    print("Press Ctrl+C to stop mining\n")
    start = time.time()
    last_print = start
    
    # Execute operations in batches and monitor real-time performance
    # Execute operations in batches and monitor real-time performance
    try:
        all_futures = []  # Track all futures
        batch_size = 1000  # Process in smaller batches
        current_pos = 0
        current_batch = program[current_pos:current_pos + batch_size]
        
        # Get initial stats
        stats_data = stats.get_stats()
        print(f"\nMining Statistics:")
        print(f"Total Hashes: {stats_data['total_hashes']:,}")
        print(f"Blocks Found: {stats_data['blocks_found']}")
        print(f"Hashrate: {stats_data['hashrate']:.2f} H/s")
        
        while time.time() < END_TIME:
            try:
                # Submit new batch
                futures = driver.execute_batch(current_batch)
                all_futures.extend(futures)
                
                # Check completed futures and remove them
                completed_futures = []
                for future in all_futures:
                    try:
                        thread_futures = future.get(timeout=0.1)
                        if thread_futures:  # Check if thread_futures exists
                            for thread_future in thread_futures:
                                thread_future.result(timeout=0.1)
                        completed_futures.append(future)
                    except TimeoutError:
                        continue
                    except Exception as e:
                        print(f"Error in batch: {e}")
                        completed_futures.append(future)
                
                # Remove completed futures and move to next batch
                for future in completed_futures:
                    all_futures.remove(future)
                
                # Move to next batch if current was completed
                if completed_futures:
                    current_pos = (current_pos + batch_size) % len(program)
                    current_batch = program[current_pos:current_pos + batch_size]
                    if not current_batch:  # Wrap around
                        current_batch = program[:batch_size]
                        current_pos = 0
                
                # Print stats every second
                now = time.time()
                if now - last_print >= 1.0:
                    stats_data = stats.get_stats()
                    print("\033[H\033[J")  # Clear screen
                    print(f"Mining Statistics:")
                    print(f"Runtime: {stats_data['duration']:.1f}s / 60s")
                    print(f"Total Hashes: {stats_data['total_hashes']:,}")
                    print(f"Current Hashrate: {stats_data['hashrate']:,.0f} H/s")
                    print(f"Blocks Found: {stats_data['blocks_found']}")
                    if stats_data['blocks_found'] > 0:
                        print(f"Avg Hashes per Block: {stats_data['total_hashes']/stats_data['blocks_found']:,.0f}")
                    print(f"\nActive Futures: {len(all_futures)}")
                    print(f"Current Position: {current_pos}/{len(program)}")
                    print(f"\nTime Remaining: {max(0, END_TIME - now):.1f}s")
                    last_print = now
                    
            except Exception as e:
                print(f"Error in main loop: {e}")
                # Continue execution despite errors
                
    except KeyboardInterrupt:
        print("\nMining stopped by user")
    finally:
        # Final cleanup
        if all_futures:
            print("\nWaiting for remaining tasks to complete...")
            for future in all_futures:
                try:
                    thread_futures = future.get(timeout=1.0)
                    if thread_futures:
                        for thread_future in thread_futures:
                            try:
                                thread_future.result(timeout=1.0)
                            except Exception as e:
                                print(f"Error in thread cleanup: {e}")
                except Exception as e:
                    print(f"Error in process cleanup: {e}")
        
        # Calculate final statistics
        end = time.time()
        duration = end - start
        final_stats = stats.get_stats()
        
        # Calculate performance metrics
        hashes_per_sec = final_stats['hashrate']
        hashes_per_core = hashes_per_sec / NUM_CORES
        hashes_per_thread = hashes_per_sec / (NUM_CORES * THREADS_PER_CORE)
    
    print(f"\nFinal Mining Results:")
    print(f"- Total Hashes: {final_stats['total_hashes']:,}")
    print(f"- Blocks Found: {final_stats['blocks_found']}")
    if final_stats['blocks_found'] > 0:
        print(f"- Block Find Rate: 1 block per {final_stats['total_hashes']/final_stats['blocks_found']:,.1f} hashes")
    print(f"- Execution time: {duration:.2f} seconds")
    print(f"- Hashrate: {hashes_per_sec:,.2f} H/s")
    print(f"- Hashrate/core: {hashes_per_core:,.2f} H/s")
    print(f"- Hashrate/thread: {hashes_per_thread:,.2f} H/s")
    
    # Calculate and display mining efficiency
    peak_hashrate = 100e6  # Assume 100 MH/s per thread peak
    theory_total = peak_hashrate * NUM_CORES * THREADS_PER_CORE
    efficiency = (hashes_per_sec / theory_total) * 100
    
    print(f"\nFinal Mining Results:")
    print(f"- Runtime: {duration:.1f}s")
    print(f"- Total Hashes: {stats.total_hashes.value:,}")
    print(f"- Average Hashrate: {hashes_per_sec:,.0f} H/s")
    print(f"- Hashrate per Core: {hashes_per_core:,.0f} H/s")
    print(f"- Hashrate per Thread: {hashes_per_thread:,.0f} H/s")
    print(f"- Blocks Found: {stats.blocks_found.value}")
    if stats.blocks_found.value > 0:
        print(f"- Avg Hashes per Block: {stats.total_hashes.value/stats.blocks_found.value:,.0f}")
    
    # Calculate and show performance rating
    rating = calculate_performance_rating(hashes_per_sec)
    print(f"\nPerformance Rating: {rating}")
    
    # Show efficiency vs theoretical peak
    print(f"\nEfficiency Analysis:")
    print(f"- Peak Hashrate: {theory_total:,.0f} H/s")
    print(f"- Achieved: {hashes_per_sec:,.0f} H/s")
    print(f"- Efficiency: {efficiency:.1f}%")

if __name__ == '__main__':
    main()