#!/usr/bin/env fish

# Function to check if a CPU core is free based on utilization
function is_core_free
    set -l core_id $argv[1]
    set -l threshold $argv[2]
    
    # Get initial CPU stats
    set -l initial (grep "^cpu$core_id " /proc/stat)
    set -l initial_values (string split ' ' $initial)
    
    # Extract user, nice, system, idle times
    set -l initial_user $initial_values[2]
    set -l initial_nice $initial_values[3]
    set -l initial_system $initial_values[4]
    set -l initial_idle $initial_values[5]
    
    # Wait a moment for measurement
    sleep 0.2
    
    # Get updated CPU stats
    set -l current (grep "^cpu$core_id " /proc/stat)
    set -l current_values (string split ' ' $current)
    
    # Extract updated times
    set -l current_user $current_values[2]
    set -l current_nice $current_values[3]
    set -l current_system $current_values[4]
    set -l current_idle $current_values[5]
    
    # Calculate deltas
    set -l delta_user (math "$current_user - $initial_user")
    set -l delta_nice (math "$current_nice - $initial_nice")
    set -l delta_system (math "$current_system - $initial_system")
    set -l delta_idle (math "$current_idle - $initial_idle")
    
    # Calculate total time and busy time
    set -l total_time (math "$delta_user + $delta_nice + $delta_system + $delta_idle")
    set -l busy_time (math "$delta_user + $delta_nice + $delta_system")
    
    # Avoid division by zero
    if test $total_time -eq 0
        set total_time 1
    end
    
    # Calculate CPU usage percentage
    set -l usage_pct (math "100 * $busy_time / $total_time")
    
    # echo "Core $core_id usage: $usage_pct%" >&2
    
    # Check if usage is below threshold
    if test "$usage_pct" -lt "$threshold"
        return 0  # Core is free
    else
        return 1  # Core is busy
    end
end

# Function to extract output file path from a command
function extract_output_path
    set -l cmd $argv[1]
    
    # Look for the --min-errors-output parameter and extract the path after it
    if string match -q "*--min-errors-output*" "$cmd"
        # Split the command by spaces
        set -l parts (string split ' ' "$cmd")
        
        # Find the index of --min-errors-output
        set -l idx 1
        set -l output_idx 0
        
        for part in $parts
            if test "$part" = "--min-errors-output"
                set output_idx (math $idx + 1)
                break
            end
            set idx (math $idx + 1)
        end
        
        # If we found it and there's a next parameter, return it
        if test $output_idx -gt 0 -a $output_idx -le (count $parts)
            echo $parts[$output_idx]
            return 0
        end
    end
    
    # No output path found
    return 1
end

