#!/usr/bin/bash

set -euo pipefail

# Source common utilities and appropriate device support library
SCRIPT_DIR="$(dirname "${BASH_SOURCE[0]}")"

# shellcheck source=../lib/common.sh disable=SC1091
source "${SCRIPT_DIR}/../lib/common.sh"

if [[ -n "${TEST_LOGFILE:-}" ]]; then
    # Test mode - use mock device support
    # shellcheck source=../lib/mock-device-support.sh disable=SC1091
    source "${SCRIPT_DIR}/../lib/mock-device-support.sh"
else
    # Normal mode - use standard device support
    # shellcheck source=../lib/device-support.sh disable=SC1091
    source "${SCRIPT_DIR}/../lib/device-support.sh"
fi

# Configuration
LOGFILE="${TEST_LOGFILE:-/var/log/qm-wayland-client-devices.log}"
# shellcheck disable=SC2034  # Used by log() function in device-support.sh
HOOK_NAME="qm-wayland-client-devices"

input="-"
CONTAINER_CONFIG=$(cat "$input")

if [[ -z "$CONTAINER_CONFIG" ]]; then
    log "ERROR" "Failed to read OCI spec from stdin"
    exit 1
fi

GPU_ENABLED=$(echo "$CONTAINER_CONFIG" | jq -r '.annotations["org.containers.qm.wayland-client.gpu"] // empty')

DEVNAME_LIST=()

if [[ "$GPU_ENABLED" =~ ^(true|1|yes|True|TRUE|YES)$ ]]; then
    log "INFO" "Processing Wayland client GPU annotation: $GPU_ENABLED (enabled)"

    # Find all the render devices available
    RENDER_DEVICES=$(discover_gpu_devices)
    log "INFO" "Scanning for GPU render devices"

    for RENDER_DEVICE in $RENDER_DEVICES; do
        DEVNAME_LIST+=("$RENDER_DEVICE")
        log "INFO" "Adding GPU render device: $RENDER_DEVICE"
    done

    log "INFO" "Found ${#DEVNAME_LIST[@]} GPU render devices"
elif [ -n "$GPU_ENABLED" ]; then
    log "INFO" "Wayland client GPU annotation present but disabled: $GPU_ENABLED"
else
    log "INFO" "No Wayland client GPU annotation found"
fi

# Iterate over the DEVNAME_LIST to include the required information in the CONTAINER_CONFIG
if [ ${#DEVNAME_LIST[@]} -gt 0 ]; then
    log "INFO" "Processing ${#DEVNAME_LIST[@]} GPU devices for Wayland client"

    for DEVICE in "${DEVNAME_LIST[@]}"; do
        if ! jq -e ".linux.devices[] | select(.path == \"$DEVICE\")" <<<"$CONTAINER_CONFIG" >/dev/null 2>&1; then
            # Get device info using device support library
            if ! device_info=$(get_device_info "$DEVICE"); then
                log "WARNING" "Failed to get device info for $DEVICE, skipping"
                continue
            fi

            # Parse device info: type:major:minor:file_mode:uid:gid
            IFS=':' read -r dev_type major minor filemode uid gid <<<"$device_info"

            NEW_DEVICE=$(jq -n --arg path "$DEVICE" \
                --arg dev_type "$dev_type" \
                --arg major "$major" \
                --arg minor "$minor" \
                --arg filemode "$filemode" \
                --arg uid "$uid" \
                --arg gid "$gid" \
                '{
                                "path": $path,
                                "type": $dev_type,
                                "major": $major|tonumber,
                                "minor": $minor|tonumber,
                                "fileMode": $filemode|tonumber,
                                "uid": $uid|tonumber,
                                "gid": $gid|tonumber,
                            }')

            NEW_DEV_RESOURCE=$(jq -n \
                --arg dev_type "$dev_type" \
                --arg major "$major" \
                --arg minor "$minor" \
                '{
                                "allow": true,
                                "type": $dev_type,
                                "major": $major|tonumber,
                                "minor": $minor|tonumber,
                                "access": "rwm"
                            }')

            CONTAINER_CONFIG=$(jq ".linux.devices += [$NEW_DEVICE]" <<<"$CONTAINER_CONFIG")
            CONTAINER_CONFIG=$(jq ".linux.resources.devices += [$NEW_DEV_RESOURCE]" <<<"$CONTAINER_CONFIG")
            log "INFO" "Added GPU device: $DEVICE"
        else
            log "INFO" "GPU device already exists in spec: $DEVICE"
        fi
    done

    log "INFO" "Successfully processed all GPU devices for Wayland client"
else
    log "INFO" "No GPU devices to process for Wayland client"
fi

# Initialize log file directory
mkdir -p "$(dirname "$LOGFILE")"
touch "$LOGFILE"

# Count total devices in final spec
total_devices=$(echo "$CONTAINER_CONFIG" | jq '.linux.devices | length // 0')
log "INFO" "Total devices in final spec: $total_devices"
log "INFO" "QM Wayland Client Devices hook completed successfully"

echo "$CONTAINER_CONFIG" | jq .
