//go:build linux

package microvm

import (
	"context"
	"crypto/rand"
	"errors"
	"log/slog"
	"os"
	"os/exec"
	"testing"

	"tangled.org/core/spindle/models"
)

type shutdownTestVM struct {
	exited      bool
	waitErr     error
	shutdownErr error
	closed      bool
}

func (v *shutdownTestVM) Shutdown(context.Context) error {
	v.exited = true
	return v.shutdownErr
}

func (v *shutdownTestVM) WaitContext(ctx context.Context) error {
	if v.exited {
		return v.waitErr
	}
	return ctx.Err()
}

func (v *shutdownTestVM) Close() error {
	v.closed = true
	return nil
}

func (*shutdownTestVM) Logs() VMLogs    { return VMLogs{} }
func (*shutdownTestVM) CID() uint32     { return 0 }
func (*shutdownTestVM) WorkDir() string { return "" }
func (*shutdownTestVM) OOMKilled() bool { return false }

func TestShutdownVM_TreatsExitedVMAsCleanedUp(t *testing.T) {
	for name, vm := range map[string]*shutdownTestVM{
		"already exited": {
			exited:      true,
			waitErr:     errors.New("qemu exited"),
			shutdownErr: errors.New("qmp broken pipe"),
		},
		"exits during fallback": {
			waitErr:     errors.New("qemu exited"),
			shutdownErr: errors.New("qmp broken pipe"),
		},
	} {
		t.Run(name, func(t *testing.T) {
			e := &Engine{l: slog.Default()}
			state := &workflowState{VM: vm}

			if err := e.shutdownVM(context.Background(), models.WorkflowId{}, state); err != nil {
				t.Fatalf("shutdownVM: %v", err)
			}
			if !vm.closed {
				t.Fatal("expected vm handle to be closed")
			}
		})
	}
}

func TestVolumeUsage(t *testing.T) {
	tmpFile, err := os.CreateTemp("", "spindle-test-volume")
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(tmpFile.Name())
	defer tmpFile.Close()

	data := make([]byte, 1024*1024)
	if _, err := rand.Read(data); err != nil {
		t.Fatal(err)
	}
	if _, err := tmpFile.Write(data); err != nil {
		t.Fatal(err)
	}
	if err := tmpFile.Sync(); err != nil {
		t.Fatal(err)
	}

	h := &QEMUVMHandle{
		volumePaths: map[string]string{
			"vol1": tmpFile.Name(),
		},
	}

	usages, err := h.VolumeUsage()
	if err != nil {
		t.Fatal(err)
	}

	allocated, ok := usages[tmpFile.Name()]
	if !ok {
		t.Fatalf("expected path %q in usages, got %+v", tmpFile.Name(), usages)
	}

	if allocated < 1024*1024 {
		t.Errorf("expected at least 1MB allocated, got %d", allocated)
	}
}

func TestQEMUCloseSnapshotsAfterProcessExit(t *testing.T) {
	cmd := exec.Command("sleep", "30")
	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}
	h := &QEMUVMHandle{
		Process: cmd.Process,
		done:    make(chan struct{}),
	}
	go func() {
		err := cmd.Wait()
		h.waitErrMu.Lock()
		h.waitErr = err
		h.waitErrMu.Unlock()
		close(h.done)
	}()

	var waitedDuringSnapshot bool
	if err := h.close(func() {
		select {
		case <-h.done:
			waitedDuringSnapshot = true
		default:
		}
	}); err != nil {
		t.Fatal(err)
	}
	if !waitedDuringSnapshot {
		t.Fatal("resource snapshot ran before QEMU wait completed")
	}
}

func TestCaptureResourceUsageMarksPartialVolumeSnapshotUnavailable(t *testing.T) {
	tmpFile, err := os.CreateTemp("", "spindle-test-volume")
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(tmpFile.Name())
	defer tmpFile.Close()
	if _, err := tmpFile.Write(make([]byte, 4096)); err != nil {
		t.Fatal(err)
	}
	if err := tmpFile.Sync(); err != nil {
		t.Fatal(err)
	}

	state := &workflowState{VM: &QEMUVMHandle{
		volumePaths: map[string]string{
			"present": tmpFile.Name(),
			"missing": tmpFile.Name() + "-missing",
		},
	}}
	e := &Engine{l: slog.New(slog.NewTextHandler(os.Stderr, nil))}
	e.captureResourceUsage(models.WorkflowId{}, state)

	if !state.ResourceUsageAvailable || state.ResourceUsage == nil {
		t.Fatal("partial resource snapshot was discarded")
	}
	if state.ResourceUsage.VolumeAvailable {
		t.Fatal("partial volume snapshot marked available")
	}
	if state.ResourceUsage.VolumeAllocatedBytes == 0 {
		t.Fatal("partial volume snapshot lost successful measurements")
	}
}

func TestCIDRegistryAvoidsCollisionsAndReusesReleasedCID(t *testing.T) {
	registry := newCIDRegistry()
	first, err := registry.allocate()
	if err != nil {
		t.Fatal(err)
	}
	second, err := registry.allocate()
	if err != nil {
		t.Fatal(err)
	}
	if first == second {
		t.Fatalf("duplicate allocated CID %d", first)
	}
	registry.release(first)
	registry.mu.Lock()
	registry.next = first
	registry.mu.Unlock()
	reused, err := registry.allocate()
	if err != nil {
		t.Fatal(err)
	}
	if reused != first {
		t.Fatalf("released CID: got %d, want %d", reused, first)
	}
}