// Package partition detects and manages A/B system partitions.
//
// It identifies System A and System B partitions by label (KSOLOA, KSOLOB)
// and provides mount/write operations for the update process.
package partition

import (
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const (
	LabelSystemA = "KSOLOA"
	LabelSystemB = "KSOLOB"
	LabelData    = "KSOLODATA"
	LabelEFI     = "KSOLOEFI"
)

// Info contains information about a partition.
type Info struct {
	Device     string // e.g. /dev/sda2
	Label      string // e.g. KSOLOA
	MountPoint string // current mount point, empty if not mounted
	Slot       string // "A" or "B"
}

// FindByLabel locates a block device by its filesystem label.
func FindByLabel(label string) (string, error) {
	cmd := exec.Command("blkid", "-L", label)
	output, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("partition with label %q not found: %w", label, err)
	}
	return strings.TrimSpace(string(output)), nil
}

// GetSlotPartition returns the partition info for the given slot ("A" or "B").
func GetSlotPartition(slot string) (*Info, error) {
	var label string
	switch slot {
	case "A":
		label = LabelSystemA
	case "B":
		label = LabelSystemB
	default:
		return nil, fmt.Errorf("invalid slot: %q", slot)
	}

	dev, err := FindByLabel(label)
	if err != nil {
		return nil, err
	}

	return &Info{
		Device: dev,
		Label:  label,
		Slot:   slot,
	}, nil
}

// MountReadOnly mounts a partition read-only at the given mount point.
func MountReadOnly(dev, mountPoint string) error {
	if err := os.MkdirAll(mountPoint, 0o755); err != nil {
		return fmt.Errorf("creating mount point: %w", err)
	}
	cmd := exec.Command("mount", "-o", "ro", dev, mountPoint)
	if output, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("mounting %s at %s: %w\n%s", dev, mountPoint, err, output)
	}
	slog.Debug("mounted", "device", dev, "mountpoint", mountPoint, "mode", "ro")
	return nil
}

// MountReadWrite mounts a partition read-write at the given mount point.
func MountReadWrite(dev, mountPoint string) error {
	if err := os.MkdirAll(mountPoint, 0o755); err != nil {
		return fmt.Errorf("creating mount point: %w", err)
	}
	cmd := exec.Command("mount", dev, mountPoint)
	if output, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("mounting %s at %s: %w\n%s", dev, mountPoint, err, output)
	}
	slog.Debug("mounted", "device", dev, "mountpoint", mountPoint, "mode", "rw")
	return nil
}

// Unmount unmounts a mount point.
func Unmount(mountPoint string) error {
	cmd := exec.Command("umount", mountPoint)
	if output, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("unmounting %s: %w\n%s", mountPoint, err, output)
	}
	return nil
}

// ReadVersion reads the version file from a mounted system partition.
func ReadVersion(mountPoint string) (string, error) {
	data, err := os.ReadFile(filepath.Join(mountPoint, "version"))
	if err != nil {
		return "", fmt.Errorf("reading version: %w", err)
	}
	return strings.TrimSpace(string(data)), nil
}

// WriteSystemImage copies vmlinuz and initramfs to a mounted partition.
func WriteSystemImage(mountPoint, vmlinuzPath, initramfsPath, version string) error {
	// Copy vmlinuz
	if err := copyFile(vmlinuzPath, filepath.Join(mountPoint, "vmlinuz")); err != nil {
		return fmt.Errorf("writing vmlinuz: %w", err)
	}

	// Copy initramfs
	if err := copyFile(initramfsPath, filepath.Join(mountPoint, "kubesolo-os.gz")); err != nil {
		return fmt.Errorf("writing initramfs: %w", err)
	}

	// Write version
	if err := os.WriteFile(filepath.Join(mountPoint, "version"), []byte(version+"\n"), 0o644); err != nil {
		return fmt.Errorf("writing version: %w", err)
	}

	// Sync to ensure data is flushed to disk
	exec.Command("sync").Run()

	slog.Info("system image written", "mountpoint", mountPoint, "version", version)
	return nil
}

func copyFile(src, dst string) error {
	data, err := os.ReadFile(src)
	if err != nil {
		return err
	}
	return os.WriteFile(dst, data, 0o644)
}