//go:build windows

package uvm

import (
	"context"
	"fmt"

	"github.com/Microsoft/go-winio/pkg/guid"
	"github.com/sirupsen/logrus"

	"github.com/Microsoft/hcsshim/internal/hcs/resourcepaths"
	hcsschema "github.com/Microsoft/hcsshim/internal/hcs/schema2"
	"github.com/Microsoft/hcsshim/internal/log"
	"github.com/Microsoft/hcsshim/internal/logfields"
	"github.com/Microsoft/hcsshim/internal/protocol/guestrequest"
	"github.com/Microsoft/hcsshim/internal/protocol/guestresource"
	"github.com/Microsoft/hcsshim/osversion"
)

const (
	GPUDeviceIDType         = "gpu"
	VPCILocationPathIDType  = "vpci-location-path"
	VPCIClassGUIDTypeLegacy = "class"
	VPCIClassGUIDType       = "vpci-class-guid"
	VPCIDeviceIDTypeLegacy  = "vpci"
	VPCIDeviceIDType        = "vpci-instance-id"
)

type VPCIDeviceID struct {
	deviceInstanceID     string
	virtualFunctionIndex uint16
}

func NewVPCIDeviceID(deviceInstanceID string, virtualFunctionIndex uint16) VPCIDeviceID {
	return VPCIDeviceID{
		deviceInstanceID:     deviceInstanceID,
		virtualFunctionIndex: virtualFunctionIndex,
	}
}

// VPCIDevice represents a vpci device. Holds its guid and a handle to the uvm it
// belongs to.
type VPCIDevice struct {
	// vm is the handle to the UVM that this device belongs to
	vm *UtilityVM
	// VMBusGUID is the instance ID for this device when it is exposed via VMBus
	VMBusGUID string
	// deviceInstanceID is the instance ID of the device on the host
	deviceInstanceID string
	// virtualFunctionIndex is the function index for the pci device to assign
	virtualFunctionIndex uint16
	// refCount stores the number of references to this device in the UVM
	refCount uint32
}

// Release frees the resources of the corresponding vpci device
func (vpci *VPCIDevice) Release(ctx context.Context) error {
	if err := vpci.vm.RemoveDevice(ctx, vpci.deviceInstanceID, vpci.virtualFunctionIndex); err != nil {
		return fmt.Errorf("failed to remove VPCI device: %w", err)
	}
	return nil
}

func IsValidDeviceType(deviceType string) bool {
	return (deviceType == VPCIDeviceIDType) ||
		(deviceType == VPCIDeviceIDTypeLegacy) ||
		(deviceType == GPUDeviceIDType)
}

// AssignDevice assigns a vpci device to a uvm.
// If the device already exists, the stored VPCIDevice's ref count is increased
// and the VPCIDevice is returned.
// Otherwise, a new request is made to assign the target device indicated by the deviceID
// onto the UVM. A new VPCIDevice entry is made on the UVM and the VPCIDevice is returned
// to the caller.
// Allow callers to specify the vmbus guid they want the device to show up with.
func (uvm *UtilityVM) AssignDevice(ctx context.Context, deviceID string, index uint16, vmBusGUID string) (*VPCIDevice, error) {
	if vmBusGUID == "" {
		guid, err := guid.NewV4()
		if err != nil {
			return nil, err
		}
		vmBusGUID = guid.String()
	}

	ctx, entry := log.SetEntry(ctx, logrus.Fields{
		logfields.UVMID: uvm.ID(),
		"deviceID":      deviceID,
		"funcIndex":     index,
	})

	key := VPCIDeviceID{
		deviceInstanceID:     deviceID,
		virtualFunctionIndex: index,
	}

	uvm.m.Lock()
	defer uvm.m.Unlock()

	existingVPCIDevice := uvm.vpciDevices[key]
	if existingVPCIDevice != nil {
		existingVPCIDevice.refCount++
		entry.WithFields(logrus.Fields{
			"vmBusGUID": existingVPCIDevice.VMBusGUID,
			"refCount":  existingVPCIDevice.refCount,
		}).Debug("reusing existing vPCI device")
		return existingVPCIDevice, nil
	}

	var propagateAffinity *bool
	T := true
	if osversion.Get().Build >= osversion.V25H1Server {
		propagateAffinity = &T
	}

	targetDevice := hcsschema.VirtualPciDevice{
		Functions: []hcsschema.VirtualPciFunction{
			{
				DeviceInstancePath: deviceID,
				VirtualFunction:    index,
			},
		},
		PropagateNumaAffinity: propagateAffinity,
	}

	request := &hcsschema.ModifySettingRequest{
		ResourcePath: fmt.Sprintf(resourcepaths.VirtualPCIResourceFormat, vmBusGUID),
		RequestType:  guestrequest.RequestTypeAdd,
		Settings:     targetDevice,
	}

	// WCOW (when supported) does not require a guest request as part of the
	// device assignment
	if uvm.operatingSystem != "windows" {
		// for LCOW, we need to make sure that specific paths relating to the
		// device exist so they are ready to be used by later
		// work in openGCS
		request.GuestRequest = guestrequest.ModificationRequest{
			ResourceType: guestresource.ResourceTypeVPCIDevice,
			RequestType:  guestrequest.RequestTypeAdd,
			Settings: guestresource.LCOWMappedVPCIDevice{
				VMBusGUID: vmBusGUID,
			},
		}
	}

	entry.WithField("vmBusGUID", vmBusGUID).Debug("assigning vPCI device")

	if err := uvm.modify(ctx, request); err != nil {
		return nil, err
	}
	device := &VPCIDevice{
		vm:                   uvm,
		VMBusGUID:            vmBusGUID,
		deviceInstanceID:     key.deviceInstanceID,
		virtualFunctionIndex: key.virtualFunctionIndex,
		refCount:             1,
	}
	uvm.vpciDevices[key] = device

	return device, nil
}

// RemoveDevice removes a vpci device from a uvm when there are
// no more references to a given VPCIDevice. Otherwise, decrements
// the reference count of the stored VPCIDevice and returns nil.
func (uvm *UtilityVM) RemoveDevice(ctx context.Context, deviceInstanceID string, index uint16) error {
	key := VPCIDeviceID{
		deviceInstanceID:     deviceInstanceID,
		virtualFunctionIndex: index,
	}

	ctx, entry := log.SetEntry(ctx, logrus.Fields{
		logfields.UVMID: uvm.ID(),
		"deviceID":      deviceInstanceID,
		"funcIndex":     index,
	})

	uvm.m.Lock()
	defer uvm.m.Unlock()

	vpci := uvm.vpciDevices[key]
	if vpci == nil {
		entry.Error("failed to remove nonexistent vPCI device ")
		return fmt.Errorf("no device with ID %s and index %d is present on the uvm %s", deviceInstanceID, index, uvm.ID())
	}

	vpci.refCount--
	entry = entry.WithFields(logrus.Fields{
		"vmBusGUID": vpci.VMBusGUID,
		"refCount":  vpci.refCount,
	})

	entry.Trace("vPCI device refcount updated")
	if vpci.refCount == 0 {
		entry.Debug("removing vPCI device from uVM")

		delete(uvm.vpciDevices, key)
		return uvm.modify(ctx, &hcsschema.ModifySettingRequest{
			ResourcePath: fmt.Sprintf(resourcepaths.VirtualPCIResourceFormat, vpci.VMBusGUID),
			RequestType:  guestrequest.RequestTypeRemove,
		})
	}

	return nil
}