package microvm

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"math"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync/atomic"
	"time"

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

const (
	minGuestCID         = 3
	vmCrashLogTailBytes = 4096
)

func AllocateCID() (uint32, error) {
	var data [4]byte
	if _, err := rand.Read(data[:]); err != nil {
		return 0, fmt.Errorf("allocate guest CID: %w", err)
	}
	return minGuestCID + binary.BigEndian.Uint32(data[:])%60000, nil
}

func prepareWorkDir(workDir string) (string, bool, error) {
	if workDir == "" {
		dir, err := os.MkdirTemp("", "spindle-microvm-*")
		if err != nil {
			return "", false, fmt.Errorf("create microvm work directory: %w", err)
		}
		return dir, true, nil
	}
	if err := os.MkdirAll(workDir, 0o755); err != nil {
		return "", false, fmt.Errorf("create microvm work directory: %w", err)
	}
	return workDir, false, nil
}

func prepareVolumes(ctx context.Context, workDir string, volumes []Volume, mkfsExt4 string) (map[string]string, error) {
	paths := make(map[string]string, len(volumes))
	for _, volume := range volumes {
		if volume.ReadOnly {
			return nil, fmt.Errorf("read-only microvm volume %q is not supported yet", volume.Image)
		}
		if volume.FSType != "ext4" {
			return nil, fmt.Errorf("microvm volume %q uses unsupported fsType %q", volume.Image, volume.FSType)
		}
		if volume.ImageType != "" && volume.ImageType != "raw" {
			return nil, fmt.Errorf("microvm volume %q uses unsupported imageType %q", volume.Image, volume.ImageType)
		}

		path := filepath.Join(workDir, filepath.Base(volume.Image))
		if err := createSparseFile(path, volume.SizeMiB); err != nil {
			return nil, err
		}
		noJournal := volume.MountPoint == "/workspace"
		if err := runMkfsExt4(ctx, mkfsExt4, path, noJournal); err != nil {
			return nil, err
		}
		paths[volume.Image] = path
	}
	return paths, nil
}

func createSparseFile(path string, sizeMiB int64) error {
	if sizeMiB <= 0 {
		return fmt.Errorf("sparse file %q size must be positive", path)
	}
	if sizeMiB > math.MaxInt64/(1024*1024) {
		return fmt.Errorf("sparse file %q size is too large", path)
	}
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return fmt.Errorf("create sparse file %q: %w", path, err)
	}
	defer file.Close()

	if err := file.Truncate(sizeMiB * 1024 * 1024); err != nil {
		return fmt.Errorf("resize sparse file %q: %w", path, err)
	}
	return nil
}

func runMkfsExt4(ctx context.Context, mkfsExt4, path string, noJournal bool) error {
	if mkfsExt4 == "" {
		return fmt.Errorf("mkfs.ext4 path is required")
	}
	args := []string{"-F"}
	if noJournal {
		args = append(args, "-O", "^has_journal")
	}
	args = append(args, path)

	cmd := exec.CommandContext(ctx, mkfsExt4, args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("mkfs.ext4 %q: %w: %s", path, err, strings.TrimSpace(string(output)))
	}
	return nil
}

func createParentedFile(path string) (*os.File, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return nil, fmt.Errorf("create log directory: %w", err)
	}
	file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		return nil, fmt.Errorf("create log file %q: %w", path, err)
	}
	return file, nil
}

func validateHostDevices(enableKVM bool, spec ImageSpec) error {
	if _, err := os.Stat("/dev/vhost-vsock"); err != nil {
		return fmt.Errorf("microvm requires /dev/vhost-vsock for vhost-vsock-device: %w", err)
	}
	if len(spec.NetworkInterfaces) > 0 {
		if _, err := os.Stat("/dev/net/tun"); err != nil {
			return fmt.Errorf("microvm slirp4netns networking requires /dev/net/tun: %w", err)
		}
	}
	if enableKVM {
		if _, err := os.Stat("/dev/kvm"); err != nil {
			return fmt.Errorf("microvm KVM was requested but /dev/kvm is not accessible: %w", err)
		}
	}
	return nil
}

type VMLogs struct {
	Serial string
	Extra  map[string]string
}

