#!/bin/bash

log() {
    logger -t "pve-cpu-pin" "$1"
    echo "$(date): $1" >> /var/log/pve-cpu-pin.log
}

get_cpu_affinity() {
    local vmid=$1
    local affinity=$(qm config $vmid | grep ^affinity: | cut -d' ' -f2)
    echo "$affinity"
}

parse_cpu_range() {
    local range=$1
    local result=()
    
    IFS=',' read -ra PARTS <<< "$range"
    for part in "${PARTS[@]}"; do
        if [[ $part == *-* ]]; then
            local start=${part%-*}
            local end=${part#*-}
            for ((i=start; i<=end; i++)); do
                result+=($i)
            done
        else
            result+=($part)
        fi
    done
    
    echo "${result[@]}"
}

main() {
    local vmid=$1
    local phase=$2
    
    if [ "$phase" != "post-start" ]; then
        exit 0
    fi
    
    log "Processing CPU pinning for VM $vmid"
    
    local affinity=$(get_cpu_affinity $vmid)
    if [ -z "$affinity" ]; then
        log "No CPU affinity configuration found for VM $vmid"
        exit 0
    fi
    
    local -a allowed_cpus=($(parse_cpu_range "$affinity"))
    log "Allowed CPUs: ${allowed_cpus[*]}"
    
    local num_allowed=${#allowed_cpus[@]}
    
    local qemu_pid=$(pgrep -f "^/usr/bin/kvm.*-id $vmid\b")
    if [ -z "$qemu_pid" ]; then
        log "Cannot find QEMU process for VM $vmid"
        exit 1
    fi
    
    log "Found QEMU process: $qemu_pid"
    
    sleep 2
    
    local vcpu_count=0
    local management_cpu=${allowed_cpus[-1]}
    
    while read -r tid comm; do
        if [[ $comm =~ ^CPU\ ([0-9]+)/KVM$ ]]; then
            if [ $vcpu_count -lt $((num_allowed-1)) ]; then
                local target_cpu=${allowed_cpus[$vcpu_count]}
                log "Pinning vCPU thread $tid ($comm) to CPU $target_cpu"
                taskset -p -c $target_cpu $tid > /dev/null 2>&1
                ((vcpu_count++))
            fi
        else
            log "Pinning management thread $tid ($comm) to CPU $management_cpu"
            taskset -p -c $management_cpu $tid > /dev/null 2>&1
        fi
    done < <(ps -T -p $qemu_pid -o tid=,comm=)
    
    log "Completed CPU pinning for VM $vmid"
}

if [ $# -lt 2 ]; then
    echo "Usage: $0 <vmid> <phase>"
    exit 1
fi

main "$1" "$2"
