//go:build linux

package main

import (
	"encoding/base64"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"unsafe"

	"syscall"
)

// BLKSSZGET returns the logical (soft) sector size of a block device.
const blkSSZGet = 0x1268

func logicalSectorSize(path string) (int64, error) {
	fi, err := os.Stat(path)
	if err == nil && fi.Mode().IsRegular() {
		return 512, nil
	}
	f, err := os.Open(path)
	if err != nil {
		return 512, err
	}
	defer f.Close()
	var ssz int32
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), blkSSZGet, uintptr(unsafe.Pointer(&ssz)))
	if errno != 0 || ssz <= 0 {
		return 512, nil
	}
	return int64(ssz), nil
}

// diskAttachment is a probe-time view of a whole disk: the block path whose
// partitions can be opened (the disk itself for a real device, a loop device
// for an image file) plus a cleanup to release a loop device.
type diskAttachment struct {
	blockPath string
	cleanup   func()
}

// attachDisk returns a block path whose partition nodes exist. For a regular
// file it sets up a partition-scanning loop device; for a real block device
// it is a no-op.
func attachDisk(path string, writable bool) (*diskAttachment, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return nil, err
	}
	if fi.Mode()&os.ModeDevice != 0 {
		return &diskAttachment{blockPath: path, cleanup: func() {}}, nil
	}
	args := []string{"--find", "--show", "--partscan"}
	if !writable {
		args = append(args, "--read-only")
	}
	args = append(args, path)
	out, err := exec.Command("losetup", args...).Output()
	if err != nil {
		return nil, fmt.Errorf("losetup %s: %w", path, cmdErr(err))
	}
	lo := strings.TrimSpace(string(out))
	return &diskAttachment{
		blockPath: lo,
		cleanup:   func() { _ = exec.Command("losetup", "--detach", lo).Run() },
	}, nil
}

// partNode maps a whole-disk block path + 1-based partition number to its
// partition device node (sda -> sda1, nvme0n1 -> nvme0n1p1, loop0 -> loop0p1).
func partNode(disk string, num int) string {
	if l := len(disk); l > 0 && disk[l-1] >= '0' && disk[l-1] <= '9' {
		return fmt.Sprintf("%sp%d", disk, num)
	}
	return fmt.Sprintf("%s%d", disk, num)
}

var diskTools = []string{
	"sfdisk", "sgdisk", "partprobe", "blockdev", "blkid", "lsblk", "losetup",
	"ntfsclone", "ntfsresize", "ntfsfix", "mkntfs",
	"e2image", "resize2fs", "dumpe2fs", "mkfs.ext4", "e2fsck",
	"partclone.dd", "partclone.extfs", "partclone.ntfs", "partclone.fat",
	"partclone.exfat", "partclone.xfs", "partclone.btrfs", "partclone.restore",
	"xfs_copy", "mkfs.vfat", "mkfs.xfs", "mkswap", "rsync",
	"mount", "umount", "grub-install", "update-grub", "grub-mkconfig",
}

func probeToolset() map[string]bool {
	m := make(map[string]bool, len(diskTools))
	for _, t := range diskTools {
		if _, err := exec.LookPath(t); err == nil {
			m[t] = true
		}
	}
	return m
}

// probeDiskLayout gathers everything clone-disk needs from a source disk.
func probeDiskLayout(path string) (*DiskLayout, error) {
	d := &DiskLayout{DiskPath: path, Tools: probeToolset()}

	info, err := statPath(path)
	if err != nil {
		return nil, err
	}
	if !info.Exists {
		return nil, fmt.Errorf("%s does not exist", path)
	}
	d.DiskSize = info.Size
	if d.LogicalSector, err = logicalSectorSize(path); err != nil {
		d.LogicalSector = 512
	}
	d.PhysicalSector = d.LogicalSector

	// Partition table via sfdisk -d (works on a file directly).
	dump, derr := exec.Command("sfdisk", "-d", path).Output()
	if derr != nil {
		d.Scheme = "raw"
		d.Warnings = append(d.Warnings, fmt.Sprintf("no partition table (%v)", cmdErr(derr)))
		return d, nil
	}
	if err := d.parseSfdiskDump(string(dump)); err != nil {
		return nil, err
	}
	if d.LogicalSector == 0 {
		d.LogicalSector = 512
	}
	if d.Scheme == "gpt" {
		if b, err := exec.Command("sgdisk", "--backup=/dev/stdout", path).Output(); err == nil {
			d.SgdiskB64 = base64.StdEncoding.EncodeToString(b)
		}
	}

	// Attach so per-partition nodes exist, then enrich each row.
	att, err := attachDisk(path, false)
	if err != nil {
		d.Warnings = append(d.Warnings, fmt.Sprintf("cannot attach for fs probe: %v", err))
	} else {
		defer att.cleanup()
		for i := range d.Partitions {
			enrichPartition(&d.Partitions[i], att.blockPath, d.LogicalSector)
		}
	}

	d.BootRegions = computeBootRegions(d)
	return d, nil
}

