package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// fsCloner is the pair of argv templates that stream one filesystem as an
// image: save reads the fs on dev and writes an image to stdout; restore
// reads that image from stdin and writes it onto dev. The image format is
// whatever the chosen standard tool uses (ntfsclone, partclone, e2image) —
// clonetool only moves the bytes between the two ends.
type fsCloner struct {
	name    string // tool family, for logs
	save    []string
	restore []string
}

// fsCloneOpts is the JSON blob carried in CtrlMsg.Options for a fs-image
// partition clone so both ends build the identical argv.
type fsCloneOpts struct {
	FSType    string `json:"fsType"`
	Tool      string `json:"tool"`
	SizeBytes int64  `json:"sizeBytes"`
	// ShrinkToBytes > 0: shrink the SOURCE filesystem to this size in place
	// before imaging it (so it fits a smaller target partition). This
	// mutates the source and only runs when the operator passed
	// --allow-shrink --yes.
	ShrinkToBytes int64 `json:"shrinkToBytes,omitempty"`
}

func (o fsCloneOpts) encode() string { b, _ := json.Marshal(o); return string(b) }
func decodeFSCloneOpts(s string) (fsCloneOpts, error) {
	var o fsCloneOpts
	err := json.Unmarshal([]byte(s), &o)
	return o, err
}

// pickCloner returns the fs-image cloner for (fsType, tool) or nil if the
// partition should be handled some other way. tool is the family name chosen
// by chooseCloneMethod ("ntfsclone", "partclone", "e2image").
func pickCloner(dev, fsType, tool string) *fsCloner {
	switch tool {
	case "ntfsclone":
		return &fsCloner{
			name:    "ntfsclone",
			save:    []string{"ntfsclone", "--save-image", "--force", "--output", "-", dev},
			restore: []string{"ntfsclone", "--restore-image", "--overwrite", dev, "-"},
		}
	case "e2image":
		return &fsCloner{
			name: "e2image",
			// -ra: raw + skip unallocated; stream to/from stdio.
			save:    []string{"e2image", "-ra", "-p", dev, "-"},
			restore: []string{"e2image", "-ra", "-p", "-", dev},
		}
	case "partclone":
		pc := "partclone." + partcloneSuffix(fsType)
		return &fsCloner{
			name:    pc,
			save:    []string{pc, "-c", "-s", dev, "-O", "-", "-L", "/dev/null"},
			restore: []string{"partclone.restore", "-s", "-", "-O", dev, "-L", "/dev/null"},
		}
	}
	return nil
}

func partcloneSuffix(fsType string) string {
	switch fsType {
	case "ext2", "ext3", "ext4":
		return "extfs"
	case "vfat", "fat", "fat12", "fat16", "fat32", "msdos":
		return "fat"
	case "ntfs":
		return "ntfs"
	default:
		return fsType // exfat, xfs, btrfs, f2fs, hfsplus, ...
	}
}

// chooseCloneMethod decides how a partition's data is copied.
//
//	"raw"        -> the block-diff engine over an offset window (also the
//	                re-run-friendly path; used for swap, unknown fs, --raw)
//	"fs-image"   -> stream a filesystem image with a standard tool
//	"file-level" -> mkfs on the target + rsync
func chooseCloneMethod(p Partition, tools map[string]bool, forceRaw, forceFile map[int]bool, fileAuto bool) (method, tool string) {
	if forceRaw[p.Num] {
		return "raw", ""
	}
	if forceFile[p.Num] {
		return "file-level", ""
	}
	if p.isBIOSBoot() {
		return "raw", "" // tiny, and holds GRUB core - always verbatim
	}
	switch p.FSType {
	case "", "swap", "crypto_LUKS", "LVM2_member":
		return "raw", ""
	case "ntfs":
		if tools["ntfsclone"] {
			return "fs-image", "ntfsclone"
		}
	case "ext2", "ext3", "ext4":
		if tools["partclone.extfs"] && tools["partclone.restore"] {
			return "fs-image", "partclone"
		}
		if tools["e2image"] {
			return "fs-image", "e2image"
		}
	case "vfat", "fat", "fat12", "fat16", "fat32", "msdos":
		if tools["partclone.fat"] && tools["partclone.restore"] {
			return "fs-image", "partclone"
		}
	case "exfat", "xfs", "btrfs", "f2fs", "hfsplus":
		if tools["partclone."+partcloneSuffix(p.FSType)] && tools["partclone.restore"] {
			return "fs-image", "partclone"
		}
	}
	if fileAuto && p.FSType != "" && tools["rsync"] {
		return "file-level", ""
	}
	return "raw", ""
}

// progressScanner reads a cloner's stderr line by line, pulls a percentage
// out of the tool's own progress chatter, and calls emit with an estimated
// byte count (fraction * sizeBytes). It swallows everything else.
var pctRe = regexp.MustCompile(`([0-9]{1,3}(?:\.[0-9]+)?)\s*%|([0-9]{1,3}(?:\.[0-9]+)?)\s*percent`)

func scanCloneProgress(r io.Reader, sizeBytes int64, emit func(bytesDone int64)) {
	br := bufio.NewReader(r)
	var buf []byte
	flush := func() {
		line := strings.TrimSpace(string(buf))
		buf = buf[:0]
		if line == "" {
			return
		}
		m := pctRe.FindStringSubmatch(line)
		if m == nil {
			return
		}
		s := m[1]
		if s == "" {
			s = m[2]
		}
		if f, err := strconv.ParseFloat(s, 64); err == nil && f >= 0 && f <= 100 {
			emit(int64(f / 100 * float64(sizeBytes)))
		}
	}
	for {
		c, err := br.ReadByte()
		if err != nil {
			flush()
			return
		}
		if c == '\n' || c == '\r' {
			flush()
			continue
		}
		buf = append(buf, c)
		if len(buf) > 4096 {
			buf = buf[:0]
		}
	}
}

// preShrinkFS shrinks the filesystem on dev to (at least) toBytes in place.
// Used only on the source, only with --allow-shrink --yes. Best effort with a
// hard error if the resize tool fails.
func preShrinkFS(dev, fsType string, toBytes int64) error {
	if toBytes <= 0 {
		return nil
	}
	switch {
	case isExtFS(fsType):
		if out, err := runCmd("e2fsck", "-f", "-y", dev); err != nil {
			return fmt.Errorf("e2fsck %s: %v: %s", dev, err, out)
		}
		kib := toBytes / 1024
		if out, err := runCmd("resize2fs", dev, fmt.Sprintf("%dK", kib)); err != nil {
			return fmt.Errorf("resize2fs %s %dK: %v: %s", dev, kib, err, out)
		}
		return nil
	case fsType == "ntfs":
		// ntfsresize prompts for confirmation on a real resize; feed it "y".
		if out, err := runCmdStdin("y\n", "ntfsresize", "--force", "--size", strconv.FormatInt(toBytes, 10), dev); err != nil {
			return fmt.Errorf("ntfsresize %s --size %d: %v: %s", dev, toBytes, err, out)
		}
		return nil
	default:
		return fmt.Errorf("don't know how to shrink a %s filesystem", fsType)
	}
}