type VMHandle interface {
	Shutdown(ctx context.Context) error
	WaitContext(ctx context.Context) error
	Close() error
	Logs() VMLogs
	CID() uint32
	WorkDir() string
	OOMKilled() bool
}

type QEMUVMConfig struct {
	Binary        string
	LogPath       string
	QMPPath       string
	SerialLogPath string
}

type VMConfig struct {
	Image       ImageSpec
	CID         uint32
	EnableKVM   bool
	KeepWorkDir bool
	WorkDir     string
	Cgroup      CgroupLimits

	QEMU QEMUVMConfig

	BootTimeout time.Duration
	MkfsExt4    string
}

type workflowState struct {
	ImageSpec              ImageSpec
	ImageSpecPath          string
	Config                 manifestConfig
	ConfigKey              string
	Image                  string
	CacheReadURLs          []string
	CacheTrustedPublicKeys []string
	VM                     VMHandle
	Agent                  *AgentSession
	ReadCache              *ReadCacheProxy
	UploadCache            *UploadCacheProxy
	WorkDir                string
	NixOSToplevelCache     nixosToplevelCacheStore
}

func (e *Engine) cleanupState(ctx context.Context, wid models.WorkflowId, state *workflowState) error {
	if state == nil {
		return nil
	}

	ctx = context.WithoutCancel(ctx)

	var err error
	err = errors.Join(err, e.drainNixCache(ctx, state))
	err = errors.Join(err, e.shutdownVM(ctx, wid, state))
	err = errors.Join(err, closeAgent(state))
	err = errors.Join(err, closeReadCache(state))
	err = errors.Join(err, closeUploadCache(state))
	err = errors.Join(err, removeWorkDir(state))
	return err
}

func (e *Engine) drainNixCache(ctx context.Context, state *workflowState) error {
	if state.Agent == nil || e.cfg.NixCache.UploadURL == "" {
		return nil
	}

	drainCtx, cancel := context.WithTimeout(ctx, cacheDrainTimeout)
	defer cancel()

	if _, err := state.Agent.Drain(drainCtx); err != nil {
		return fmt.Errorf("drain nix cache: %w", err)
	}
	return nil
}

func (e *Engine) shutdownVM(ctx context.Context, wid models.WorkflowId, state *workflowState) error {
	if state.VM == nil {
		return nil
	}

	var err error

	if state.Agent != nil {
		gracefulCtx, cancel := context.WithTimeout(ctx, vmShutdownTimeout)
		poweredOff, poweroffErr := e.poweroffViaAgent(gracefulCtx, wid, state)
		cancel()

		err = errors.Join(err, poweroffErr)
		if poweredOff {
			return errors.Join(err, closeVM(state))
		}
	}

	fallbackCtx, cancel := context.WithTimeout(ctx, vmShutdownTimeout)
	defer cancel()

	if shutdownErr := state.VM.Shutdown(fallbackCtx); shutdownErr != nil {
		e.l.Warn("microVM shutdown fallback failed", "workflow", wid, "error", shutdownErr)
		err = errors.Join(err, shutdownErr)
	}

	return errors.Join(err, closeVM(state))
}

func (e *Engine) poweroffViaAgent(ctx context.Context, wid models.WorkflowId, state *workflowState) (bool, error) {
	if err := state.Agent.Poweroff(ctx); err != nil {
		e.l.Warn("agent poweroff request failed", "workflow", wid, "error", err)
		return false, err
	}

	if err := state.VM.WaitContext(ctx); err != nil {
		e.l.Warn("agent poweroff did not stop microVM", "workflow", wid, "error", err)
		return false, nil
	}

	return true, nil
}

func closeAgent(state *workflowState) error {
	if state.Agent == nil {
		return nil
	}

	err := state.Agent.Close()
	state.Agent = nil
	return err
}

func closeVM(state *workflowState) error {
	if state.VM == nil {
		return nil
	}

	err := state.VM.Close()
	state.VM = nil
	return err
}

func closeReadCache(state *workflowState) error {
	if state.ReadCache == nil {
		return nil
	}

	err := state.ReadCache.Close()
	state.ReadCache = nil
	return err
}

func closeUploadCache(state *workflowState) error {
	if state.UploadCache == nil {
		return nil
	}

	err := state.UploadCache.Close()
	state.UploadCache = nil
	return err
}

