package main

import (
	"fmt"
	"strings"
)

// plannedPart is one partition's copy plan: where it lives on each side (in
// bytes) and how its data moves.
type plannedPart struct {
	Num                 int
	SrcStartB, SrcSizeB int64
	DstStartB, DstSizeB int64
	Method              string // "raw" | "fs-image" | "file-level"
	Tool                string // fs-image tool family
	FSType              string
	ShrinkToB           int64 // >0: shrink the fs to this many bytes before imaging
}

type targetPlan struct {
	Script      string // sfdisk restore script ("" when the source has no table / Windows-native)
	IsFile      bool
	ImageSize   int64 // file target: truncate to this
	Parts       []plannedPart
	BootRegions []Region
	RootPart    int
	ESPPart     int
	UEFI        bool
	Notes       []string
}

// bootOpts is the JSON blob passed to reinstall_boot.
type bootOpts struct {
	RootPart int    `json:"rootPart"` // 1-based partition number holding "/"
	ESPPart  int    `json:"espPart"`  // 1-based ESP partition number, 0 if BIOS-only
	UEFI     bool   `json:"uefi"`
	DiskPath string `json:"diskPath"` // whole-disk path grub-install targets
}

type planOpts struct {
	Parts       map[int]bool // empty = all
	Raw         map[int]bool
	FileLevel   map[int]bool
	FileAuto    bool
	AllowShrink bool
	NewIDs      bool
	ImageSize   int64 // explicit file-target size, 0 = auto
}

const gptTailSectors = 33 // secondary GPT header + entries

// planTargetLayout turns a probed source layout + the target's nature/size
// into a concrete copy plan, applying the sizing rules:
//
//   - device target: never grow; if everything fits, keep the layout and
//     leave trailing space untouched; if the last partition(s) overflow,
//     shrink them (filesystem + partition) from the last inward, or fail.
//   - file target: size the image to hold the layout (or ImageSize).
func planTargetLayout(src *DiskLayout, isFile bool, destSize int64, o planOpts) (*targetPlan, error) {
	sector := src.LogicalSector
	if sector <= 0 {
		sector = 512
	}
	parts := sortedParts(src.Partitions)
	var sel []Partition
	for _, p := range parts {
		if len(o.Parts) == 0 || o.Parts[p.Num] {
			sel = append(sel, p)
		}
	}
	if len(sel) == 0 && src.Scheme != "raw" {
		return nil, fmt.Errorf("no partitions selected")
	}

	plan := &targetPlan{IsFile: isFile, BootRegions: append([]Region(nil), src.BootRegions...)}

	// With --new-ids on an MBR disk, keep the raw bootstrap copy clear of the
	// 4-byte disk signature at offset 0x1B8 so sfdisk's freshly generated one
	// survives (GPT regenerates its GUID in the header, which we never
	// raw-copy, so no clamp is needed there).
	if o.NewIDs && src.Scheme == "mbr" {
		for i := range plan.BootRegions {
			r := &plan.BootRegions[i]
			if r.Offset == 0 && r.Length > 0x1B8 {
				r.Length = 0x1B8
			}
		}
	}

	// Bytes the layout needs on disk = end of the last selected partition
	// (+ secondary GPT for GPT disks).
	lastEndB := int64(0)
	for _, p := range sel {
		if e := (p.Start + p.Size) * sector; e > lastEndB {
			lastEndB = e
		}
	}
	neededB := lastEndB
	if src.Scheme == "gpt" {
		neededB += gptTailSectors * sector
	}
	if src.Scheme == "raw" {
		neededB = src.DiskSize
	}

	if isFile {
		plan.ImageSize = o.ImageSize
		if plan.ImageSize == 0 {
			plan.ImageSize = roundUp(neededB+(1<<20), 1<<20)
		}
		destSize = plan.ImageSize
	}
	if destSize <= 0 {
		return nil, fmt.Errorf("could not determine target size")
	}

	// --- shrink pass ----------------------------------------------------
	sizeOverride := map[int]int64{}
	if destSize < neededB {
		short := neededB - destSize
		for i := len(sel) - 1; i >= 0 && short > 0; i-- {
			p := &sel[i]
			curB := p.Size * sector
			minB := p.FSMinBytes
			shrinkable := minB > 0 && (isExtFS(p.FSType) || p.FSType == "ntfs")
			if !shrinkable {
				continue
			}
			minB = roundUp(minB+(16<<20), sector) // 16 MiB slack
			if minB >= curB {
				continue
			}
			reduce := curB - minB
			if reduce > short {
				reduce = roundUp(short, sector)
			}
			newB := curB - reduce
			sizeOverride[p.Num] = newB / sector
			plan.Notes = append(plan.Notes,
				fmt.Sprintf("partition p%d (%s) shrinks %s -> %s to fit the target",
					p.Num, p.FSType, humanBytes(curB), humanBytes(newB)))
			short -= reduce
		}
		if short > 0 {
			return nil, fmt.Errorf(
				"target (%s) is smaller than the source layout needs (%s); the trailing partition(s) "+
					"cannot shrink to make up the missing %s; use a larger target",
				humanBytes(destSize), humanBytes(neededB), humanBytes(short))
		}
		if !o.AllowShrink {
			var names []string
			for n := range sizeOverride {
				names = append(names, fmt.Sprintf("p%d", n))
			}
			return nil, fmt.Errorf(
				"target is smaller than the source; partition(s) %s would have to be shrunk. "+
					"Re-run with --allow-shrink (this resizes the SOURCE filesystem in place before imaging) "+
					"or use --no-shrink off / a larger target",
				strings.Join(names, ", "))
		}
	}

	// --- partition-table script --------------------------------------
	if src.Scheme != "raw" && strings.TrimSpace(src.SfdiskDump) != "" {
		plan.Script = src.sfdiskRestoreScript(o.NewIDs, sizeOverride)
	}

	// --- per-partition method --------------------------------------
	for _, p := range sel {
		pp := plannedPart{
			Num:       p.Num,
			SrcStartB: p.Start * sector,
			SrcSizeB:  p.Size * sector,
			DstStartB: p.Start * sector,
			DstSizeB:  p.Size * sector,
			FSType:    p.FSType,
		}
		if ov, ok := sizeOverride[p.Num]; ok {
			pp.DstSizeB = ov * sector
			pp.ShrinkToB = ov * sector
		}
		method, tool := chooseCloneMethod(p, src.Tools, o.Raw, o.FileLevel, o.FileAuto)
		pp.Method, pp.Tool = method, tool
		if pp.ShrinkToB > 0 && method != "fs-image" {
			return nil, fmt.Errorf(
				"partition p%d (%s) must shrink to fit but no image cloner is available for it; "+
					"install ntfsclone / partclone or use a larger target", p.Num, orDash(p.FSType))
		}
		plan.Parts = append(plan.Parts, pp)
	}

	// --- boot roles ------------------------------------------------
	plan.RootPart, plan.ESPPart = pickBootParts(sel)
	plan.UEFI = plan.ESPPart != 0
	return plan, nil
}

func isExtFS(t string) bool { return t == "ext2" || t == "ext3" || t == "ext4" }

// pickBootParts guesses the "/" partition (largest Linux fs) and the ESP.
func pickBootParts(sel []Partition) (root, esp int) {
	var bestSize int64
	for _, p := range sel {
		if p.isESP() && esp == 0 {
			esp = p.Num
		}
		switch p.FSType {
		case "ext2", "ext3", "ext4", "btrfs", "xfs":
			if sz := p.Size; sz > bestSize {
				bestSize, root = sz, p.Num
			}
		}
	}
	return root, esp
}