vctp2/internal/vcenter/vcenter.go

package vcenter

import (
	"context"
	"fmt"
	"log"
	"log/slog"
	"net/url"
	"path"
	"strings"

	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/find"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/view"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

type Vcenter struct {
	Logger      *slog.Logger
	Vurl        string
	ctx         context.Context
	client      *govmomi.Client
	credentials *VcenterLogin
}

type VcenterLogin struct {
	Username string
	Password string
	Insecure bool
}

type VmProperties struct {
	Vm           mo.VirtualMachine
	ResourcePool string
}

type HostLookup struct {
	Cluster    string
	Datacenter string
}

type FolderLookup map[string]string

// New creates a new Vcenter with the given logger
func New(logger *slog.Logger, creds *VcenterLogin) *Vcenter {

	//ctx, cancel := context.WithCancel(context.Background())
	//defer cancel()

	return &Vcenter{
		Logger:      logger,
		ctx:         context.Background(),
		credentials: creds,
	}
}

func (v *Vcenter) Login(vUrl string) error {
	// Connect to vCenter
	u, err := soap.ParseURL(vUrl)
	if err != nil {
		log.Fatalf("Error parsing vCenter URL: %s", err)
	}
	v.Vurl = vUrl

	u.User = url.UserPassword(v.credentials.Username, v.credentials.Password)

	/*
		c, err := govmomi.NewClient(ctx, u, insecure)
		if err != nil {
			log.Fatalf("Error connecting to vCenter: %s", err)
		}
	*/

	c, err := govmomi.NewClient(v.ctx, u, v.credentials.Insecure)
	if err != nil {
		v.Logger.Error("Unable to connect to vCenter", "error", err)
		return fmt.Errorf("unable to connect to vCenter : %s", err)
	}

	//defer c.Logout(v.ctx)

	v.client = c

	v.Logger.Debug("successfully connected to vCenter", "url", vUrl, "username", v.credentials.Username)

	return nil
}

func (v *Vcenter) Logout() error {
	//v.Logger.Debug("vcenter logging out")

	if v.ctx == nil {
		v.Logger.Warn("Nil context, unable to logout")
		return nil
	}

	if v.client.Valid() {
		//v.Logger.Debug("vcenter client is valid. Logging out")
		return v.client.Logout(v.ctx)
	} else {
		v.Logger.Debug("vcenter client is not valid")
		return nil
	}

}

func (v *Vcenter) GetAllVmReferences() ([]*object.VirtualMachine, error) {
	var results []*object.VirtualMachine
	finder := find.NewFinder(v.client.Client, true)

	m := view.NewManager(v.client.Client)

	vms, err := m.CreateContainerView(v.ctx, v.client.ServiceContent.RootFolder, []string{"VirtualMachine"}, true)
	if err != nil {
		return nil, err
	}
	defer vms.Destroy(v.ctx)

	// List all datacenters
	datacenters, err := finder.DatacenterList(v.ctx, "*")
	if err != nil {
		return nil, fmt.Errorf("failed to list datacenters: %w", err)
	}

	for _, dc := range datacenters {
		v.Logger.Debug("Getting VMs in", "datacenter", dc.Name())
		// Set the current datacenter
		finder.SetDatacenter(dc)

		// Get the list of all virtual machines in the current datacenter
		vms, err := finder.VirtualMachineList(v.ctx, "*")
		if err != nil {
			v.Logger.Error("Failed to list VMs in", "datacenter", dc.Name(), "error", err)
			continue
		}

		for _, vm := range vms {
			//vmRef := vm.Reference()
			//v.Logger.Debug("result", "vm", vm, "MoRef", vmRef, "path", vm.InventoryPath)
			results = append(results, vm)
		}

	}

	v.Logger.Debug("Found VM references", "count", len(results))

	return results, err
}

// GetAllVMsWithProps returns all VMs with the properties needed for snapshotting in a single property-collector call.
func (v *Vcenter) GetAllVMsWithProps() ([]mo.VirtualMachine, error) {
	m := view.NewManager(v.client.Client)
	cv, err := m.CreateContainerView(v.ctx, v.client.ServiceContent.RootFolder, []string{"VirtualMachine"}, true)
	if err != nil {
		return nil, fmt.Errorf("failed to create VM container view: %w", err)
	}
	defer cv.Destroy(v.ctx)

	var vms []mo.VirtualMachine
	props := []string{
		"name",
		"parent",
		"config.uuid",
		"config.createDate",
		"config.hardware",
		"config.managedBy",
		"config.template",
		"runtime.powerState",
		"runtime.host",
		"resourcePool",
	}
	if err := cv.Retrieve(v.ctx, []string{"VirtualMachine"}, props, &vms); err != nil {
		return nil, fmt.Errorf("failed to retrieve VMs: %w", err)
	}
	return vms, nil
}

func (v *Vcenter) BuildHostLookup() (map[string]HostLookup, error) {
	finder := find.NewFinder(v.client.Client, true)
	datacenters, err := finder.DatacenterList(v.ctx, "*")
	if err != nil {
		return nil, fmt.Errorf("failed to list datacenters: %w", err)
	}

	lookup := make(map[string]HostLookup)
	clusterCache := make(map[string]string)

	for _, dc := range datacenters {
		finder.SetDatacenter(dc)
		hosts, err := finder.HostSystemList(v.ctx, "*")
		if err != nil {
			v.Logger.Warn("failed to list hosts for datacenter", "datacenter", dc.Name(), "error", err)
			continue
		}

		for _, host := range hosts {
			ref := host.Reference()
			var moHost mo.HostSystem
			if err := v.client.RetrieveOne(v.ctx, ref, []string{"parent"}, &moHost); err != nil {
				v.Logger.Warn("failed to retrieve host info", "host", host.Name(), "error", err)
				continue
			}

			clusterName := ""
			if moHost.Parent != nil {
				if cached, ok := clusterCache[moHost.Parent.Value]; ok {
					clusterName = cached
				} else {
					var moCompute mo.ComputeResource
					if err := v.client.RetrieveOne(v.ctx, *moHost.Parent, []string{"name"}, &moCompute); err == nil {
						clusterName = moCompute.Name
						clusterCache[moHost.Parent.Value] = clusterName
					}
				}
			}

			lookup[ref.Value] = HostLookup{
				Cluster:    clusterName,
				Datacenter: dc.Name(),
			}
		}
	}

	return lookup, nil
}

func (v *Vcenter) BuildFolderPathLookup() (FolderLookup, error) {
	m := view.NewManager(v.client.Client)
	folders, err := m.CreateContainerView(v.ctx, v.client.ServiceContent.RootFolder, []string{"Folder"}, true)
	if err != nil {
		return nil, err
	}
	defer folders.Destroy(v.ctx)

	var results []mo.Folder
	if err := folders.Retrieve(v.ctx, []string{"Folder"}, []string{"name", "parent"}, &results); err != nil {
		return nil, err
	}

	nameByID := make(map[string]string, len(results))
	parentByID := make(map[string]*types.ManagedObjectReference, len(results))
	for _, folder := range results {
		nameByID[folder.Reference().Value] = folder.Name
		parentByID[folder.Reference().Value] = folder.Parent
	}

	paths := make(FolderLookup, len(results))
	var buildPath func(id string) string
	buildPath = func(id string) string {
		if pathValue, ok := paths[id]; ok {
			return pathValue
		}
		name, ok := nameByID[id]
		if !ok {
			return ""
		}
		parent := parentByID[id]
		if parent == nil || parent.Type == "Datacenter" {
			paths[id] = path.Join("/", name)
			return paths[id]
		}
		if parent.Type != "Folder" {
			paths[id] = path.Join("/", name)
			return paths[id]
		}
		parentPath := buildPath(parent.Value)
		if parentPath == "" {
			paths[id] = path.Join("/", name)
			return paths[id]
		}
		paths[id] = path.Join(parentPath, name)
		return paths[id]
	}

	for id := range nameByID {
		_ = buildPath(id)
	}

	return paths, nil
}

func (v *Vcenter) GetVMFolderPathFromLookup(vm mo.VirtualMachine, lookup FolderLookup) (string, bool) {
	if vm.Parent == nil || lookup == nil {
		return "", false
	}
	pathValue, ok := lookup[vm.Parent.Value]
	return pathValue, ok
}

func (v *Vcenter) ConvertObjToMoVM(vmObj *object.VirtualMachine) (*mo.VirtualMachine, error) {
	// Use the InventoryPath to extract the datacenter name and VM path
	inventoryPath := vmObj.InventoryPath
	parts := strings.SplitN(inventoryPath, "/", 3)

	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid InventoryPath: %s", inventoryPath)
	}

	// The first part of the path is the datacenter name
	datacenterName := parts[1]

	// Finder to search for datacenter and VM
	finder := find.NewFinder(v.client.Client, true)

	// Find the specific datacenter by name
	datacenter, err := finder.Datacenter(v.ctx, fmt.Sprintf("/%s", datacenterName))
	if err != nil {
		return nil, fmt.Errorf("failed to find datacenter %s: %w", datacenterName, err)
	}

	// Set the found datacenter in the finder
	finder.SetDatacenter(datacenter)

	// Now retrieve the VM using its ManagedObjectReference
	vmRef := vmObj.Reference()

	// Retrieve the full mo.VirtualMachine object for the reference
	var moVM mo.VirtualMachine
	err = v.client.RetrieveOne(v.ctx, vmRef, nil, &moVM)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve VM %s in datacenter %s: %w", vmObj.Name(), datacenterName, err)
	}

	// Return the found mo.VirtualMachine object
	//v.Logger.Debug("Found VM in datacenter", "vm_name", moVM.Name, "dc_name", datacenterName)
	return &moVM, nil
}