func removeWorkDir(state *workflowState) error {
	if state.WorkDir == "" {
		return nil
	}

	err := os.RemoveAll(state.WorkDir)
	state.WorkDir = ""
	return err
}

// returns a context derived from ctx that is cancelled either when ctx itself
// is cancelled or when the microVM exits on its own. the returned flag reports
// whether the VM exited (as opposed to ctx being cancelled for another reason,
// e.g. the workflow timeout), letting callers tell a crash apart from a
// timeout. cancel must be called to release the watcher goroutine.
func watchVMExit(ctx context.Context, vm VMHandle) (context.Context, *atomic.Bool, context.CancelFunc) {
	exited := &atomic.Bool{}
	watchCtx, cancel := context.WithCancel(ctx)
	if vm == nil {
		return watchCtx, exited, cancel
	}
	go func() {
		_ = vm.WaitContext(watchCtx) // returns when VM exits or watchCtx is cancelled
		if watchCtx.Err() == nil {
			exited.Store(true)
			cancel() // don't forget to cancel the watchCtx...
		}
	}()
	return watchCtx, exited, cancel
}

func vmCrashLog(vm VMHandle) string {
	if vm == nil {
		return ""
	}
	logs := vm.Logs()

	var b strings.Builder
	if tail := tailFile(logs.Serial, vmCrashLogTailBytes); tail != "" {
		fmt.Fprintf(&b, "==== serial log ====\n%s\n", tail)
	}
	if qemuLog := logs.Extra["qemu"]; qemuLog != "" {
		if tail := tailFile(qemuLog, vmCrashLogTailBytes); tail != "" {
			fmt.Fprintf(&b, "==== qemu log ====\n%s\n", tail)
		}
	}
	return strings.TrimRight(b.String(), "\n")
}

func tailFile(path string, max int64) string {
	if path == "" {
		return ""
	}
	f, err := os.Open(path)
	if err != nil {
		return ""
	}
	defer f.Close()
	if info, err := f.Stat(); err == nil && info.Size() > max {
		if _, err := f.Seek(-max, io.SeekEnd); err != nil {
			return ""
		}
	}
	data, err := io.ReadAll(f)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(data))
}

func waitAgentConn(ctx context.Context, connCh <-chan net.Conn) (net.Conn, error) {
	select {
	case conn := <-connCh:
		if conn == nil {
			return nil, fmt.Errorf("agent connection closed before setup")
		}
		return conn, nil
	case <-ctx.Done():
		return nil, fmt.Errorf("waiting for agent: %w", ctx.Err())
	}
}

func StartVM(ctx context.Context, cfg VMConfig, logger *slog.Logger) (VMHandle, error) {
	runnerType := cfg.Image.RunnerType
	if runnerType == "" {
		runnerType = "qemu"
	}
	switch runnerType {
	case "qemu":
		qemuLogPath := cfg.QEMU.LogPath
		if qemuLogPath == "" && cfg.WorkDir != "" {
			qemuLogPath = filepath.Join(cfg.WorkDir, "qemu.log")
		}
		serialLogPath := cfg.QEMU.SerialLogPath
		if serialLogPath == "" && cfg.WorkDir != "" {
			serialLogPath = filepath.Join(cfg.WorkDir, "serial.log")
		}
		bootTimeout := cfg.BootTimeout
		if bootTimeout == 0 {
			bootTimeout = 10 * time.Second
		}
		return StartQEMU(ctx, QEMUConfig{
			Image:         cfg.Image,
			BootTimeout:   bootTimeout,
			CID:           cfg.CID,
			EnableKVM:     cfg.EnableKVM,
			KeepWorkDir:   cfg.KeepWorkDir,
			QEMUBinary:    cfg.QEMU.Binary,
			QEMULogPath:   qemuLogPath,
			QMPPath:       cfg.QEMU.QMPPath,
			SerialLogPath: serialLogPath,
			WorkDir:       cfg.WorkDir,
			MkfsExt4:      cfg.MkfsExt4,
			Cgroup:        cfg.Cgroup,
		}, logger)
	case "firecracker":
		return nil, fmt.Errorf("runner type %q not implemented yet", runnerType)
	default:
		return nil, fmt.Errorf("unsupported runner type %q", runnerType)
	}
}
