#!/bin/bash

# -------------------------------------------------------------------------- #
# Copyright 2002-2025, OpenNebula Project, OpenNebula Systems                #
#                                                                            #
# Licensed under the Apache License, Version 2.0 (the "License"); you may    #
# not use this file except in compliance with the License. You may obtain    #
# a copy of the License at                                                   #
#                                                                            #
# http://www.apache.org/licenses/LICENSE-2.0                                 #
#                                                                            #
# Unless required by applicable law or agreed to in writing, software        #
# distributed under the License is distributed on an "AS IS" BASIS,          #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #
# See the License for the specific language governing permissions and        #
# limitations under the License.                                             #
#--------------------------------------------------------------------------- #

# ------------------------------------------------------------------------------
# HELPERS
# ------------------------------------------------------------------------------

# Get PCI devices UUID
function get_uuids() {
    uuids="$(xmllint --format --xpath '/VM/TEMPLATE/PCI/UUID/text()' "$1" 2>/dev/null)"
    echo "$uuids" | sed -e 's/<!\[CDATA\[//g; s/\]\]>//g'
}

# Get value from XML and remove CDATA part
function get_xpath_val() {
    echo "$1" | xmllint --format --xpath "$2/text()" - | sed -e 's/<!\[CDATA\[//g; s/\]\]>//g'
}

# Get mdev path used to (de)activate mediated device
function vgpuctl() {
    pci="$(xmllint --format --xpath "/VM/TEMPLATE/PCI[UUID='$1']" "$2" 2>/dev/null)"

    # Get specific information about the PCI
    domain=$(get_xpath_val "$pci" "/PCI/DOMAIN")
    bus=$(get_xpath_val "$pci" "/PCI/BUS")
    slot=$(get_xpath_val "$pci" "/PCI/SLOT")
    func=$(get_xpath_val "$pci" "/PCI/FUNCTION")
    profile=$(get_xpath_val "$pci" "/PCI/PROFILE")
    mode=$(get_xpath_val "$pci" "/PCI/MDEV_MODE")

    if [[ "$mode" == "legacy" || -z "$mode" ]]; then
        # Generate mdev path
        mdev="/sys/class/mdev_bus/$domain:$bus:$slot.$func"

        if [[ ! -d $mdev ]]; then
            error_message "Directory '$mdev' does not exist"
            exit 1
        fi

        if [ -z "$profile" ]; then
            profile="$(ls "$mdev/mdev_supported_types" | head -n1)"
        fi

        mdev="$mdev/mdev_supported_types/$profile/"

        case "$3" in
        "create")
            if ! echo "$1" > "$mdev/create"; then
                error_message "Error creating mediated device"
                exit 1
            fi
            ;;
        "remove")
            if ! echo "1" > "$mdev/devices/$1/remove"; then
                error_message "Error removing mediated device"
            fi
            ;;
        esac
    else
        pci="/sys/bus/pci/devices/$domain:$bus:$slot.$func"

        if [[ ! -d "${pci}" ]]; then
            error_message "Directory '$pci' does not exist"
            exit 1
        fi

        ppath="${pci}/nvidia/creatable_vgpu_types"

        if [ -z "$profile" ]; then
            profile=$(sed -n '2p' ${ppath} | cut -f1 -d':' | tr -d '[:blank:]')
        else
            profile=${profile%% *}

            if [[ "$3" == "create" ]] && ! grep -q "${profile}" ${ppath}; then
                error_message "Profile '$profile' not supported by vGPU"
                exit 1
            fi
        fi

        case "$3" in
        "create")
            if ! echo "${profile}" > "${pci}/nvidia/current_vgpu_type"; then
                error_message "Error activating vgpu with profile ${profile}"
                exit 1
            fi
            ;;
        "remove")
            if ! echo "0" > "${pci}/nvidia/current_vgpu_type"; then
                error_message "Error deactivating vgpu"
            fi
            ;;
        esac
    fi
}

function update_pci_info() {
    # Get host ID
    host_id="$(xmllint --xpath '/VM/HISTORY_RECORDS/HISTORY[last()]/HID/text()' "$VM" 2>/dev/null)"

    if [ -n "$host_id" ]; then

        PROBE_PATH="/var/tmp/one/im/kvm-probes.d/host/system/pci.rb"

        # Execute probe and collect the output
        if [ -f "$PROBE_PATH" ] && [ -x "$PROBE_PATH" ]; then
            payload="$("$PROBE_PATH" "kvm" < "$VM")"

            # Send monitoring data to monitord
            if [ $? -eq 0 ] && [ -n "$payload" ]; then
                send_to_monitor "SYSTEM_HOST" "SUCCESS" "$host_id" "$payload"
            fi
        fi
    fi
}

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------

ACTION=${1,,}

# create -> vm.xml path
# delete -> vm.xml path
# datastore -> vm deploy ID
VM="$2"

XPATH="$(dirname $0)/datastore/xpath.rb --stdin"

source "$(dirname $0)/etc/vmm/kvm/kvmrc"
source "$(dirname $0)/scripts_common.sh"

case "$ACTION" in
    "create")
        uuids="$(get_uuids "$VM")"

        if [ -n "$uuids" ]; then
            for uuid in $uuids; do
                vgpuctl "$uuid" "$VM" "create"
            done
        fi

        update_pci_info
    ;;
    "delete")
        # Not exit with error, just log
        uuids="$(get_uuids "$VM")"

        if [ -n "$uuids" ]; then
            for uuid in $uuids; do
                vgpuctl "$uuid" "$VM" "remove"
            done
        fi

        update_pci_info
    ;;
    "datastore")
        METADATA_XML="$(virsh --connect "$LIBVIRT_URI" metadata "$VM" "$LIBVIRT_MD_URI" "$LIBVIRT_MD_KEY")"

        unset i XPATH_ELEMENTS

        while IFS= read -r -d '' element; do
            XPATH_ELEMENTS[i++]="$element"
        done < <(echo "$METADATA_XML" | $XPATH /vm/system_datastore/)

        unset i

        DATASTORE_PATH="${XPATH_ELEMENTS[i++]}"

        if [ -z "$DATASTORE_PATH" ]; then
            error_message "Datastore path not found"
            exit 1
        fi

        echo "$DATASTORE_PATH"
    ;;
    *)
        error_message "Unsupported action '$ACTION'"
        exit 1
esac