func (v *Vcenter) ConvertObjToMoHost(hostObj *object.HostSystem) (*mo.HostSystem, error) {
	// Use the InventoryPath to extract the datacenter name and Host path
	inventoryPath := hostObj.InventoryPath
	parts := strings.SplitN(inventoryPath, "/", 3)
	v.Logger.Debug("inventory path", "parts", parts)

	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid InventoryPath: %s", inventoryPath)
	}

	// The first part of the path is the datacenter name
	datacenterName := parts[1]

	// Finder to search for datacenter and VM
	finder := find.NewFinder(v.client.Client, true)

	// Find the specific datacenter by name
	datacenter, err := finder.Datacenter(v.ctx, fmt.Sprintf("/%s", datacenterName))
	if err != nil {
		return nil, fmt.Errorf("failed to find datacenter %s: %w", datacenterName, err)
	}

	// Set the found datacenter in the finder
	finder.SetDatacenter(datacenter)

	// Now retrieve the VM using its ManagedObjectReference
	hostRef := hostObj.Reference()

	// Retrieve the full mo.HostSystem object for the reference
	var moHost mo.HostSystem
	err = v.client.RetrieveOne(v.ctx, hostRef, nil, &moHost)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve Host %s in datacenter %s: %w", hostObj.Name(), datacenterName, err)
	}

	// Return the found mo.HostSystem object
	v.Logger.Debug("Found Host in datacenter", "host_name", moHost.Name, "dc_name", datacenterName)
	return &moHost, nil
}

