package main

import (
	"strings"
	"testing"
)

const gptDump = `label: gpt
label-id: 1D2E3F4A-1111-2222-3333-444455556666
device: /dev/sda
unit: sectors
first-lba: 2048
last-lba: 8388574
sector-size: 512

/dev/sda1 : start=        2048, size=      204800, type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B, uuid=AAAAAAAA-1111-2222-3333-444455556666, name="EFI System Partition"
/dev/sda2 : start=      206848, size=     8179712, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=BBBBBBBB-1111-2222-3333-444455556666
`

const mbrDump = `label: dos
label-id: 0x1a2b3c4d
device: /dev/sdb
unit: sectors
sector-size: 512

/dev/sdb1 : start=        2048, size=      512000, type=83, bootable
/dev/sdb2 : start=      514048, size=    41734144, type=83
`

func TestParseSfdiskDumpGPT(t *testing.T) {
	var d DiskLayout
	if err := d.parseSfdiskDump(gptDump); err != nil {
		t.Fatal(err)
	}
	if d.Label != "gpt" || d.Scheme != "gpt" {
		t.Fatalf("label/scheme = %q/%q", d.Label, d.Scheme)
	}
	if d.LabelID != "1D2E3F4A-1111-2222-3333-444455556666" {
		t.Fatalf("labelID = %q", d.LabelID)
	}
	if d.FirstLBA != 2048 || d.LastLBA != 8388574 {
		t.Fatalf("lba = %d/%d", d.FirstLBA, d.LastLBA)
	}
	if len(d.Partitions) != 2 {
		t.Fatalf("parts = %d", len(d.Partitions))
	}
	p1 := d.Partitions[0]
	if p1.Num != 1 || p1.Start != 2048 || p1.Size != 204800 {
		t.Fatalf("p1 = %+v", p1)
	}
	if p1.Type != "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" || p1.UUID != "AAAAAAAA-1111-2222-3333-444455556666" {
		t.Fatalf("p1 type/uuid = %q/%q", p1.Type, p1.UUID)
	}
	if p1.Name != "EFI System Partition" {
		t.Fatalf("p1 name = %q", p1.Name)
	}
	if !p1.isESP() {
		t.Fatalf("p1 should be ESP")
	}
	if d.Partitions[1].Num != 2 || d.Partitions[1].Size != 8179712 {
		t.Fatalf("p2 = %+v", d.Partitions[1])
	}
}

func TestParseSfdiskDumpMBR(t *testing.T) {
	var d DiskLayout
	if err := d.parseSfdiskDump(mbrDump); err != nil {
		t.Fatal(err)
	}
	if d.Scheme != "mbr" || d.LabelID != "0x1a2b3c4d" {
		t.Fatalf("scheme/id = %q/%q", d.Scheme, d.LabelID)
	}
	if len(d.Partitions) != 2 || d.Partitions[0].Attrs != "bootable" {
		t.Fatalf("parts = %+v", d.Partitions)
	}
}

func TestSfdiskRestoreScript(t *testing.T) {
	var d DiskLayout
	if err := d.parseSfdiskDump(gptDump); err != nil {
		t.Fatal(err)
	}

	// keep IDs
	s := d.sfdiskRestoreScript(false, nil)
	if !strings.Contains(s, "label-id: 1D2E3F4A-1111-2222-3333-444455556666") {
		t.Fatalf("expected label-id kept:\n%s", s)
	}
	if !strings.Contains(s, "uuid=AAAAAAAA-1111-2222-3333-444455556666") {
		t.Fatalf("expected part uuid kept:\n%s", s)
	}
	if strings.Contains(s, "device:") || strings.Contains(s, "/dev/sda1") {
		t.Fatalf("script must be device-independent:\n%s", s)
	}
	if strings.Contains(s, "last-lba") {
		t.Fatalf("script must drop last-lba:\n%s", s)
	}

	// new IDs
	s2 := d.sfdiskRestoreScript(true, nil)
	if strings.Contains(s2, "label-id:") || strings.Contains(s2, "uuid=") {
		t.Fatalf("--new-ids must strip ids:\n%s", s2)
	}

	// size override on p2
	s3 := d.sfdiskRestoreScript(false, map[int]int64{2: 4000000})
	if !strings.Contains(s3, "size=4000000") {
		t.Fatalf("expected overridden size:\n%s", s3)
	}
	if strings.Contains(s3, "size=8179712") {
		t.Fatalf("old size should be gone:\n%s", s3)
	}
}