import multiprocessing
import time
import argparse

def cpu_load(duration=None, intensity=1.0):
    """
    Generate CPU load on a single core.
    
    Args:
        duration (float, optional): Duration in seconds to run the load.
                                   If None, runs until interrupted.
        intensity (float): Value between 0.0 and 1.0 to control load intensity.
                          1.0 means 100% load, 0.5 means approximately 50% load.
    """
    start_time = time.time()
    
    while True:
        # Busy work
        x = 0
        # Adjust end range to control intensity
        end = int(10000000 * intensity)
        for i in range(end):
            x += i
            x *= 1.0000001
        
        # Sleep to reduce CPU usage if intensity < 1.0
        if intensity < 1.0:
            time.sleep((1.0 - intensity) * 0.1)
            
        # Check if duration has been reached
        if duration is not None and time.time() - start_time >= duration:
            break

def run_load_test(cores=None, duration=None, intensity=1.0):
    """
    Run CPU load test on multiple cores.
    
    Args:
        cores (int, optional): Number of cores to utilize. 
                              If None, uses all available cores.
        duration (float, optional): Duration in seconds to run the test.
                                   If None, runs until interrupted.
        intensity (float): Value between 0.0 and 1.0 to control load intensity.
    """
    if cores is None:
        cores = multiprocessing.cpu_count()
    
    print(f"Starting CPU load test on {cores} cores with {intensity*100}% intensity")
    if duration:
        print(f"Test will run for {duration} seconds")
    else:
        print("Test will run until interrupted with Ctrl+C")
    
    # Create and start processes
    processes = []
    try:
        for _ in range(cores):
            p = multiprocessing.Process(target=cpu_load, args=(duration, intensity))
            processes.append(p)
            p.start()
        
        # Wait for processes to complete if duration is specified
        if duration:
            for p in processes:
                p.join()
            print("CPU load test completed")
        else:
            # Keep the main process running until interrupted
            while True:
                time.sleep(1)
                
    except KeyboardInterrupt:
        print("Stopping CPU load test...")
        for p in processes:
            if p.is_alive():
                p.terminate()
        print("CPU load test stopped")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate CPU load for testing")
    parser.add_argument("-c", "--cores", type=int, default=None, 
                        help="Number of cores to utilize (default: all available)")
    parser.add_argument("-d", "--duration", type=float, default=None,
                        help="Duration in seconds (default: run until interrupted)")
    parser.add_argument("-i", "--intensity", type=float, default=1.0,
                        help="Load intensity from 0.0 to 1.0 (default: 1.0)")
    
    args = parser.parse_args()
    
    run_load_test(args.cores, args.duration, args.intensity)