func (v *Vcenter) GetHostSystemObject(hostRef types.ManagedObjectReference) (*mo.HostSystem, error) {
	finder := find.NewFinder(v.client.Client, true)

	// List all datacenters
	datacenters, err := finder.DatacenterList(v.ctx, "*")
	if err != nil {
		return nil, fmt.Errorf("failed to list datacenters: %w", err)
	}

	for _, dc := range datacenters {
		v.Logger.Debug("Checking dc for host", "name", dc.Name(), "hostRef", hostRef.String())
		// Set the current datacenter
		finder.SetDatacenter(dc)

		var hs mo.HostSystem
		err := v.client.RetrieveOne(v.ctx, hostRef, nil, &hs)
		if err != nil {
			return nil, err
		} else {
			v.Logger.Debug("Found hostsystem", "name", hs.Name)
			return &hs, nil
		}
	}

	return nil, nil
}

// Function to find the cluster or compute resource from a host reference
func (v *Vcenter) GetClusterFromHost(hostRef *types.ManagedObjectReference) (string, error) {
	// Get the host object
	host, err := v.GetHostSystemObject(*hostRef)
	if err != nil {
		v.Logger.Error("cant get host", "error", err)
		return "", err
	}

	v.Logger.Debug("host parent", "parent", host.Parent)

	if host.Parent.Type == "ClusterComputeResource" {
		// Retrieve properties of the compute resource
		var moCompute mo.ComputeResource
		err = v.client.RetrieveOne(v.ctx, *host.Parent, nil, &moCompute)
		if err != nil {
			return "", fmt.Errorf("failed to retrieve compute resource: %w", err)
		}
		v.Logger.Debug("VM is on host in cluster/compute resource", "name", moCompute.Name)
		return moCompute.Name, nil
	}

	return "", nil
}

