package mount

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"metis/pkg/util"
)

// LoopMount describes a mounted image with boot/root paths.
type LoopMount struct {
	LoopDevice string // only set when losetup created it
	BootPath   string
	RootPath   string
}

// Setup attaches an image as a loop device with partitions (-P) OR mounts an existing /dev path
// by assuming p1=boot, p2=root. Intended for Linux hosts only.
func Setup(path string) (*LoopMount, error) {
	device := path
	loopDevice := ""
	if !strings.HasPrefix(path, "/dev/") {
		var err error
		device, err = createLoop(path)
		if err != nil {
			return nil, err
		}
		loopDevice = device
	}
	bootDir, err := os.MkdirTemp("", "metis-boot-")
	if err != nil {
		return nil, err
	}
	rootDir, err := os.MkdirTemp("", "metis-root-")
	if err != nil {
		return nil, err
	}
	// Assume p1=boot, p2=root (Raspberry Pi style images)
	if err := util.Run("mount", partitionPath(device, 1), bootDir); err != nil {
		_ = Teardown(&LoopMount{LoopDevice: loopDevice, BootPath: bootDir, RootPath: rootDir})
		return nil, fmt.Errorf("mount boot: %w", err)
	}
	if err := util.Run("mount", partitionPath(device, 2), rootDir); err != nil {
		_ = util.Run("umount", bootDir)
		_ = Teardown(&LoopMount{LoopDevice: loopDevice, BootPath: bootDir, RootPath: rootDir})
		return nil, fmt.Errorf("mount root: %w", err)
	}
	return &LoopMount{LoopDevice: loopDevice, BootPath: bootDir, RootPath: rootDir}, nil
}

// Teardown unmounts and detaches the loop device.
func Teardown(m *LoopMount) error {
	if m == nil {
		return nil
	}
	if m.BootPath != "" {
		_ = util.Run("umount", m.BootPath)
		_ = os.RemoveAll(m.BootPath)
	}
	if m.RootPath != "" {
		_ = util.Run("umount", m.RootPath)
		_ = os.RemoveAll(m.RootPath)
	}
	if m.LoopDevice != "" {
		_ = util.Run("losetup", "-d", m.LoopDevice)
	}
	return nil
}

func partitionPath(base string, idx int) string {
	p := fmt.Sprintf("%sp%d", base, idx)
	if _, err := os.Stat(p); err == nil {
		return p
	}
	return fmt.Sprintf("%s%d", base, idx)
}

func createLoop(imagePath string) (string, error) {
	// losetup -Pf --show <image>
	out, err := util.RunLogged("losetup", "-Pf", "--show", filepath.Clean(imagePath))
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(out), nil
}