# Main function to process commands
function core_runner
    set -l command_file $argv[1]
    set -l threshold $argv[2]
    set -l time_limit $argv[3]
    set -l memory_limit $argv[4]
    
    # Set default threshold if not provided
    if test -z "$threshold"
        set threshold 30
    end
    
    # Set default time limit if not provided (0 means no limit)
    if test -z "$time_limit"
        set time_limit 0
    end
    
    # Set default memory limit if not provided (0 means no limit)
    if test -z "$memory_limit"
        set memory_limit 0
    end
    
    # Create output directory
    set -l output_dir "./task_outputs"
    mkdir -p $output_dir
    
    # Get number of CPU cores
    set -l num_cores (nproc)
    
    echo "Processing commands with the following settings:"
    echo "  CPU usage threshold: $threshold%"
    echo "  Number of cores: $num_cores"
    echo "  Time limit: "(test $time_limit -eq 0 && echo "None" || echo "$time_limit seconds")
    echo "  Memory limit: "(test $memory_limit -eq 0 && echo "None" || echo "$memory_limit MB")
    echo "  Command outputs will be saved to $output_dir/"
    echo "  Commands will be skipped if their output files already exist"
    
    # Read commands from file
        set -l commands
    set -l skipped_count 0
    set -l timeout_skipped_count 0
    set -l included_count 0
    
    while read -l line
        if test -n "$line"  # Skip empty lines
            # Extract output path from command
            set -l output_path (extract_output_path "$line")
            
            # If output path exists, check if file or timeout marker exists
            if test -n "$output_path"
                if test -e "$output_path"
                    echo "Skipping command: Output file already exists: $output_path"
                    set skipped_count (math $skipped_count + 1)
                    continue
                end
                
                # Check for timeout marker
                if test -e "$output_path.timeout"
                    echo "Skipping command: Previously timed out: $output_path"
                    set timeout_skipped_count (math $timeout_skipped_count + 1)
                    continue
                end
            end
            
            # Add command to the list
            set -a commands "$line"
            set included_count (math $included_count + 1)
        end
    end < $command_file
    
    set -l total_commands (count $commands)
    set -l next_command 1
    
    echo "Commands summary:"
    echo "  To execute: $included_count"
    echo "  Skipped (output exists): $skipped_count"
    echo "  Skipped (previous timeout): $timeout_skipped_count"
    
    # Main processing loop
    while test $next_command -le $total_commands
        for core in (seq 0 (math $num_cores - 1))
            # Check if core is free and we have commands left to run
            if test $next_command -le $total_commands; and is_core_free $core $threshold
                set -l cmd $commands[$next_command]
                
                # Create unique output filename for logs
                set -l timestamp (date +%Y%m%d_%H%M%S_%N)
                set -l output_file "$output_dir/task_{$next_command}_core_{$core}_{$timestamp}.log"
                
                # Extract the actual output path for reference
                set -l program_output_path (extract_output_path "$cmd")
                
                echo "Starting command #$next_command on core $core:"
                echo "  Command: $cmd"
                echo "  Log:     $output_file"
                if test -n "$program_output_path"
                    echo "  Output:  $program_output_path"
                end
                
                # Build the wrapper command with time and memory limits
                set -l wrapper_cmd ""
                
                # Add memory limit if specified
                if test $memory_limit -gt 0
                    set -l memory_kb (math "$memory_limit * 1024")
                    set wrapper_cmd "ulimit -v $memory_kb; $wrapper_cmd"
                end
                
                # Add time limit if specified
                if test $time_limit -gt 0
                    set wrapper_cmd "timeout $time_limit $wrapper_cmd"
                end
                
                # Add the actual command
                set wrapper_cmd "$wrapper_cmd $cmd"
                
                # Run the command with appropriate limits on the specified core
                taskset -c $core fish -c "$wrapper_cmd > '$output_file' 2>&1; 
                    set -l exit_code \$status
                    if test \$exit_code -eq 124
                        echo 'TIMEOUT: Command exceeded time limit of $time_limit seconds' >> '$output_file'
                        
                        # Create a timeout marker file if we have an output path
                        set -l program_output_path '$program_output_path'
                        if test -n \"\$program_output_path\"
                            echo \"Command timed out after $time_limit seconds\" > \"\$program_output_path.timeout\"
                            echo \"Command: $cmd\" >> \"\$program_output_path.timeout\"
                            echo \"Date: \$(date)\" >> \"\$program_output_path.timeout\"
                        end
                    else if test \$exit_code -ne 0
                        echo 'ERROR: Command failed with exit code \$exit_code' >> '$output_file'
                    end" &
                
                set next_command (math $next_command + 1)
                echo "Progress: $next_command of $total_commands commands started"
                
                # Small delay to prevent overwhelming the system
                sleep 0.5
            end
        end
        
        # Pause briefly before checking cores again
        sleep 1
    end
    
    echo "All commands have been scheduled. Waiting for completion..."
    
    # Wait for all background jobs to complete
    wait
    
    echo "All commands completed successfully"
    echo "Output logs are available in $output_dir/"
    
    # Report on any commands that exceeded limits
    set -l timeout_count (grep -l "TIMEOUT: Command exceeded time limit" $output_dir/*.log | wc -l)
    if test $timeout_count -gt 0
        echo "WARNING: $timeout_count commands exceeded the time limit"
    end
    
    set -l memory_error_count (grep -l "Cannot allocate memory" $output_dir/*.log | wc -l)
    if test $memory_error_count -gt 0
        echo "WARNING: $memory_error_count commands likely exceeded the memory limit"
    end
end

# Check arguments
if test (count $argv) -lt 1
    echo "Usage: "(status current-filename)" COMMAND_FILE [CPU_THRESHOLD [TIME_LIMIT [MEMORY_LIMIT]]]"
    echo "  COMMAND_FILE:   File containing commands to execute (one per line)"
    echo "  CPU_THRESHOLD:  CPU utilization threshold to consider a core free (default: 30%)"
    echo "  TIME_LIMIT:     Maximum execution time in seconds (default: 0, no limit)"
    echo "  MEMORY_LIMIT:   Maximum memory usage in MB (default: 0, no limit)"
    exit 1
end

# Validate command file
if not test -f $argv[1]
    echo "Error: Command file '"$argv[1]"' not found"
    exit 1
end

# Run the executor
core_runner $argv
