// Package dmsnap wraps the host-side device-mapper snapshot operations used
// to give each VM a copy-on-write view over a shared rootfs image. It issues
// losetup/dmsetup via a system.CommandRunner-compatible runner.
package dmsnap

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"
)

// Runner is the narrow command-runner surface dmsnap needs. system.Runner
// satisfies it.
type Runner interface {
	RunSudo(ctx context.Context, args ...string) ([]byte, error)
}

// Handles records the loop devices and dm target allocated for a snapshot.
// Callers pass it back to Cleanup to unwind in the right order.
type Handles struct {
	BaseLoop string
	COWLoop  string
	DMName   string
	DMDev    string
}

// Create sets up a dm-snapshot named dmName layering cowPath over rootfsPath.
// On failure it cleans up whatever it had attached so far.
func Create(ctx context.Context, runner Runner, rootfsPath, cowPath, dmName string) (handles Handles, err error) {
	defer func() {
		if err == nil {
			return
		}
		if cleanupErr := Cleanup(context.Background(), runner, handles); cleanupErr != nil {
			err = errors.Join(err, cleanupErr)
		}
	}()

	baseBytes, err := runner.RunSudo(ctx, "losetup", "-f", "--show", "--read-only", rootfsPath)
	if err != nil {
		return handles, err
	}
	handles.BaseLoop = strings.TrimSpace(string(baseBytes))

	cowBytes, err := runner.RunSudo(ctx, "losetup", "-f", "--show", cowPath)
	if err != nil {
		return handles, err
	}
	handles.COWLoop = strings.TrimSpace(string(cowBytes))

	sectorsBytes, err := runner.RunSudo(ctx, "blockdev", "--getsz", handles.BaseLoop)
	if err != nil {
		return handles, err
	}
	sectors := strings.TrimSpace(string(sectorsBytes))

	if _, err := runner.RunSudo(ctx, "dmsetup", "create", dmName, "--table", fmt.Sprintf("0 %s snapshot %s %s P 8", sectors, handles.BaseLoop, handles.COWLoop)); err != nil {
		return handles, err
	}
	handles.DMName = dmName
	handles.DMDev = "/dev/mapper/" + dmName
	return handles, nil
}

// Cleanup tears down a snapshot: remove the dm target, then detach the loops.
// Missing-handle errors (already cleaned up) are ignored.
func Cleanup(ctx context.Context, runner Runner, handles Handles) error {
	var cleanupErr error

	switch {
	case handles.DMName != "":
		if err := Remove(ctx, runner, handles.DMName); err != nil {
			cleanupErr = errors.Join(cleanupErr, err)
		}
	case handles.DMDev != "":
		if err := Remove(ctx, runner, handles.DMDev); err != nil {
			cleanupErr = errors.Join(cleanupErr, err)
		}
	}

	if handles.COWLoop != "" {
		if _, err := runner.RunSudo(ctx, "losetup", "-d", handles.COWLoop); err != nil {
			if !isMissing(err) {
				cleanupErr = errors.Join(cleanupErr, err)
			}
		}
	}
	if handles.BaseLoop != "" {
		if _, err := runner.RunSudo(ctx, "losetup", "-d", handles.BaseLoop); err != nil {
			if !isMissing(err) {
				cleanupErr = errors.Join(cleanupErr, err)
			}
		}
	}

	return cleanupErr
}

// Remove retries dmsetup remove while the device is briefly busy after
// detach. Missing targets succeed.
func Remove(ctx context.Context, runner Runner, target string) error {
	deadline := time.Now().Add(15 * time.Second)
	for {
		if _, err := runner.RunSudo(ctx, "dmsetup", "remove", target); err != nil {
			if isMissing(err) {
				return nil
			}
			if strings.Contains(err.Error(), "Device or resource busy") && time.Now().Before(deadline) {
				time.Sleep(100 * time.Millisecond)
				continue
			}
			return err
		}
		return nil
	}
}

func isMissing(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "No such device or address") ||
		strings.Contains(msg, "not found") ||
		strings.Contains(msg, "does not exist")
}