// Function to determine the datacenter a VM belongs to
func (v *Vcenter) GetDatacenterForVM(vm mo.VirtualMachine) (string, error) {
	// Start with the VM's parent reference
	ref := vm.Reference()

	// Traverse the inventory hierarchy upwards to find the datacenter
	for {
		// Get the parent reference of the current object
		parentRef, err := v.getParent(ref)
		if err != nil {
			return "", fmt.Errorf("failed to get parent object: %w", err)
		}

		// If we get a nil parent reference, it means we've hit the root without finding the datacenter
		if parentRef == nil {
			return "", fmt.Errorf("failed to find datacenter for VM")
		}

		// Check if the parent is a Datacenter
		switch parentRef.Type {
		case "Datacenter":
			// If we found a Datacenter, retrieve its properties
			datacenter := object.NewDatacenter(v.client.Client, *parentRef)
			var moDC mo.Datacenter
			err = v.client.RetrieveOne(v.ctx, datacenter.Reference(), nil, &moDC)
			if err != nil {
				return "", fmt.Errorf("failed to retrieve datacenter: %w", err)
			}

			//log.Printf("VM is in datacenter: %s", moDC.Name)
			v.Logger.Debug("VM datacenter found", "vm_name", vm.Name, "dc_name", moDC.Name)
			return moDC.Name, nil

		default:
			// Continue traversing upwards if not a Datacenter
			ref = *parentRef
		}
	}
}

// Helper function to get the parent ManagedObjectReference of a given object
func (v *Vcenter) getParent(ref types.ManagedObjectReference) (*types.ManagedObjectReference, error) {
	// Retrieve the object's properties
	var obj mo.ManagedEntity
	err := v.client.RetrieveOne(v.ctx, ref, []string{"parent"}, &obj)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve parent of object: %w", err)
	}

	// Return the parent reference
	if obj.Parent != nil {
		return obj.Parent, nil
	}
	return nil, nil
}

func (v *Vcenter) FindVMByName(vmName string) ([]mo.VirtualMachine, error) {
	m := view.NewManager(v.client.Client)

	vms, err := m.CreateContainerView(v.ctx, v.client.ServiceContent.RootFolder, []string{"VirtualMachine"}, true)
	if err != nil {
		return nil, err
	}
	defer vms.Destroy(v.ctx)

	var matchingVMs []mo.VirtualMachine

	err = vms.Retrieve(v.ctx, []string{"VirtualMachine"}, []string{"name"}, &matchingVMs)
	if err != nil {
		return nil, err
	}

	// Temporarily just return all VMs
	//return matchingVMs, nil

	var result []mo.VirtualMachine

	for _, vm := range matchingVMs {
		if vm.Name == vmName {
			result = append(result, vm)
		}
	}

	return result, nil
}

func (v *Vcenter) FindVMByID(vmID string) (*VmProperties, error) {

	v.Logger.Debug("searching for vm id", "vm_id", vmID)

	finder := find.NewFinder(v.client.Client, true)

	// List all datacenters
	datacenters, err := finder.DatacenterList(v.ctx, "*")
	if err != nil {
		return nil, fmt.Errorf("failed to list datacenters: %w", err)
	}

	for _, dc := range datacenters {
		// Set the current datacenter
		finder.SetDatacenter(dc)

		// Create a ManagedObjectReference for the VM
		vmRef := types.ManagedObjectReference{
			Type:  "VirtualMachine",
			Value: vmID,
		}

		// Try to find the VM by ID in the current datacenter
		//vm, err := finder.ObjectReference(v.ctx, vmRef)

		var vm mo.VirtualMachine
		err := v.client.RetrieveOne(v.ctx, vmRef, []string{"config", "name"}, &vm)

		if err == nil {
			return &VmProperties{
				//Datacenter: dc.Name(),
				Vm: vm,
			}, nil
		} else if _, ok := err.(*find.NotFoundError); !ok {
			// If the error is not a NotFoundError, return it
			//return nil, fmt.Errorf("failed to retrieve VM with ID %s in datacenter %s: %w", vmID, dc.Name(), err)
			v.Logger.Debug("Couldn't find vm in datacenter", "vm_id", vmID, "datacenter_name", dc.Name())
		} else {
			return nil, fmt.Errorf("failed to retrieve VM: %w", err)
		}
	}

	return nil, fmt.Errorf("VM with ID %s not found in any datacenter", vmID)
}

