spinifex/hive/daemon/daemon_handlers_instance.go

package daemon

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mulgadc/hive/hive/awserrors"
	"github.com/mulgadc/hive/hive/qmp"
	"github.com/mulgadc/hive/hive/types"
	"github.com/mulgadc/hive/hive/utils"
	"github.com/mulgadc/hive/hive/vm"
	"github.com/nats-io/nats.go"
)

// handleEC2RunInstances processes incoming EC2 RunInstances requests
func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) {
	slog.Debug("Received message on subject", "subject", msg.Subject)
	slog.Debug("Message data", "data", string(msg.Data))

	// Extract account ID from NATS header
	accountID := utils.AccountIDFromMsg(msg)
	if accountID == "" {
		slog.Error("handleEC2RunInstances: missing account ID in NATS header")
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	// Initialize runInstancesInput before unmarshaling into it
	runInstancesInput := &ec2.RunInstancesInput{}
	errResp := utils.UnmarshalJsonPayload(runInstancesInput, msg.Data)

	if errResp != nil {
		if err := msg.Respond(errResp); err != nil {
			slog.Error("Failed to respond to NATS request", "err", err)
		}
		slog.Error("Request does not match RunInstancesInput")
		return
	}

	slog.Info("Processing RunInstances request for instance type", "instanceType", *runInstancesInput.InstanceType)

	// Check if instance type is supported
	instanceType, exists := d.resourceMgr.instanceTypes[*runInstancesInput.InstanceType]
	if !exists {
		slog.Error("handleEC2RunInstances instance lookup", "err", awserrors.ErrorInvalidInstanceType, "InstanceType", *runInstancesInput.InstanceType)
		respondWithError(msg, awserrors.ErrorInvalidInstanceType)
		return
	}

	// Validate AMI exists before allocating resources
	if runInstancesInput.ImageId == nil || *runInstancesInput.ImageId == "" {
		slog.Error("handleEC2RunInstances missing ImageId")
		respondWithError(msg, awserrors.ErrorMissingParameter)
		return
	}
	if d.imageService == nil {
		slog.Error("handleEC2RunInstances image service not initialized")
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}
	amiMeta, err := d.imageService.GetAMIConfig(*runInstancesInput.ImageId)
	if err != nil {
		slog.Error("handleEC2RunInstances AMI not found", "imageId", *runInstancesInput.ImageId, "err", err)
		respondWithError(msg, awserrors.ErrorInvalidAMIIDNotFound)
		return
	}
	// Verify the caller can use this AMI: must own it or it must be a system/pre-phase4 AMI.
	// System AMIs have non-account-ID owner aliases (e.g. "self", "hive", empty).
	amiOwner := amiMeta.ImageOwnerAlias
	if amiOwner != "" && amiOwner != accountID {
		if utils.IsAccountID(amiOwner) {
			slog.Warn("handleEC2RunInstances AMI not owned by caller", "imageId", *runInstancesInput.ImageId, "amiOwner", amiOwner, "accountID", accountID)
			respondWithError(msg, awserrors.ErrorInvalidAMIIDNotFound)
			return
		}
	}

	// Validate key pair exists (if specified)
	if runInstancesInput.KeyName != nil && *runInstancesInput.KeyName != "" {
		if err := d.keyService.ValidateKeyPairExists(accountID, *runInstancesInput.KeyName); err != nil {
			slog.Error("handleEC2RunInstances key pair not found", "keyName", *runInstancesInput.KeyName, "err", err)
			respondWithError(msg, awserrors.ErrorInvalidKeyPairNotFound)
			return
		}
	}

	// Determine how many instances to launch based on MinCount/MaxCount
	minCount := int(*runInstancesInput.MinCount)
	maxCount := int(*runInstancesInput.MaxCount)

	// Check how many we can actually launch
	allocatableCount := d.resourceMgr.canAllocate(instanceType, maxCount)

	if allocatableCount < minCount {
		// Cannot satisfy MinCount requirement - fail entirely
		slog.Error("handleEC2RunInstances insufficient capacity", "requested", minCount, "available", allocatableCount, "InstanceType", *runInstancesInput.InstanceType)
		respondWithError(msg, awserrors.ErrorInsufficientInstanceCapacity)
		return
	}

	// Launch up to MaxCount, capped by available capacity
	// Note: canAllocate() already caps at maxCount, so allocatableCount <= maxCount
	launchCount := allocatableCount

	slog.Info("Instance count determined", "min", minCount, "max", maxCount, "launching", launchCount)

	// Allocate resources for all instances upfront
	var allocatedCount int
	for i := 0; i < launchCount; i++ {
		if err := d.resourceMgr.allocate(instanceType); err != nil {
			slog.Error("handleEC2RunInstances allocate failed mid-allocation", "allocated", allocatedCount, "err", err)
			break
		}
		allocatedCount++
	}

	// Check if we still meet MinCount after allocation
	if allocatedCount < minCount {
		// Rollback allocations
		for i := 0; i < allocatedCount; i++ {
			d.resourceMgr.deallocate(instanceType)
		}
		slog.Error("handleEC2RunInstances insufficient capacity after allocation", "allocated", allocatedCount, "minCount", minCount)
		respondWithError(msg, awserrors.ErrorInsufficientInstanceCapacity)
		return
	}

	// Update launchCount to what we actually allocated
	launchCount = allocatedCount

	// Delegate to service for business logic (volume creation, cloud-init, etc.)
	instanceTypeName := ""
	if instanceType.InstanceType != nil {
		instanceTypeName = *instanceType.InstanceType
	}
	slog.Info("Launching EC2 instances", "instanceType", instanceTypeName, "count", launchCount)

	// Create all instances
	var instances []*vm.VM
	var allEC2Instances []*ec2.Instance
	var lastRunErr error

	for i := 0; i < launchCount; i++ {
		instance, ec2Instance, err := d.instanceService.RunInstance(runInstancesInput)
		if err != nil {
			slog.Error("handleEC2RunInstances service.RunInstance failed", "index", i, "err", err)
			lastRunErr = err
			// Deallocate this instance's resources
			d.resourceMgr.deallocate(instanceType)
			continue
		}

		// Auto-create ENI when SubnetId is provided (matches AWS behavior)
		if runInstancesInput.SubnetId != nil && *runInstancesInput.SubnetId != "" && d.vpcService != nil {
			eniOut, eniErr := d.vpcService.CreateNetworkInterface(&ec2.CreateNetworkInterfaceInput{
				SubnetId:    runInstancesInput.SubnetId,
				Description: aws.String("Primary network interface for " + instance.ID),
			}, accountID)
			if eniErr != nil {
				slog.Error("handleEC2RunInstances auto-create ENI failed", "instanceId", instance.ID, "subnetId", *runInstancesInput.SubnetId, "err", eniErr)
				lastRunErr = eniErr
				d.resourceMgr.deallocate(instanceType)
				continue
			}

			eni := eniOut.NetworkInterface
			instance.ENIId = *eni.NetworkInterfaceId
			instance.ENIMac = *eni.MacAddress
			ec2Instance.SetPrivateIpAddress(*eni.PrivateIpAddress)
			ec2Instance.SetSubnetId(*runInstancesInput.SubnetId)
			ec2Instance.SetVpcId(*eni.VpcId)
			ec2Instance.NetworkInterfaces = []*ec2.InstanceNetworkInterface{
				{
					NetworkInterfaceId: eni.NetworkInterfaceId,
					PrivateIpAddress:   eni.PrivateIpAddress,
					MacAddress:         eni.MacAddress,
					SubnetId:           runInstancesInput.SubnetId,
					VpcId:              eni.VpcId,
					Status:             aws.String("in-use"),
					Attachment: &ec2.InstanceNetworkInterfaceAttachment{
						DeviceIndex: aws.Int64(0),
						Status:      aws.String("attached"),
					},
				},
			}

			slog.Info("Auto-created ENI for VPC instance",
				"instanceId", instance.ID,
				"eniId", instance.ENIId,
				"privateIp", *eni.PrivateIpAddress,
				"mac", instance.ENIMac,
			)
		}

		instances = append(instances, instance)
		allEC2Instances = append(allEC2Instances, ec2Instance)
	}

	// Check if we still have enough instances after creation errors
	if len(instances) < minCount {
		// Rollback: deallocate resources for successfully created instances
		// (failed instances already deallocated their resources above)
		for range instances {
			d.resourceMgr.deallocate(instanceType)
		}
		// Propagate the service-layer error if it's a known AWS error code
		errCode := awserrors.ErrorServerInternal
		if lastRunErr != nil {
			if _, ok := awserrors.ErrorLookup[lastRunErr.Error()]; ok {
				errCode = lastRunErr.Error()
			}
		}
		slog.Error("handleEC2RunInstances failed to create minimum instances", "created", len(instances), "minCount", minCount, "err", errCode)
		respondWithError(msg, errCode)
		return
	}

	// Build reservation with all instances
	reservation := ec2.Reservation{}
	reservation.SetReservationId(utils.GenerateResourceID("r"))
	reservation.SetOwnerId(accountID)
	reservation.Instances = allEC2Instances

	// Store reservation reference and account ID in all VMs
	for _, instance := range instances {
		instance.Reservation = &reservation
		instance.AccountID = accountID
	}

	// Respond to NATS immediately with reservation (instances are provisioning)
	jsonResponse, err := json.Marshal(reservation)
	if err != nil {
		slog.Error("handleEC2RunInstances failed to marshal reservation", "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		// Deallocate all resources
		for range instances {
			d.resourceMgr.deallocate(instanceType)
		}
		return
	}
	if err := msg.Respond(jsonResponse); err != nil {
		slog.Error("Failed to respond to NATS request", "err", err)
	}

	// Add all instances to state immediately so DescribeInstances can find them
	// while volumes are being prepared and VMs are launching
	d.Instances.Mu.Lock()
	for _, instance := range instances {
		d.Instances.VMS[instance.ID] = instance
	}
	d.Instances.Mu.Unlock()

	if err := d.WriteState(); err != nil {
		slog.Error("handleEC2RunInstances failed to write initial state", "err", err)
	}

	slog.Info("Instances added to state with pending status", "count", len(instances))

	// Launch all instances (volumes and VMs)
	var successCount int
	for _, instance := range instances {
		// Pre-compute dev MAC so cloud-init can generate per-interface netplan
		// that suppresses the default route on the dev/hostfwd NIC.
		if d.config.Daemon.DevNetworking && instance.ENIId != "" {
			instance.DevMAC = generateDevMAC(instance.ID)
		}

		// Prepare the root volume, cloud-init, EFI drives via NBD (AMI clone to new volume)
		volumeInfos, err := d.instanceService.GenerateVolumes(runInstancesInput, instance)
		if err != nil {
			slog.Error("handleEC2RunInstances GenerateVolumes failed", "instanceId", instance.ID, "err", err)
			d.resourceMgr.deallocate(instanceType)
			d.markInstanceFailed(instance, "volume_preparation_failed")
			continue
		}

		// Populate BlockDeviceMappings on the ec2.Instance
		instance.Instance.BlockDeviceMappings = make([]*ec2.InstanceBlockDeviceMapping, 0, len(volumeInfos))
		for _, vi := range volumeInfos {
			mapping := &ec2.InstanceBlockDeviceMapping{}
			mapping.SetDeviceName(vi.DeviceName)
			mapping.Ebs = &ec2.EbsInstanceBlockDevice{}
			mapping.Ebs.SetVolumeId(vi.VolumeId)
			mapping.Ebs.SetAttachTime(vi.AttachTime)
			mapping.Ebs.SetDeleteOnTermination(vi.DeleteOnTermination)
			mapping.Ebs.SetStatus("attached")
			instance.Instance.BlockDeviceMappings = append(instance.Instance.BlockDeviceMappings, mapping)
		}

		// Launch the instance infrastructure (QEMU, QMP, NATS subscriptions)
		err = d.LaunchInstance(instance)
		if err != nil {
			slog.Error("handleEC2RunInstances LaunchInstance failed", "instanceId", instance.ID, "err", err)
			d.resourceMgr.deallocate(instanceType)
			d.markInstanceFailed(instance, "launch_failed")
			continue
		}

		// Discover actual guest device names via QMP query-block
		d.updateGuestDeviceNames(instance)

		successCount++
		slog.Info("handleEC2RunInstances launched instance", "instanceId", instance.ID)
	}

	slog.Info("handleEC2RunInstances completed", "requested", launchCount, "created", len(instances), "launched", successCount)
}

func (d *Daemon) handleRebootInstance(msg *nats.Msg, command types.EC2InstanceCommand, instance *vm.VM) {
	slog.Info("Rebooting instance", "id", command.ID)

	d.Instances.Mu.Lock()
	status := instance.Status
	d.Instances.Mu.Unlock()

	if status != vm.StateRunning {
		slog.Error("RebootInstance: instance not in running state", "instanceId", command.ID, "status", status)
		respondWithError(msg, awserrors.ErrorIncorrectInstanceState)
		return
	}

	_, err := d.SendQMPCommand(instance.QMPClient, qmp.QMPCommand{Execute: "system_reset"}, command.ID)
	if err != nil {
		slog.Error("RebootInstance: QMP system_reset failed", "instanceId", command.ID, "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	slog.Info("Instance rebooted", "instanceId", command.ID)

	if err := msg.Respond([]byte(`{}`)); err != nil {
		slog.Error("Failed to respond to NATS request", "err", err)
	}
}

func (d *Daemon) handleStartInstance(msg *nats.Msg, command types.EC2InstanceCommand, instance *vm.VM) {
	slog.Info("Starting instance", "id", command.ID)

	// Validate instance is in stopped state
	d.Instances.Mu.Lock()
	status := instance.Status
	d.Instances.Mu.Unlock()

	if status != vm.StateStopped {
		slog.Error("StartInstance: instance not in stopped state", "instanceId", command.ID, "status", status)
		respondWithError(msg, awserrors.ErrorIncorrectInstanceState)
		return
	}

	// Allocate resources
	instanceType, ok := d.resourceMgr.instanceTypes[instance.InstanceType]
	if ok {
		if err := d.resourceMgr.allocate(instanceType); err != nil {
			slog.Error("Failed to allocate resources for start command", "id", command.ID, "err", err)
			respondWithError(msg, awserrors.ErrorInsufficientInstanceCapacity)
			return
		}
	}

	// Clear stop attribute before launch so WriteState inside LaunchInstance
	// persists the correct attributes. Without this, a daemon restart after
	// a stop→start cycle would see StopInstance=true and skip reconnecting QEMU.
	d.Instances.Mu.Lock()
	instance.Attributes = command.Attributes
	d.Instances.Mu.Unlock()

	// Launch the instance infrastructure (QEMU, QMP, NATS subscriptions)
	err := d.LaunchInstance(instance)

	if err != nil {
		slog.Error("handleEC2RunInstances LaunchInstance failed", "err", err)
		// Free the resource on failure
		if ok {
			d.resourceMgr.deallocate(instanceType)
		}
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	// Discover actual guest device names via QMP query-block
	d.updateGuestDeviceNames(instance)

	slog.Info("Instance started", "instanceId", instance.ID)

	if err := msg.Respond(fmt.Appendf(nil, `{"status":"running","instanceId":"%s"}`, instance.ID)); err != nil {
		slog.Error("Failed to respond to NATS request", "err", err)
	}
}

func (d *Daemon) handleStopOrTerminateInstance(msg *nats.Msg, command types.EC2InstanceCommand, instance *vm.VM) {
	isTerminate := command.Attributes.TerminateInstance
	action := "Stopping"
	initialState := vm.StateStopping
	finalState := vm.StateStopped
	if isTerminate {
		action = "Terminating"
		initialState = vm.StateShuttingDown
		finalState = vm.StateTerminated
	}

	slog.Info(action+" instance", "id", command.ID)

	// Check state validity before attempting transition — return the correct
	// AWS error code when the instance is already stopped/terminated/etc.
	d.Instances.Mu.Lock()
	currentState := instance.Status
	d.Instances.Mu.Unlock()
	if !vm.IsValidTransition(currentState, initialState) {
		slog.Warn("Instance in incorrect state for "+strings.ToLower(action),
			"instanceId", instance.ID, "currentState", string(currentState))
		respondWithError(msg, awserrors.ErrorIncorrectInstanceState)
		return
	}

	// Transition to the initial transitional state
	if err := d.TransitionState(instance, initialState); err != nil {
		slog.Error("Failed to transition to "+string(initialState), "instanceId", instance.ID, "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	// Respond immediately - operation will complete in background
	// stopInstance() handles the QMP shutdown command, so we don't send it here
	if err := msg.Respond([]byte(`{}`)); err != nil {
		slog.Error("Failed to respond to NATS request", "err", err)
	}

	// Run cleanup in goroutine to not block NATS
	go func(inst *vm.VM, attrs types.EC2CommandAttributes) {
		stopErr := d.stopInstance(map[string]*vm.VM{inst.ID: inst}, isTerminate)

		if stopErr != nil {
			slog.Error("Failed to "+strings.ToLower(action)+" instance", "err", stopErr, "id", inst.ID)
			if err := d.TransitionState(inst, vm.StateError); err != nil {
				slog.Error("Failed to transition to error state", "instanceId", inst.ID, "err", err)
			}
		} else {
			d.Instances.Mu.Lock()
			inst.Attributes = attrs
			inst.LastNode = d.node
			d.Instances.Mu.Unlock()

			if err := d.TransitionState(inst, finalState); err != nil {
				slog.Error("Failed to transition to final state", "instanceId", inst.ID, "err", err)
			}
			slog.Info("Instance "+string(finalState), "id", inst.ID)

			if d.jsManager != nil {
				if isTerminate {
					// Write to terminated KV bucket (auto-expires after 1 hour via TTL).
					// If this fails, keep the instance in local state so DescribeInstances
					// still sees it and restoreInstances can retry the KV migration.
					if err := d.jsManager.WriteTerminatedInstance(inst.ID, inst); err != nil {
						slog.Error("Failed to write terminated instance to KV, keeping in local state for retry",
							"instanceId", inst.ID, "err", err)
						return
					}
				} else {
					// Write to shared KV first — if daemon crashes after this but
					// before local cleanup, restoreInstances handles the overlap.
					if err := d.jsManager.WriteStoppedInstance(inst.ID, inst); err != nil {
						slog.Error("Failed to write stopped instance to shared KV, keeping local ownership",
							"instanceId", inst.ID, "err", err)
						return
					}
				}

				// Guard + delete must be atomic under the same lock hold.
				// A concurrent ec2.start handler may have loaded the instance
				// from stopped KV, re-added it to VMS with a new pointer, and
				// launched it. Deleting here would destroy the running instance's
				// state — creating a "ghost instance" visible nowhere.
				d.Instances.Mu.Lock()
				current, exists := d.Instances.VMS[inst.ID]
				if !exists || current != inst {
					d.Instances.Mu.Unlock()
					slog.Info("Instance was reclaimed by another handler, skipping local cleanup",
						"instanceId", inst.ID, "state", string(finalState))
					return
				}
				delete(d.Instances.VMS, inst.ID)
				d.Instances.Mu.Unlock()

				// Unsubscribe from per-instance NATS topic. Safe to do after
				// the delete — LaunchInstance already unsubscribes stale entries
				// before creating new ones (daemon.go:1658-1664).
				d.mu.Lock()
				if sub, ok := d.natsSubscriptions[inst.ID]; ok {
					if err := sub.Unsubscribe(); err != nil {
						slog.Error("Failed to unsubscribe instance", "instanceId", inst.ID, "err", err)
					}
					delete(d.natsSubscriptions, inst.ID)
				}
				d.mu.Unlock()

				// Persist local state without the instance
				if err := d.WriteState(); err != nil {
					slog.Error("Failed to persist state after releasing instance, re-adding to local map for consistency",
						"instanceId", inst.ID, "err", err)
					// Only re-add if another handler hasn't claimed the slot
					d.Instances.Mu.Lock()
					if _, occupied := d.Instances.VMS[inst.ID]; !occupied {
						d.Instances.VMS[inst.ID] = inst
					}
					d.Instances.Mu.Unlock()
				} else {
					slog.Info("Released instance ownership to KV",
						"instanceId", inst.ID, "state", string(finalState), "lastNode", d.node)
				}
			}
		}
	}(instance, command.Attributes)
}

// handleEC2DescribeInstances processes incoming EC2 DescribeInstances requests
// This handler responds with instances running on this node owned by the caller's account
func (d *Daemon) handleEC2DescribeInstances(msg *nats.Msg) {
	slog.Debug("Received message", "subject", msg.Subject, "data", string(msg.Data))

	// Extract account ID from NATS header for scoping
	accountID := utils.AccountIDFromMsg(msg)

	// Initialize describeInstancesInput before unmarshaling into it
	describeInstancesInput := &ec2.DescribeInstancesInput{}
	errResp := utils.UnmarshalJsonPayload(describeInstancesInput, msg.Data)

	if errResp != nil {
		if err := msg.Respond(errResp); err != nil {
			slog.Error("Failed to respond to NATS request", "err", err)
		}
		slog.Error("Request does not match DescribeInstancesInput")
		return
	}

	slog.Info("Processing DescribeInstances request from this node", "accountID", accountID)

	d.Instances.Mu.Lock()
	defer d.Instances.Mu.Unlock()

	// Validate and filter instances if specific instance IDs were requested
	instanceIDFilter := make(map[string]bool)
	if len(describeInstancesInput.InstanceIds) > 0 {
		for _, id := range describeInstancesInput.InstanceIds {
			if id != nil && *id != "" {
				if !strings.HasPrefix(*id, "i-") {
					respondWithError(msg, awserrors.ErrorInvalidInstanceIDMalformed)
					return
				}
				instanceIDFilter[*id] = true
			}
		}
	}

	// Group instances by reservation ID (AWS returns instances grouped by reservation)
	reservationMap := make(map[string]*ec2.Reservation)

	// Iterate through all instances on this node
	for _, instance := range d.Instances.VMS {
		// Skip instances not owned by the caller's account.
		// Pre-Phase4 instances (empty AccountID) are only visible to root.
		if !isInstanceVisible(accountID, instance.AccountID) {
			continue
		}

		// Skip if filtering by instance IDs and this instance is not in the filter
		if len(instanceIDFilter) > 0 && !instanceIDFilter[instance.ID] {
			continue
		}

		// Use stored reservation metadata if available
		if instance.Reservation != nil && instance.Instance != nil {
			resID := ""
			if instance.Reservation.ReservationId != nil {
				resID = *instance.Reservation.ReservationId
			}

			// Create reservation entry if it doesn't exist
			if _, exists := reservationMap[resID]; !exists {
				reservation := &ec2.Reservation{}
				reservation.SetReservationId(resID)
				if instance.Reservation.OwnerId != nil {
					reservation.SetOwnerId(*instance.Reservation.OwnerId)
				}
				reservation.Instances = []*ec2.Instance{}
				reservationMap[resID] = reservation
			}

			// Update the instance state to current state
			instanceCopy := *instance.Instance
			instanceCopy.State = &ec2.InstanceState{}

			// Map internal status to EC2 state codes using the centralized mapping
			if info, ok := vm.EC2StateCodes[instance.Status]; ok {
				instanceCopy.State.SetCode(info.Code)
				instanceCopy.State.SetName(info.Name)
			} else {
				slog.Warn("Instance has unmapped status, reporting as pending",
					"instanceId", instance.ID, "status", string(instance.Status))
				instanceCopy.State.SetCode(0)
				instanceCopy.State.SetName("pending")
			}

			// Add instance to its reservation
			reservationMap[resID].Instances = append(reservationMap[resID].Instances, &instanceCopy)
		}
	}

	// Convert map to slice
	reservations := make([]*ec2.Reservation, 0, len(reservationMap))
	for _, reservation := range reservationMap {
		reservations = append(reservations, reservation)
	}

	// Create the response
	output := &ec2.DescribeInstancesOutput{
		Reservations: reservations,
	}

	respondWithJSON(msg, output)
	slog.Info("handleEC2DescribeInstances completed", "count", len(reservations))
}

// handleEC2DescribeInstanceTypes processes incoming EC2 DescribeInstanceTypes requests
// This handler responds with instance types that can currently be provisioned on this node
// based on available resources (CPU and memory not already allocated to running instances)
func (d *Daemon) handleEC2DescribeInstanceTypes(msg *nats.Msg) {
	slog.Debug("Received message", "subject", msg.Subject)

	// Initialize input
	describeInput := &ec2.DescribeInstanceTypesInput{}
	errResp := utils.UnmarshalJsonPayload(describeInput, msg.Data)
	if errResp != nil {
		if err := msg.Respond(errResp); err != nil {
			slog.Error("Failed to respond to NATS request", "err", err)
		}
		slog.Error("Request does not match DescribeInstanceTypesInput")
		return
	}

	slog.Info("Processing DescribeInstanceTypes request from this node")

	// Check if "capacity" filter is set to "true"
	showCapacity := false
	for _, f := range describeInput.Filters {
		if f.Name != nil && *f.Name == "capacity" {
			for _, v := range f.Values {
				if v != nil && *v == "true" {
					showCapacity = true
					break
				}
			}
		}
	}

	// Get instance types based on capacity and the showCapacity flag
	filteredTypes := d.resourceMgr.GetAvailableInstanceTypeInfos(showCapacity)

	// Create the response
	output := &ec2.DescribeInstanceTypesOutput{
		InstanceTypes: filteredTypes,
	}

	respondWithJSON(msg, output)
	slog.Info("handleEC2DescribeInstanceTypes completed", "count", len(filteredTypes))
}

// startStoppedInstanceRequest is the payload for ec2.start topic
type startStoppedInstanceRequest struct {
	InstanceID string `json:"instance_id"`
}

// handleEC2StartStoppedInstance picks up a stopped instance from shared KV,
// re-launches it on this daemon node, and removes it from shared KV.
func (d *Daemon) handleEC2StartStoppedInstance(msg *nats.Msg) {
	var req startStoppedInstanceRequest
	if err := json.Unmarshal(msg.Data, &req); err != nil {
		slog.Error("handleEC2StartStoppedInstance: failed to unmarshal request", "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	if req.InstanceID == "" {
		slog.Error("handleEC2StartStoppedInstance: missing instance_id")
		respondWithError(msg, awserrors.ErrorMissingParameter)
		return
	}

	if d.jsManager == nil {
		slog.Error("handleEC2StartStoppedInstance: JetStream not available")
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	// Load instance from shared KV
	instance, err := d.jsManager.LoadStoppedInstance(req.InstanceID)
	if err != nil {
		slog.Error("handleEC2StartStoppedInstance: failed to load stopped instance", "instanceId", req.InstanceID, "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}
	if instance == nil {
		slog.Warn("handleEC2StartStoppedInstance: instance not found in shared KV", "instanceId", req.InstanceID)
		respondWithError(msg, awserrors.ErrorInvalidInstanceIDNotFound)
		return
	}

	if instance.Status != vm.StateStopped {
		slog.Error("handleEC2StartStoppedInstance: instance not in stopped state", "instanceId", req.InstanceID, "status", instance.Status)
		respondWithError(msg, awserrors.ErrorIncorrectInstanceState)
		return
	}

	// Verify the caller owns this instance
	if !checkInstanceOwnership(msg, req.InstanceID, instance.AccountID) {
		return
	}

	// Reset node-local fields that are stale after cross-node migration
	instance.ResetNodeLocalState()

	// Allocate resources
	instanceType, ok := d.resourceMgr.instanceTypes[instance.InstanceType]
	if !ok {
		slog.Error("handleEC2StartStoppedInstance: instance type not available on this node",
			"instanceId", req.InstanceID, "instanceType", instance.InstanceType)
		respondWithError(msg, awserrors.ErrorInsufficientInstanceCapacity)
		return
	}
	if err := d.resourceMgr.allocate(instanceType); err != nil {
		slog.Error("handleEC2StartStoppedInstance: failed to allocate resources", "instanceId", req.InstanceID, "err", err)
		respondWithError(msg, awserrors.ErrorInsufficientInstanceCapacity)
		return
	}

	// Add instance to local map and clear stop attribute before launch
	d.Instances.Mu.Lock()
	d.Instances.VMS[instance.ID] = instance
	instance.Attributes = types.EC2CommandAttributes{StartInstance: true}
	d.Instances.Mu.Unlock()

	// Launch the instance infrastructure (QEMU, QMP, NATS subscriptions)
	err = d.LaunchInstance(instance)
	if err != nil {
		slog.Error("handleEC2StartStoppedInstance: LaunchInstance failed", "instanceId", req.InstanceID, "err", err)
		// Rollback: deallocate resources and remove from local map
		if ok {
			d.resourceMgr.deallocate(instanceType)
		}
		d.Instances.Mu.Lock()
		delete(d.Instances.VMS, instance.ID)
		d.Instances.Mu.Unlock()
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	// Discover actual guest device names via QMP query-block
	d.updateGuestDeviceNames(instance)

	// Remove from shared KV now that it's running locally.
	// Retry once on failure — a stale KV entry risks duplicate starts.
	if err := d.jsManager.DeleteStoppedInstance(req.InstanceID); err != nil {
		slog.Warn("handleEC2StartStoppedInstance: first KV delete failed, retrying",
			"instanceId", req.InstanceID, "err", err)
		if retryErr := d.jsManager.DeleteStoppedInstance(req.InstanceID); retryErr != nil {
			slog.Error("handleEC2StartStoppedInstance: KV delete failed after retry, instance is running locally but stale entry remains in shared KV",
				"instanceId", req.InstanceID, "err", retryErr)
		}
	}

	slog.Info("Started stopped instance from shared KV", "instanceId", instance.ID, "node", d.node)

	if err := msg.Respond(fmt.Appendf(nil, `{"status":"running","instanceId":"%s"}`, instance.ID)); err != nil {
		slog.Error("Failed to respond to NATS request", "err", err)
	}
}

// terminateStoppedInstanceRequest is the payload for ec2.terminate topic
type terminateStoppedInstanceRequest struct {
	InstanceID string `json:"instance_id"`
}

// handleEC2TerminateStoppedInstance picks up a stopped instance from shared KV,
// deletes its volumes, and removes it from shared KV.
func (d *Daemon) handleEC2TerminateStoppedInstance(msg *nats.Msg) {
	var req terminateStoppedInstanceRequest
	if err := json.Unmarshal(msg.Data, &req); err != nil {
		slog.Error("handleEC2TerminateStoppedInstance: failed to unmarshal request", "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	if req.InstanceID == "" {
		slog.Error("handleEC2TerminateStoppedInstance: missing instance_id")
		respondWithError(msg, awserrors.ErrorMissingParameter)
		return
	}

	if d.jsManager == nil {
		slog.Error("handleEC2TerminateStoppedInstance: JetStream not available")
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	// Load instance from shared KV
	instance, err := d.jsManager.LoadStoppedInstance(req.InstanceID)
	if err != nil {
		slog.Error("handleEC2TerminateStoppedInstance: failed to load stopped instance", "instanceId", req.InstanceID, "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}
	if instance == nil {
		slog.Warn("handleEC2TerminateStoppedInstance: instance not found in shared KV", "instanceId", req.InstanceID)
		respondWithError(msg, awserrors.ErrorInvalidInstanceIDNotFound)
		return
	}

	if instance.Status != vm.StateStopped {
		slog.Error("handleEC2TerminateStoppedInstance: instance not in stopped state", "instanceId", req.InstanceID, "status", instance.Status)
		respondWithError(msg, awserrors.ErrorIncorrectInstanceState)
		return
	}

	// Verify the caller owns this instance
	if !checkInstanceOwnership(msg, req.InstanceID, instance.AccountID) {
		return
	}

	// Delete volumes — no QEMU shutdown or unmount needed (already done during stop)
	instance.EBSRequests.Mu.Lock()
	for _, ebsRequest := range instance.EBSRequests.Requests {
		// Internal volumes (EFI, cloud-init) are always cleaned up via ebs.delete
		if ebsRequest.EFI || ebsRequest.CloudInit {
			ebsDeleteData, err := json.Marshal(types.EBSDeleteRequest{Volume: ebsRequest.Name})
			if err != nil {
				slog.Error("handleEC2TerminateStoppedInstance: failed to marshal ebs.delete request", "name", ebsRequest.Name, "err", err)
				continue
			}
			deleteMsg, err := d.natsConn.Request("ebs.delete", ebsDeleteData, 30*time.Second)
			if err != nil {
				slog.Warn("handleEC2TerminateStoppedInstance: ebs.delete failed for internal volume", "name", ebsRequest.Name, "err", err)
			} else {
				slog.Info("handleEC2TerminateStoppedInstance: ebs.delete sent for internal volume", "name", ebsRequest.Name, "data", string(deleteMsg.Data))
			}
			continue
		}

		// User-visible volumes: respect DeleteOnTermination flag
		if !ebsRequest.DeleteOnTermination {
			slog.Info("handleEC2TerminateStoppedInstance: volume has DeleteOnTermination=false, skipping", "name", ebsRequest.Name)
			continue
		}

		slog.Info("handleEC2TerminateStoppedInstance: deleting volume with DeleteOnTermination=true", "name", ebsRequest.Name)
		_, err := d.volumeService.DeleteVolume(&ec2.DeleteVolumeInput{
			VolumeId: &ebsRequest.Name,
		}, instance.AccountID)
		if err != nil {
			slog.Error("handleEC2TerminateStoppedInstance: failed to delete volume", "name", ebsRequest.Name, "err", err)
		}
	}
	instance.EBSRequests.Mu.Unlock()

	// Write to terminated KV bucket FIRST so the instance is visible in DescribeInstances.
	// If this fails, the instance remains in the stopped bucket (safe to retry).
	instance.Status = vm.StateTerminated
	if err := d.jsManager.WriteTerminatedInstance(req.InstanceID, instance); err != nil {
		slog.Error("handleEC2TerminateStoppedInstance: failed to write to terminated KV, aborting", "instanceId", req.InstanceID, "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	// Now safe to remove from shared stopped KV — instance already exists in terminated bucket.
	// Retry once on failure to avoid duplicate entries in DescribeInstances.
	if err := d.jsManager.DeleteStoppedInstance(req.InstanceID); err != nil {
		slog.Warn("handleEC2TerminateStoppedInstance: first stopped KV delete failed, retrying",
			"instanceId", req.InstanceID, "err", err)
		if retryErr := d.jsManager.DeleteStoppedInstance(req.InstanceID); retryErr != nil {
			slog.Error("handleEC2TerminateStoppedInstance: stopped KV delete failed after retry, instance may appear in both buckets",
				"instanceId", req.InstanceID, "err", retryErr)
		}
	}

	slog.Info("Terminated stopped instance from shared KV", "instanceId", req.InstanceID)

	if err := msg.Respond(fmt.Appendf(nil, `{"status":"terminated","instanceId":"%s"}`, req.InstanceID)); err != nil {
		slog.Error("Failed to respond to NATS request", "err", err)
	}
}

// handleEC2DescribeStoppedInstances returns stopped instances from shared KV.
func (d *Daemon) handleEC2DescribeStoppedInstances(msg *nats.Msg) {
	d.describeInstancesFromKV(msg, d.jsManager.ListStoppedInstances, 80, "stopped", "handleEC2DescribeStoppedInstances")
}

// handleEC2DescribeTerminatedInstances returns terminated instances from the terminated KV bucket.
func (d *Daemon) handleEC2DescribeTerminatedInstances(msg *nats.Msg) {
	d.describeInstancesFromKV(msg, d.jsManager.ListTerminatedInstances, 48, "terminated", "handleEC2DescribeTerminatedInstances")
}

// describeInstancesFromKV is a shared helper for DescribeStopped/TerminatedInstances handlers.
// It lists instances from a KV bucket, filters by account/instance ID, and responds with reservations.
func (d *Daemon) describeInstancesFromKV(msg *nats.Msg, listFn func() ([]*vm.VM, error), fallbackCode int64, fallbackName, handlerName string) {
	if d.jsManager == nil {
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	accountID := utils.AccountIDFromMsg(msg)

	describeInput := &ec2.DescribeInstancesInput{}
	if len(msg.Data) > 0 {
		if errResp := utils.UnmarshalJsonPayload(describeInput, msg.Data); errResp != nil {
			if err := msg.Respond(errResp); err != nil {
				slog.Error("Failed to respond to NATS request", "err", err)
			}
			return
		}
	}

	instanceIDFilter := make(map[string]bool)
	for _, id := range describeInput.InstanceIds {
		if id != nil {
			instanceIDFilter[*id] = true
		}
	}

	instances, err := listFn()
	if err != nil {
		slog.Error(handlerName+": failed to list instances", "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	reservationMap := make(map[string]*ec2.Reservation)

	for _, instance := range instances {
		if !isInstanceVisible(accountID, instance.AccountID) {
			continue
		}
		if len(instanceIDFilter) > 0 && !instanceIDFilter[instance.ID] {
			continue
		}
		if instance.Reservation == nil || instance.Instance == nil {
			slog.Warn(handlerName+": skipping instance with nil Reservation/Instance (data integrity issue)",
				"instanceId", instance.ID)
			continue
		}

		resID := ""
		if instance.Reservation.ReservationId != nil {
			resID = *instance.Reservation.ReservationId
		}

		if _, exists := reservationMap[resID]; !exists {
			reservation := &ec2.Reservation{}
			reservation.SetReservationId(resID)
			if instance.Reservation.OwnerId != nil {
				reservation.SetOwnerId(*instance.Reservation.OwnerId)
			}
			reservation.Instances = []*ec2.Instance{}
			reservationMap[resID] = reservation
		}

		instanceCopy := *instance.Instance
		instanceCopy.State = &ec2.InstanceState{}
		if info, ok := vm.EC2StateCodes[instance.Status]; ok {
			instanceCopy.State.SetCode(info.Code)
			instanceCopy.State.SetName(info.Name)
		} else {
			instanceCopy.State.SetCode(fallbackCode)
			instanceCopy.State.SetName(fallbackName)
		}

		reservationMap[resID].Instances = append(reservationMap[resID].Instances, &instanceCopy)
	}

	reservations := make([]*ec2.Reservation, 0, len(reservationMap))
	for _, reservation := range reservationMap {
		reservations = append(reservations, reservation)
	}

	respondWithJSON(msg, &ec2.DescribeInstancesOutput{Reservations: reservations})
	slog.Info(handlerName+" completed", "count", len(reservations))
}

// handleEC2ModifyInstanceAttribute modifies attributes of a stopped instance in shared KV.
// All supported attributes (InstanceType, UserData) require the instance to be stopped.
func (d *Daemon) handleEC2ModifyInstanceAttribute(msg *nats.Msg) {
	var input ec2.ModifyInstanceAttributeInput
	if err := json.Unmarshal(msg.Data, &input); err != nil {
		slog.Error("handleEC2ModifyInstanceAttribute: failed to unmarshal request", "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	if input.InstanceId == nil || *input.InstanceId == "" {
		slog.Error("handleEC2ModifyInstanceAttribute: missing instance_id")
		respondWithError(msg, awserrors.ErrorMissingParameter)
		return
	}

	instanceID := *input.InstanceId

	if d.jsManager == nil {
		slog.Error("handleEC2ModifyInstanceAttribute: JetStream not available")
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	instance, err := d.jsManager.LoadStoppedInstance(instanceID)
	if err != nil {
		slog.Error("handleEC2ModifyInstanceAttribute: failed to load stopped instance", "instanceId", instanceID, "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}
	if instance == nil {
		slog.Warn("handleEC2ModifyInstanceAttribute: instance not found in shared KV", "instanceId", instanceID)
		respondWithError(msg, awserrors.ErrorInvalidInstanceIDNotFound)
		return
	}

	if instance.Status != vm.StateStopped {
		slog.Error("handleEC2ModifyInstanceAttribute: instance not in stopped state", "instanceId", instanceID, "status", instance.Status)
		respondWithError(msg, awserrors.ErrorIncorrectInstanceState)
		return
	}

	// Verify the caller owns this instance
	if !checkInstanceOwnership(msg, instanceID, instance.AccountID) {
		return
	}

	// Apply the requested attribute change
	if input.InstanceType != nil && input.InstanceType.Value != nil {
		newType := *input.InstanceType.Value
		if newType == "" {
			slog.Error("handleEC2ModifyInstanceAttribute: empty instance type value", "instanceId", instanceID)
			respondWithError(msg, awserrors.ErrorInvalidInstanceAttributeValue)
			return
		}
		if instance.Instance == nil {
			slog.Error("handleEC2ModifyInstanceAttribute: instance.Instance is nil, data integrity issue", "instanceId", instanceID)
			respondWithError(msg, awserrors.ErrorServerInternal)
			return
		}
		slog.Info("handleEC2ModifyInstanceAttribute: changing instance type",
			"instanceId", instanceID, "oldType", instance.InstanceType, "newType", newType)

		instance.InstanceType = newType
		instance.Config.InstanceType = newType
		instance.Instance.InstanceType = aws.String(newType)
		// Clear StateReason — resolves capacity-unavailable state from instance-type-missing bug
		instance.Instance.StateReason = nil
	}

	if input.UserData != nil && input.UserData.Value != nil {
		slog.Info("handleEC2ModifyInstanceAttribute: changing user data", "instanceId", instanceID)

		// Value arrives as decoded bytes (JSON unmarshal handles base64 → []byte automatically)
		instance.UserData = string(input.UserData.Value)
		if instance.RunInstancesInput != nil {
			instance.RunInstancesInput.UserData = aws.String(base64.StdEncoding.EncodeToString(input.UserData.Value))
		}
	}

	if err := d.jsManager.WriteStoppedInstance(instanceID, instance); err != nil {
		slog.Error("handleEC2ModifyInstanceAttribute: failed to write modified instance to KV",
			"instanceId", instanceID, "err", err)
		respondWithError(msg, awserrors.ErrorServerInternal)
		return
	}

	slog.Info("handleEC2ModifyInstanceAttribute: completed successfully", "instanceId", instanceID)

	if err := msg.Respond([]byte(`{}`)); err != nil {
		slog.Error("Failed to respond to NATS request", "err", err)
	}
}