package main

import "testing"

// mkLayout builds a GPT layout: ESP (vfat) + root (ext4) + data (ntfs).
func mkLayout(diskBytes int64) *DiskLayout {
	d := &DiskLayout{
		DiskPath: "/dev/src", DiskSize: diskBytes, LogicalSector: 512,
		Scheme: "gpt", Label: "gpt", LabelID: "GUID",
		SfdiskDump: gptDump, // any non-empty dump so a script is produced
		Tools: map[string]bool{
			"sfdisk": true, "ntfsclone": true,
			"partclone.extfs": true, "partclone.restore": true, "e2image": true,
		},
		Partitions: []Partition{
			{Num: 1, Start: 2048, Size: 204800, Type: gptESP, FSType: "vfat"},                                   // 100 MiB
			{Num: 2, Start: 206848, Size: 10 * 2048 * 1024, FSType: "ext4", FSMinBytes: 2 << 30},                // 10 GiB, min 2 GiB
			{Num: 3, Start: 206848 + 10*2048*1024, Size: 10 * 2048 * 1024, FSType: "ntfs", FSMinBytes: 3 << 30}, // 10 GiB, min 3 GiB
		},
		BootRegions: []Region{{Offset: 0, Length: 446, Note: "MBR bootstrap"}},
	}
	return d
}

func lastEndBytes(d *DiskLayout) int64 {
	var e int64
	for _, p := range d.Partitions {
		if x := (p.Start + p.Size) * d.LogicalSector; x > e {
			e = x
		}
	}
	return e + gptTailSectors*d.LogicalSector
}

func TestPlanFitsDeviceKeepsLayout(t *testing.T) {
	d := mkLayout(40 << 30)
	plan, err := planTargetLayout(d, false, 40<<30, planOpts{})
	if err != nil {
		t.Fatal(err)
	}
	if len(plan.Parts) != 3 {
		t.Fatalf("parts = %d", len(plan.Parts))
	}
	for _, p := range plan.Parts {
		if p.SrcSizeB != p.DstSizeB || p.ShrinkToB != 0 {
			t.Fatalf("p%d resized unexpectedly: %+v", p.Num, p)
		}
		if p.SrcStartB != p.DstStartB {
			t.Fatalf("p%d start moved", p.Num)
		}
	}
	if plan.Parts[0].Method != "raw" { // vfat, no partclone.fat in tool set
		t.Fatalf("p1 method = %s", plan.Parts[0].Method)
	}
	if plan.Parts[1].Method != "fs-image" || plan.Parts[1].Tool != "partclone" {
		t.Fatalf("p2 method/tool = %s/%s", plan.Parts[1].Method, plan.Parts[1].Tool)
	}
	if plan.Parts[2].Method != "fs-image" || plan.Parts[2].Tool != "ntfsclone" {
		t.Fatalf("p3 method/tool = %s/%s", plan.Parts[2].Method, plan.Parts[2].Tool)
	}
	if plan.ESPPart != 1 || plan.RootPart != 2 || !plan.UEFI {
		t.Fatalf("boot roles: esp=%d root=%d uefi=%v", plan.ESPPart, plan.RootPart, plan.UEFI)
	}
}

func TestPlanTrailingFreeSpaceOK(t *testing.T) {
	d := mkLayout(60 << 30)
	// target smaller than the *disk* but larger than what the layout needs
	need := lastEndBytes(d)
	if _, err := planTargetLayout(d, false, need+1<<20, planOpts{}); err != nil {
		t.Fatalf("should fit into %d (need %d): %v", need+1<<20, need, err)
	}
}

func TestPlanNeedsShrinkRequiresFlag(t *testing.T) {
	d := mkLayout(40 << 30)
	need := lastEndBytes(d)
	small := need - (4 << 30) // 4 GiB short; ntfs (min 3 GiB) can give it up

	if _, err := planTargetLayout(d, false, small, planOpts{}); err == nil {
		t.Fatalf("expected error without --allow-shrink")
	}
	plan, err := planTargetLayout(d, false, small, planOpts{AllowShrink: true})
	if err != nil {
		t.Fatalf("with --allow-shrink: %v", err)
	}
	p3 := plan.Parts[2]
	if p3.ShrinkToB == 0 || p3.ShrinkToB >= p3.SrcSizeB {
		t.Fatalf("p3 should shrink: %+v", p3)
	}
	if p3.Method != "fs-image" {
		t.Fatalf("shrunk partition must be fs-image, got %s", p3.Method)
	}
}

func TestPlanCannotFit(t *testing.T) {
	d := mkLayout(40 << 30)
	// absurdly small: even shrinking ntfs+ext to their minimums can't help
	if _, err := planTargetLayout(d, false, 1<<30, planOpts{AllowShrink: true}); err == nil {
		t.Fatalf("expected cannot-fit error")
	}
}

func TestPlanFileTargetSizing(t *testing.T) {
	d := mkLayout(40 << 30)
	plan, err := planTargetLayout(d, true, 0, planOpts{})
	if err != nil {
		t.Fatal(err)
	}
	need := lastEndBytes(d)
	if plan.ImageSize < need || plan.ImageSize > need+(2<<20) {
		t.Fatalf("image size %d not ~%d", plan.ImageSize, need)
	}
	if plan.Script == "" {
		t.Fatalf("expected a partition-table script for a file target")
	}
}

func TestPlanNewIDsScript(t *testing.T) {
	d := mkLayout(40 << 30)
	plan, err := planTargetLayout(d, true, 0, planOpts{NewIDs: true})
	if err != nil {
		t.Fatal(err)
	}
	if plan.Script == "" || containsAny(plan.Script, "label-id:", "uuid=") {
		t.Fatalf("--new-ids script still has ids:\n%s", plan.Script)
	}
}

func containsAny(s string, subs ...string) bool {
	for _, x := range subs {
		if len(x) > 0 && indexOf(s, x) >= 0 {
			return true
		}
	}
	return false
}

func indexOf(s, sub string) int {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return i
		}
	}
	return -1
}