func (v *Vcenter) FindVMByIDWithDatacenter(vmID string, dcID string) (*mo.VirtualMachine, error) {
	var err error
	//v.Logger.Debug("searching for vm id", "vm_id", vmID, "datacenter_id", dcID)

	finder := find.NewFinder(v.client.Client, true)

	// Create a ManagedObjectReference for the datacenter
	dcRef := types.ManagedObjectReference{
		Type:  "Datacenter",
		Value: dcID,
	}

	// Convert the reference to a Datacenter object
	datacenter := object.NewDatacenter(v.client.Client, dcRef)
	if datacenter == nil {
		return nil, fmt.Errorf("Datacenter with id %s not found", dcID)
	}

	// Use finder.SetDatacenter to set the datacenter
	finder.SetDatacenter(datacenter)

	// Create a ManagedObjectReference for the VM
	vmRef := types.ManagedObjectReference{
		Type:  "VirtualMachine",
		Value: vmID,
	}

	var vm mo.VirtualMachine
	//err := v.client.RetrieveOne(v.ctx, vmRef, []string{"config", "name"}, &vm)
	err = v.client.RetrieveOne(v.ctx, vmRef, nil, &vm)
	if err == nil {
		//v.Logger.Debug("Found VM")

		return &vm, nil

	} else if _, ok := err.(*find.NotFoundError); !ok {
		// If the error is not a NotFoundError, return it
		//return nil, fmt.Errorf("failed to retrieve VM with ID %s in datacenter %s: %w", vmID, dc.Name(), err)
		//v.Logger.Debug("Couldn't find vm in datacenter", "vm_id", vmID, "datacenter_id", dcID)
		return nil, nil
	} else {
		return nil, fmt.Errorf("failed to retrieve VM: %w", err)
	}
}

// Helper function to retrieve the resource pool for the VM
func (v *Vcenter) GetVmResourcePool(vm mo.VirtualMachine) (string, error) {
	var resourcePool string
	if vm.ResourcePool != nil {
		rp := object.NewResourcePool(v.client.Client, *vm.ResourcePool)
		rpName, err := rp.ObjectName(v.ctx)
		if err != nil {
			v.Logger.Error("failed to get resource pool name", "error", err)
			return resourcePool, err
		} else {
			//v.Logger.Debug("Found resource pool name", "rp_name", rpName)
			resourcePool = rpName
		}
	}
	return resourcePool, nil
}

// BuildResourcePoolLookup creates a cache of resource pool MoRef -> name for fast lookups.
func (v *Vcenter) BuildResourcePoolLookup() (map[string]string, error) {
	m := view.NewManager(v.client.Client)
	cv, err := m.CreateContainerView(v.ctx, v.client.ServiceContent.RootFolder, []string{"ResourcePool"}, true)
	if err != nil {
		return nil, fmt.Errorf("failed to create resource pool view: %w", err)
	}
	defer cv.Destroy(v.ctx)

	var pools []mo.ResourcePool
	if err := cv.Retrieve(v.ctx, []string{"ResourcePool"}, []string{"name"}, &pools); err != nil {
		return nil, fmt.Errorf("failed to retrieve resource pools: %w", err)
	}

	lookup := make(map[string]string, len(pools))
	for _, pool := range pools {
		lookup[pool.Reference().Value] = pool.Name
	}
	return lookup, nil
}

// Helper function to retrieve the full folder path for the VM
func (v *Vcenter) GetVMFolderPath(vm mo.VirtualMachine) (string, error) {
	//finder := find.NewFinder(v.client.Client, true)

	v.Logger.Debug("Commencing vm folder path search")

	// Start from the VM's parent
	parentRef := vm.Parent
	if parentRef == nil {
		return "", fmt.Errorf("no parent found for the VM")
	}

	// Traverse the folder hierarchy to build the full folder path
	folderPath := ""
	//v.Logger.Debug("parent is", "parent", parentRef)

	for parentRef.Type != "Datacenter" {
		// Retrieve the parent object
		//parentObj, err := finder.ObjectReference(v.ctx, *parentRef)
		//if err != nil {
		//	return "", fmt.Errorf("failed to find parent object in inventory: %w", err)
		//}

		// Retrieve the folder name

		var parentObj mo.Folder
		err := v.client.RetrieveOne(v.ctx, *parentRef, nil, &parentObj)
		if err != nil {
			v.Logger.Error("Failed to get object for parent reference", "ref", parentRef)
			break
		}

		// Prepend the folder name to the path
		folderPath = path.Join("/", parentObj.Name, folderPath)

		// Move up to the next parent
		//if folder, ok := parentObj.(*object.Folder); ok {
		if parentObj.Parent != nil {
			parentRef = parentObj.Parent
			//v.Logger.Debug("Parent uplevel is", "ref", parentRef)
		} else {
			return "", fmt.Errorf("unexpected parent type: %s", parentObj.Reference().Type)
		}
		//break
	}

	return folderPath, nil
}