// enrichPartition fills the FS* fields of p using blkid + fs-specific probes.
func enrichPartition(p *Partition, disk string, sector int64) {
	node := partNode(disk, p.Num)
	p.DevPath = node
	if _, err := os.Stat(node); err != nil {
		return
	}
	if out, err := exec.Command("blkid", "-o", "export", node).Output(); err == nil {
		for _, line := range strings.Split(string(out), "\n") {
			k, v, ok := splitEq(strings.TrimSpace(line))
			if !ok {
				continue
			}
			switch k {
			case "TYPE":
				p.FSType = v
			case "LABEL":
				p.FSLabel = v
			case "UUID":
				p.FSUUID = v
			}
		}
	}
	switch p.FSType {
	case "ext2", "ext3", "ext4":
		p.FSUsedBytes, p.FSMinBytes = probeExtSizes(node)
	case "ntfs":
		p.FSUsedBytes, p.FSMinBytes = probeNTFSSizes(node)
	}
}

func probeExtSizes(node string) (used, min int64) {
	if out, err := exec.Command("dumpe2fs", "-h", node).Output(); err == nil {
		var bs, count, free int64
		for _, line := range strings.Split(string(out), "\n") {
			f := strings.SplitN(line, ":", 2)
			if len(f) != 2 {
				continue
			}
			v := strings.TrimSpace(f[1])
			switch strings.TrimSpace(f[0]) {
			case "Block size":
				bs, _ = strconv.ParseInt(v, 10, 64)
			case "Block count":
				count, _ = strconv.ParseInt(v, 10, 64)
			case "Free blocks":
				free, _ = strconv.ParseInt(v, 10, 64)
			}
		}
		if bs > 0 && count > 0 {
			used = (count - free) * bs
		}
	}
	// resize2fs -P needs a clean fs; ignore failures.
	if out, err := exec.Command("resize2fs", "-P", node).Output(); err == nil {
		// "Estimated minimum size of the filesystem: 123456"  (in 4k blocks)
		s := string(out)
		if i := strings.LastIndex(s, ":"); i >= 0 {
			if n, err := strconv.ParseInt(strings.TrimSpace(s[i+1:]), 10, 64); err == nil {
				min = n * 4096
			}
		}
	}
	return used, min
}

func probeNTFSSizes(node string) (used, min int64) {
	out, err := exec.Command("ntfsresize", "--info", "--force", node).CombinedOutput()
	if err != nil {
		return 0, 0
	}
	for _, line := range strings.Split(string(out), "\n") {
		line = strings.TrimSpace(line)
		switch {
		case strings.HasPrefix(line, "You might resize at "):
			// "You might resize at 1234567890 bytes ..."
			fields := strings.Fields(line)
			if len(fields) >= 5 {
				if n, err := strconv.ParseInt(fields[4], 10, 64); err == nil {
					min = n
				}
			}
		case strings.Contains(line, "Space in use") && strings.Contains(line, "bytes"):
			// "Space in use       : 1234 MB (1.2%)" - not bytes; skip
		}
	}
	return min, min
}

// computeBootRegions returns the byte ranges outside any cloned partition
// that must still be copied verbatim.
func computeBootRegions(d *DiskLayout) []Region {
	var regs []Region
	// bootstrap code before the MBR partition table on both MBR and (BIOS-)GPT disks
	regs = append(regs, Region{Offset: 0, Length: 446, Note: "MBR bootstrap"})
	if d.Scheme == "mbr" && len(d.Partitions) > 0 {
		firstStart := d.Partitions[0].Start * d.LogicalSector
		for _, p := range d.Partitions {
			s := p.Start * d.LogicalSector
			if s > 0 && s < firstStart {
				firstStart = s
			}
		}
		if firstStart > 512 {
			regs = append(regs, Region{Offset: 512, Length: firstStart - 512, Note: "MBR gap (GRUB core.img etc.)"})
		}
	}
	return regs
}

// applyPartitionTable writes script (an sfdisk restore script) to target,
// which may be a block device or an image file, and re-reads the table.
func applyPartitionTable(target, script string) error {
	cmd := exec.Command("sfdisk", "--no-reread", "--no-tell-kernel", "--force", "--wipe=always", target)
	cmd.Stdin = strings.NewReader(script)
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("sfdisk %s: %v: %s", target, err, strings.TrimSpace(string(out)))
	}
	rereadPartTable(target)
	return nil
}

func rereadPartTable(target string) {
	fi, err := os.Stat(target)
	if err != nil || fi.Mode()&os.ModeDevice == 0 {
		return
	}
	if err := exec.Command("partprobe", target).Run(); err != nil {
		_ = exec.Command("blockdev", "--rereadpt", target).Run()
	}
}

// lsblkFSType reports the filesystem type on a node ("" if none).
func lsblkFSType(node string) string {
	out, err := exec.Command("lsblk", "-ndo", "FSTYPE", node).Output()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(out))
}

// mountAt mounts node at a fresh temp dir and returns the dir + an unmount fn.
func mountAt(node string, readonly bool) (string, func(), error) {
	dir, err := os.MkdirTemp("", "clonetool-mnt-")
	if err != nil {
		return "", nil, err
	}
	args := []string{}
	if readonly {
		args = append(args, "-o", "ro")
	}
	args = append(args, node, dir)
	if out, err := exec.Command("mount", args...).CombinedOutput(); err != nil {
		os.Remove(dir)
		return "", nil, fmt.Errorf("mount %s: %v: %s", node, err, strings.TrimSpace(string(out)))
	}
	return dir, func() {
		_ = exec.Command("umount", dir).Run()
		_ = os.Remove(dir)
	}, nil
}

func cmdErr(err error) error {
	if ee, ok := err.(*exec.ExitError); ok {
		if s := strings.TrimSpace(string(ee.Stderr)); s != "" {
			return fmt.Errorf("%v: %s", err, s)
		}
	}
	return err
}