package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"testing"
	"time"
	"unicode/utf16"
)

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

func mustGeometry(t *testing.T, size int64, scheme, label string, clus int64) Geometry {
	t.Helper()
	g, err := PlanGeometry(size, scheme, label, clus)
	if err != nil {
		t.Fatalf("PlanGeometry(%d, %q): %v", size, scheme, err)
	}
	return g
}

// writeFixture materialises a map of relative path -> content under dir.
func writeFixture(t *testing.T, dir string, files map[string][]byte) {
	t.Helper()
	for rel, content := range files {
		full := filepath.Join(dir, filepath.FromSlash(rel))
		if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
			t.Fatalf("mkdir %s: %v", filepath.Dir(full), err)
		}
		if err := os.WriteFile(full, content, 0o644); err != nil {
			t.Fatalf("write %s: %v", full, err)
		}
	}
}

// buildTo writes an image from src and returns its path.
func buildTo(t *testing.T, src, out string, size int64, scheme, label string, clus int64) (Geometry, *BuildResult) {
	t.Helper()
	g := mustGeometry(t, size, scheme, label, clus)
	root, scan, err := ScanTree(src)
	if err != nil {
		t.Fatalf("ScanTree(%s): %v", src, err)
	}
	res, err := WriteImage(out, g, root, scan, false, time.Now())
	if err != nil {
		t.Fatalf("WriteImage: %v", err)
	}
	return g, res
}

// ---------------------------------------------------------------------------
// Boot sector / BPB field correctness
// ---------------------------------------------------------------------------

func TestBootSectorFieldsKnownValues(t *testing.T) {
	// A 64 MiB MBR image. Every expected value below is derived by hand from
	// the FAT32 spec and the geometry rules, not read back from the writer.
	g := mustGeometry(t, 64<<20, "mbr", "WINPE", 0)
	if g.SectorsPerCluster != 1 || g.PartSectors != 129024 || g.FATSectors != 993 || g.ClusterCount != 127006 {
		t.Fatalf("geometry drifted: spc=%d partSectors=%d fatSectors=%d clusters=%d",
			g.SectorsPerCluster, g.PartSectors, g.FATSectors, g.ClusterCount)
	}
	bs := BuildBootSector(g, 0x12345678)
	if len(bs) != 512 {
		t.Fatalf("boot sector is %d bytes, want 512", len(bs))
	}

	type check struct {
		name string
		off  int
		size int // 1, 2 or 4
		want uint32
	}
	for _, c := range []check{
		{"BS_jmpBoot[0]", 0, 1, 0xEB},
		{"BS_jmpBoot[1]", 1, 1, 0x58},
		{"BS_jmpBoot[2]", 2, 1, 0x90},
		{"BPB_BytsPerSec", 11, 2, 512},
		{"BPB_SecPerClus", 13, 1, 1},
		{"BPB_RsvdSecCnt", 14, 2, 32},
		{"BPB_NumFATs", 16, 1, 2},
		{"BPB_RootEntCnt", 17, 2, 0},
		{"BPB_TotSec16", 19, 2, 0},
		{"BPB_Media", 21, 1, 0xF8},
		{"BPB_FATSz16", 22, 2, 0},
		{"BPB_SecPerTrk", 24, 2, 63},
		{"BPB_NumHeads", 26, 2, 255},
		{"BPB_HiddSec", 28, 4, 2048},
		{"BPB_TotSec32", 32, 4, 129024},
		{"BPB_FATSz32", 36, 4, 993},
		{"BPB_ExtFlags", 40, 2, 0},
		{"BPB_FSVer", 42, 2, 0},
		{"BPB_RootClus", 44, 4, 2},
		{"BPB_FSInfo", 48, 2, 1},
		{"BPB_BkBootSec", 50, 2, 6},
		{"BS_DrvNum", 64, 1, 0x80},
		{"BS_Reserved1", 65, 1, 0},
		{"BS_BootSig", 66, 1, 0x29},
		{"BS_VolID", 67, 4, 0x12345678},
		{"signature[0]", 510, 1, 0x55},
		{"signature[1]", 511, 1, 0xAA},
	} {
		var got uint32
		switch c.size {
		case 1:
			got = uint32(bs[c.off])
		case 2:
			got = uint32(binary.LittleEndian.Uint16(bs[c.off : c.off+2]))
		case 4:
			got = binary.LittleEndian.Uint32(bs[c.off : c.off+4])
		}
		if got != c.want {
			t.Errorf("%s at offset %d = 0x%X, want 0x%X", c.name, c.off, got, c.want)
		}
	}
	if s := string(bs[3:11]); s != "MSWIN4.1" {
		t.Errorf("BS_OEMName = %q, want %q", s, "MSWIN4.1")
	}
	if s := string(bs[71:82]); s != "WINPE      " {
		t.Errorf("BS_VolLab = %q, want %q", s, "WINPE      ")
	}
	if s := string(bs[82:90]); s != "FAT32   " {
		t.Errorf("BS_FilSysType = %q, want %q", s, "FAT32   ")
	}
	for i := 52; i < 64; i++ {
		if bs[i] != 0 {
			t.Errorf("BPB_Reserved byte %d = 0x%02X, want 0", i, bs[i])
		}
	}
	// No bootloader is installed, so the code area must be empty.
	for i := 90; i < 510; i++ {
		if bs[i] != 0 {
			t.Fatalf("boot code area byte %d is 0x%02X; BootBuilder installs no bootloader", i, bs[i])
		}
	}
}

func TestFSInfoSectorFields(t *testing.T) {
	s := BuildFSInfo(116628, 10380)
	if len(s) != 512 {
		t.Fatalf("FSInfo is %d bytes, want 512", len(s))
	}
	if v := binary.LittleEndian.Uint32(s[0:4]); v != 0x41615252 {
		t.Errorf("FSI_LeadSig = 0x%08X, want 0x41615252", v)
	}
	if v := binary.LittleEndian.Uint32(s[484:488]); v != 0x61417272 {
		t.Errorf("FSI_StrucSig = 0x%08X, want 0x61417272", v)
	}
	if v := binary.LittleEndian.Uint32(s[488:492]); v != 116628 {
		t.Errorf("FSI_Free_Count = %d, want 116628", v)
	}
	if v := binary.LittleEndian.Uint32(s[492:496]); v != 10380 {
		t.Errorf("FSI_Nxt_Free = %d, want 10380", v)
	}
	if v := binary.LittleEndian.Uint32(s[508:512]); v != 0xAA550000 {
		t.Errorf("FSI_TrailSig = 0x%08X, want 0xAA550000", v)
	}
	for i := 4; i < 484; i++ {
		if s[i] != 0 {
			t.Fatalf("FSI_Reserved1 byte %d = 0x%02X, want 0", i, s[i])
		}
	}
}

// ---------------------------------------------------------------------------
// Microsoft's sectors-per-cluster table
// ---------------------------------------------------------------------------

func TestSectorsPerClusterTable(t *testing.T) {
	for _, tc := range []struct {
		name    string
		sectors int64
		want    int64
	}{
		{"32.5MB boundary, FAT32 not allowed", 66600, 0},
		{"one sector past the boundary", 66601, 1},
		{"260MB boundary", 532480, 1},
		{"one past 260MB", 532481, 8},
		{"8GB boundary", 16777216, 8},
		{"one past 8GB", 16777217, 16},
		{"16GB boundary", 33554432, 16},
		{"one past 16GB", 33554433, 32},
		{"32GB boundary", 67108864, 32},
		{"one past 32GB", 67108865, 64},
		{"1TB", 2 << 31, 64},
	} {
		t.Run(tc.name, func(t *testing.T) {
			if got := secPerClusForSectors(tc.sectors); got != tc.want {
				t.Errorf("secPerClusForSectors(%d) = %d, want %d", tc.sectors, got, tc.want)
			}
		})
	}
}

func TestGeometryPicksTableClusterSize(t *testing.T) {
	for _, tc := range []struct {
		size        int64
		wantCluster int64
	}{
		{64 << 20, 512},
		{8 << 30, 4096},
		{16 << 30, 8192},
		{40 << 30, 32768},
	} {
		g := mustGeometry(t, tc.size, "mbr", "L", 0)
		if g.ClusterBytes != tc.wantCluster {
			t.Errorf("size %s: cluster = %d bytes, want %d", humanBytes(tc.size), g.ClusterBytes, tc.wantCluster)
		}
		if !g.ClusterFromTable {
			t.Errorf("size %s: cluster size should be marked as coming from the MS table", humanBytes(tc.size))
		}
		// The FAT must be able to address every cluster it claims to have.
		if entries := g.FATBytes / 4; entries < g.ClusterCount+2 {
			t.Errorf("size %s: FAT holds %d entries but the volume has %d clusters",
				humanBytes(tc.size), entries, g.ClusterCount)
		}
		// And it must not spill past the partition.
		used := g.ReservedSectors + g.NumFATs*g.FATSectors + g.DataSectors
		if used > g.PartSectors {
			t.Errorf("size %s: layout uses %d sectors of a %d-sector partition", humanBytes(tc.size), used, g.PartSectors)
		}
		if g.SlackSectors < 0 || g.SlackSectors >= g.SectorsPerCluster+g.NumFATs {
			t.Errorf("size %s: implausible slack of %d sectors", humanBytes(tc.size), g.SlackSectors)
		}
	}
}

func TestGeometryRejectsTinyVolumes(t *testing.T) {
	if _, err := PlanGeometry(20<<20, "mbr", "L", 0); err == nil {
		t.Fatal("a 20 MiB image must be rejected: FAT32 needs more than 65525 clusters")
	} else if !errors.Is(err, errTooSmall) {
		t.Errorf("error = %v, want it to wrap errTooSmall", err)
	}
}

// ---------------------------------------------------------------------------
// 8.3 short names
// ---------------------------------------------------------------------------

func TestShortNameGeneration(t *testing.T) {
	for _, tc := range []struct {
		name      string
		inputs    []string
		wantShort []string
		wantLFN   []bool
	}{
		{
			name:      "already 8.3 and uppercase needs no long name",
			inputs:    []string{"BOOTX64.EFI"},
			wantShort: []string{"BOOTX64 EFI"},
			wantLFN:   []bool{false},
		},
		{
			name:      "lowercase keeps the short name but needs a long one",
			inputs:    []string{"readme.txt"},
			wantShort: []string{"README  TXT"},
			wantLFN:   []bool{true},
		},
		{
			name:      "no extension",
			inputs:    []string{"README"},
			wantShort: []string{"README     "},
			wantLFN:   []bool{false},
		},
		{
			name:      "too long gets a numeric tail",
			inputs:    []string{"autounattend.xml"},
			wantShort: []string{"AUTOUN~1XML"},
			wantLFN:   []bool{true},
		},
		{
			name:      "collisions number upward",
			inputs:    []string{"autounattend.xml", "autounattend2.xml", "autounattendZZ.xml"},
			wantShort: []string{"AUTOUN~1XML", "AUTOUN~2XML", "AUTOUN~3XML"},
			wantLFN:   []bool{true, true, true},
		},
		{
			name:      "case-only collision still gets a tail on the second entry",
			inputs:    []string{"Readme.txt", "README.TXT"},
			wantShort: []string{"README  TXT", "README~1TXT"},
			wantLFN:   []bool{true, true},
		},
		{
			name:      "invalid characters become underscores",
			inputs:    []string{"A+Long,File;Name=With[Invalid]Chars.txt"},
			wantShort: []string{"A_LONG~1TXT"},
			wantLFN:   []bool{true},
		},
		{
			name:      "spaces are dropped, not substituted",
			inputs:    []string{"Windows Setup Files"},
			wantShort: []string{"WINDOW~1   "},
			wantLFN:   []bool{true},
		},
		{
			name:      "multiple dots keep only the last as the extension",
			inputs:    []string{"archive.tar.gz"},
			wantShort: []string{"ARCHIV~1GZ "},
			wantLFN:   []bool{true},
		},
		{
			name:      "leading dot is stripped",
			inputs:    []string{".gitignore"},
			wantShort: []string{"GITIGN~1   "},
			wantLFN:   []bool{true},
		},
		{
			name:      "extension longer than three characters is truncated",
			inputs:    []string{"setup.config"},
			wantShort: []string{"SETUP~1 CON"},
			wantLFN:   []bool{true},
		},
		{
			name:      "exactly eight and three fits",
			inputs:    []string{"INSTALL1.WIM"},
			wantShort: []string{"INSTALL1WIM"},
			wantLFN:   []bool{false},
		},
		{
			name:      "nine characters does not fit",
			inputs:    []string{"INSTALL12.WIM"},
			wantShort: []string{"INSTAL~1WIM"},
			wantLFN:   []bool{true},
		},
		{
			name:      "non-ASCII is substituted",
			inputs:    []string{"naïve.txt"},
			wantShort: []string{"NA__VE~1TXT"},
			wantLFN:   []bool{true},
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			taken := map[string]bool{}
			for i, in := range tc.inputs {
				got, lfn, err := GenerateShortName(in, taken)
				if err != nil {
					t.Fatalf("GenerateShortName(%q): %v", in, err)
				}
				if string(got[:]) != tc.wantShort[i] {
					t.Errorf("GenerateShortName(%q) = %q, want %q", in, string(got[:]), tc.wantShort[i])
				}
				if lfn != tc.wantLFN[i] {
					t.Errorf("GenerateShortName(%q) needLFN = %v, want %v", in, lfn, tc.wantLFN[i])
				}
			}
		})
	}
}

func TestShortNameTailGrowsPastNine(t *testing.T) {
	taken := map[string]bool{}
	// 12 names that all reduce to the same basis, forcing ~1 .. ~12. At ~10 the
	// tail is three characters, so the base must shrink to five.
	var last string
	for i := 0; i < 12; i++ {
		n := fmt.Sprintf("configuration-%02d.ini", i)
		s, _, err := GenerateShortName(n, taken)
		if err != nil {
			t.Fatalf("GenerateShortName(%q): %v", n, err)
		}
		last = string(s[:])
	}
	if last != "CONFI~12INI" {
		t.Errorf("twelfth collision = %q, want %q", last, "CONFI~12INI")
	}
	if len(taken) != 12 {
		t.Errorf("expected 12 distinct short names, got %d", len(taken))
	}
}

// ---------------------------------------------------------------------------
// LFN checksums and slot ordering
// ---------------------------------------------------------------------------

func TestLFNChecksumKnownValues(t *testing.T) {
	// Values computed independently from the ChkSum() pseudo-code in fatgen103.
	for _, tc := range []struct {
		short string
		want  byte
	}{
		{"BOOTX64 EFI", 0x1D},
		{"AUTOUN~1XML", 0xAB},
		{"DEEPFI~1TXT", 0x4F},
		{"ALONGF~1CON", 0xF3},
		{"README  TXT", 0x73},
		{"A_LONG~1TXT", 0x70},
	} {
		var s [11]byte
		copy(s[:], tc.short)
		if got := LFNChecksum(s); got != tc.want {
			t.Errorf("LFNChecksum(%q) = 0x%02X, want 0x%02X", tc.short, got, tc.want)
		}
	}
}

func TestLFNSlotOrderingAndChecksums(t *testing.T) {
	for _, tc := range []struct {
		long      string
		short     string
		wantSlots int
	}{
		{"readme.txt", "README  TXT", 1},
		{"exactly-13ch.", "EXACTL~1   ", 1},                                                 // 13 units
		{"fourteen-chars", "FOURTE~1   ", 2},                                                // 14 units
		{"deep file with a very long name indeed.txt", "DEEPFI~1TXT", 4},                    // 42 units
		{"A Long File Name That Needs Several LFN Slots To Store.config", "ALONGF~1CON", 5}, // 61 units
		{strings.Repeat("z", 130) + ".dat", "ZZZZZZ~1DAT", 11},                              // 134 units
	} {
		var short [11]byte
		copy(short[:], tc.short)
		raw := BuildLFNSlots(tc.long, short)
		if len(raw) != tc.wantSlots*32 {
			t.Fatalf("%q: got %d bytes (%d slots), want %d slots", tc.long, len(raw), len(raw)/32, tc.wantSlots)
		}
		want := LFNChecksum(short)
		units := utf16.Encode([]rune(tc.long))

		for i := 0; i < tc.wantSlots; i++ {
			e := raw[i*32 : (i+1)*32]
			ord := int(e[0] & 0x3F)
			isLast := e[0]&0x40 != 0
			// On disk the slots run in REVERSE: the highest ordinal first,
			// flagged LAST_LONG_ENTRY, counting down to 1.
			if wantOrd := tc.wantSlots - i; ord != wantOrd {
				t.Errorf("%q slot %d: ordinal %d, want %d", tc.long, i, ord, wantOrd)
			}
			if isLast != (i == 0) {
				t.Errorf("%q slot %d: LAST_LONG_ENTRY = %v, want %v", tc.long, i, isLast, i == 0)
			}
			if e[11] != attrLongName {
				t.Errorf("%q slot %d: attr 0x%02X, want 0x0F", tc.long, i, e[11])
			}
			if e[12] != 0 {
				t.Errorf("%q slot %d: type byte 0x%02X, want 0", tc.long, i, e[12])
			}
			if e[13] != want {
				t.Errorf("%q slot %d: checksum 0x%02X, want 0x%02X", tc.long, i, e[13], want)
			}
			if e[26] != 0 || e[27] != 0 {
				t.Errorf("%q slot %d: first-cluster field must be zero in an LFN slot", tc.long, i)
			}
		}

		// Reassemble in ordinal order and check we get the name back.
		offs := []int{1, 3, 5, 7, 9, 14, 16, 18, 20, 22, 24, 28, 30}
		got := make([]uint16, 0, len(units))
		for ord := 1; ord <= tc.wantSlots; ord++ {
			e := raw[(tc.wantSlots-ord)*32 : (tc.wantSlots-ord+1)*32]
			for _, o := range offs {
				v := binary.LittleEndian.Uint16(e[o : o+2])
				if v == 0x0000 || v == 0xFFFF {
					continue
				}
				got = append(got, v)
			}
		}
		if string(utf16.Decode(got)) != tc.long {
			t.Errorf("reassembled %q, want %q", string(utf16.Decode(got)), tc.long)
		}
	}
}

func TestLFNPaddingUsesTerminatorThenFFFF(t *testing.T) {
	var short [11]byte
	copy(short[:], "AB~1       ")
	raw := BuildLFNSlots("ab", short) // 2 units in a 13-unit slot
	if len(raw) != 32 {
		t.Fatalf("want one slot, got %d bytes", len(raw))
	}
	offs := []int{1, 3, 5, 7, 9, 14, 16, 18, 20, 22, 24, 28, 30}
	for i, o := range offs {
		v := binary.LittleEndian.Uint16(raw[o : o+2])
		var want uint16
		switch {
		case i < 2:
			want = uint16("ab"[i])
		case i == 2:
			want = 0x0000
		default:
			want = 0xFFFF
		}
		if v != want {
			t.Errorf("character slot %d = 0x%04X, want 0x%04X", i, v, want)
		}
	}
}

// ---------------------------------------------------------------------------
// GPT CRC32, checked against a bit-by-bit reference implementation
// ---------------------------------------------------------------------------

// refCRC32 is a deliberately naive bitwise CRC-32/ISO-HDLC. It shares no code
// with hash/crc32, so agreeing with it is real evidence.
func refCRC32(data []byte) uint32 {
	crc := uint32(0xFFFFFFFF)
	for _, b := range data {
		crc ^= uint32(b)
		for i := 0; i < 8; i++ {
			if crc&1 != 0 {
				crc = (crc >> 1) ^ 0xEDB88320
			} else {
				crc >>= 1
			}
		}
	}
	return crc ^ 0xFFFFFFFF
}

func TestRefCRC32AgreesOnKnownVectors(t *testing.T) {
	for _, tc := range []struct {
		in   string
		want uint32
	}{
		{"", 0x00000000},
		{"a", 0xE8B7BE43},
		{"123456789", 0xCBF43926},
		{"The quick brown fox jumps over the lazy dog", 0x414FA339},
	} {
		if got := refCRC32([]byte(tc.in)); got != tc.want {
			t.Errorf("refCRC32(%q) = 0x%08X, want 0x%08X", tc.in, got, tc.want)
		}
	}
}

func TestGPTCRCsAndHeaderFields(t *testing.T) {
	g := mustGeometry(t, 512<<20, "gpt", "ESPTEST", 0)
	var diskGUID, partGUID [16]byte
	for i := range diskGUID {
		diskGUID[i] = byte(i + 1)
		partGUID[i] = byte(0x80 + i)
	}
	gpt := BuildGPT(g.TotalSectors, g.PartStartLBA, g.PartSectors, g.Label, diskGUID, partGUID)

	// Partition array CRC.
	if want := refCRC32(gpt.EntryArray); want != gpt.ArrayCRC {
		t.Errorf("partition array CRC32 = 0x%08X, reference says 0x%08X", gpt.ArrayCRC, want)
	}
	if len(gpt.EntryArray) != 128*128 {
		t.Errorf("partition array is %d bytes, want %d", len(gpt.EntryArray), 128*128)
	}

	for _, h := range []struct {
		name string
		buf  []byte
		crc  uint32
		my   uint64
		alt  uint64
		arr  uint64
	}{
		{"primary", gpt.PrimaryHeader, gpt.HeaderCRC, 1, uint64(g.TotalSectors - 1), 2},
		{"backup", gpt.BackupHeader, gpt.BackupHdrCRC, uint64(g.TotalSectors - 1), 1, uint64(g.TotalSectors - 1 - 32)},
	} {
		if string(h.buf[0:8]) != "EFI PART" {
			t.Errorf("%s header signature = %q", h.name, string(h.buf[0:8]))
		}
		if v := binary.LittleEndian.Uint32(h.buf[8:12]); v != 0x00010000 {
			t.Errorf("%s revision = 0x%08X, want 0x00010000", h.name, v)
		}
		if v := binary.LittleEndian.Uint32(h.buf[12:16]); v != 92 {
			t.Errorf("%s header size = %d, want 92", h.name, v)
		}
		zeroed := append([]byte(nil), h.buf[:92]...)
		binary.LittleEndian.PutUint32(zeroed[16:20], 0)
		if want := refCRC32(zeroed); want != h.crc {
			t.Errorf("%s header CRC32 = 0x%08X, reference says 0x%08X", h.name, h.crc, want)
		}
		if v := binary.LittleEndian.Uint32(h.buf[16:20]); v != h.crc {
			t.Errorf("%s header stores CRC 0x%08X but BuildGPT reported 0x%08X", h.name, v, h.crc)
		}
		if v := binary.LittleEndian.Uint64(h.buf[24:32]); v != h.my {
			t.Errorf("%s MyLBA = %d, want %d", h.name, v, h.my)
		}
		if v := binary.LittleEndian.Uint64(h.buf[32:40]); v != h.alt {
			t.Errorf("%s AlternateLBA = %d, want %d", h.name, v, h.alt)
		}
		if v := binary.LittleEndian.Uint64(h.buf[72:80]); v != h.arr {
			t.Errorf("%s PartitionEntryLBA = %d, want %d", h.name, v, h.arr)
		}
		if v := binary.LittleEndian.Uint32(h.buf[80:84]); v != 128 {
			t.Errorf("%s NumberOfPartitionEntries = %d, want 128", h.name, v)
		}
		if v := binary.LittleEndian.Uint32(h.buf[84:88]); v != 128 {
			t.Errorf("%s SizeOfPartitionEntry = %d, want 128", h.name, v)
		}
		if v := binary.LittleEndian.Uint32(h.buf[88:92]); v != gpt.ArrayCRC {
			t.Errorf("%s array CRC field = 0x%08X, want 0x%08X", h.name, v, gpt.ArrayCRC)
		}
		if v := binary.LittleEndian.Uint64(h.buf[40:48]); v != 34 {
			t.Errorf("%s FirstUsableLBA = %d, want 34", h.name, v)
		}
	}

	// The first entry must be an EFI System Partition covering the FAT volume.
	e := gpt.EntryArray[0:128]
	if !bytes.Equal(e[0:16], espTypeGUID[:]) {
		t.Errorf("partition type GUID = %s, want the ESP GUID", formatGUID(*(*[16]byte)(e[0:16])))
	}
	if formatGUID(*(*[16]byte)(e[0:16])) != "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" {
		t.Errorf("ESP GUID renders as %s", formatGUID(*(*[16]byte)(e[0:16])))
	}
	if v := binary.LittleEndian.Uint64(e[32:40]); v != uint64(g.PartStartLBA) {
		t.Errorf("entry FirstLBA = %d, want %d", v, g.PartStartLBA)
	}
	if v := binary.LittleEndian.Uint64(e[40:48]); v != uint64(g.PartStartLBA+g.PartSectors-1) {
		t.Errorf("entry LastLBA = %d, want %d", v, g.PartStartLBA+g.PartSectors-1)
	}
	// Backup structures must not overlap the partition.
	if uint64(g.PartStartLBA+g.PartSectors-1) >= uint64(g.TotalSectors-33) {
		t.Errorf("partition overlaps the backup GPT: last LBA %d, backup array starts at %d",
			g.PartStartLBA+g.PartSectors-1, g.TotalSectors-33)
	}
	if gpt.ProtectiveMBR[450] != 0xEE {
		t.Errorf("protective MBR partition type = 0x%02X, want 0xEE", gpt.ProtectiveMBR[450])
	}
	if gpt.ProtectiveMBR[510] != 0x55 || gpt.ProtectiveMBR[511] != 0xAA {
		t.Error("protective MBR is missing its 0x55AA signature")
	}
}

func TestMBRPartitionEntry(t *testing.T) {
	g := mustGeometry(t, 64<<20, "mbr", "L", 0)
	m := BuildMBR(0xDEADBEEF, g.PartStartLBA, g.PartSectors)
	e := m[446:462]
	if e[0] != 0x00 {
		t.Errorf("boot flag = 0x%02X, want 0x00 (BootBuilder installs no bootloader)", e[0])
	}
	if e[4] != 0x0C {
		t.Errorf("partition type = 0x%02X, want 0x0C (FAT32 LBA)", e[4])
	}
	if v := binary.LittleEndian.Uint32(e[8:12]); v != uint32(g.PartStartLBA) {
		t.Errorf("start LBA = %d, want %d", v, g.PartStartLBA)
	}
	if v := binary.LittleEndian.Uint32(e[12:16]); v != uint32(g.PartSectors) {
		t.Errorf("sector count = %d, want %d", v, g.PartSectors)
	}
	if v := binary.LittleEndian.Uint32(m[440:444]); v != 0xDEADBEEF {
		t.Errorf("disk signature = 0x%08X, want 0xDEADBEEF", v)
	}
	if m[510] != 0x55 || m[511] != 0xAA {
		t.Error("MBR is missing its 0x55AA signature")
	}
	// CHS for LBA 2048 with 255 heads / 63 sectors: C=0, H=32, S=33.
	if got := lbaToCHS(2048); got != [3]byte{32, 33, 0} {
		t.Errorf("lbaToCHS(2048) = %v, want [32 33 0]", got)
	}
	// Anything past the CHS limit saturates.
	if got := lbaToCHS(1 << 30); got != [3]byte{0xFE, 0xFF, 0xFF} {
		t.Errorf("lbaToCHS(2^30) = %v, want the 1023/254/63 saturation value", got)
	}
}

// ---------------------------------------------------------------------------
// Round trip: write, then read back with the independent reader
// ---------------------------------------------------------------------------

func sha(b []byte) string {
	s := sha256.Sum256(b)
	return hex.EncodeToString(s[:])
}

func TestRoundTripContentAndLongNames(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")

	rnd := rand.New(rand.NewSource(20260811))
	big := make([]byte, 1_500_000) // spans many clusters at any cluster size
	rnd.Read(big)
	medium := make([]byte, 70_000)
	rnd.Read(medium)

	files := map[string][]byte{
		"EFI/BOOT/BOOTX64.EFI":    []byte("fake efi binary\n"),
		"EFI/BOOT/bootia32.efi":   medium,
		"autounattend.xml":        []byte("<unattend/>\n"),
		"README":                  []byte("plain\n"),
		"readme.txt":              []byte("lower\n"),
		"sources/boot.wim":        big,
		"sources/zero-length.dat": {},
		"sources/one-byte.bin":    {0x41},
		"Windows Setup Files/A Long File Name That Needs Several LFN Slots To Store.config":              []byte("one\n"),
		"Windows Setup Files/a long file name that needs several lfn slots to store.CONFIG":              []byte("two\n"),
		"Windows Setup Files/A+Long,File;Name=With[Invalid]Chars.txt":                                    []byte("three\n"),
		"a very deep/nest of/directories/four/five/six/seven/deep file with a very long name indeed.txt": []byte("deep\n"),
	}
	writeFixture(t, src, files)

	out := filepath.Join(dir, "usb.img")
	_, res := buildTo(t, src, out, 64<<20, "mbr", "ROUNDTRIP", 0)
	if res.UsedClusters <= 0 {
		t.Fatal("no clusters were used")
	}

	v, err := OpenImage(out)
	if err != nil {
		t.Fatalf("OpenImage: %v", err)
	}
	defer v.Close()

	if v.BPB.VolumeLabel != "ROUNDTRIP" {
		t.Errorf("volume label read back as %q", v.BPB.VolumeLabel)
	}
	if !v.BPB.Signature55AA || !v.BPB.BackupMatches {
		t.Errorf("boot sector signature=%v backup identical=%v", v.BPB.Signature55AA, v.BPB.BackupMatches)
	}
	if !v.FATsEqual {
		t.Error("the two FAT copies are not byte-identical")
	}
	if !v.FSInfo.LeadSigOK || !v.FSInfo.StructSig || !v.FSInfo.TrailSigOK {
		t.Errorf("FSInfo signatures: %+v", v.FSInfo)
	}
	if int64(v.FSInfo.FreeCount) != res.FreeClusters {
		t.Errorf("FSInfo free count %d, writer said %d", v.FSInfo.FreeCount, res.FreeClusters)
	}

	seen := map[string]bool{}
	err = v.Walk(func(e REntry) error {
		if e.IsDir {
			return nil
		}
		want, ok := files[e.Path]
		if !ok {
			t.Errorf("image contains unexpected file %q", e.Path)
			return nil
		}
		seen[e.Path] = true
		if int(e.Size) != len(want) {
			t.Errorf("%s: size %d, want %d", e.Path, e.Size, len(want))
			return nil
		}
		got, err := v.HashFile(e)
		if err != nil {
			t.Errorf("%s: %v", e.Path, err)
			return nil
		}
		if got != sha(want) {
			t.Errorf("%s: sha256 %s, want %s", e.Path, got, sha(want))
		}
		return nil
	})
	if err != nil {
		t.Fatalf("Walk: %v", err)
	}
	for p := range files {
		if !seen[p] {
			t.Errorf("file %q is missing from the image", p)
		}
	}

	// Long names must survive exactly, including case and punctuation.
	var names []string
	_ = v.Walk(func(e REntry) error {
		names = append(names, e.Path)
		return nil
	})
	sort.Strings(names)
	for _, want := range []string{
		"Windows Setup Files/A Long File Name That Needs Several LFN Slots To Store.config",
		"Windows Setup Files/a long file name that needs several lfn slots to store.CONFIG",
		"Windows Setup Files/A+Long,File;Name=With[Invalid]Chars.txt",
		"a very deep/nest of/directories/four/five/six/seven/deep file with a very long name indeed.txt",
		"readme.txt",
		"README",
	} {
		found := false
		for _, n := range names {
			if n == want {
				found = true
				break
			}
		}
		if !found {
			t.Errorf("long name %q did not survive; got %v", want, names)
		}
	}
}

func TestDeepNestedDirectories(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")
	const depth = 24
	parts := make([]string, 0, depth)
	for i := 0; i < depth; i++ {
		parts = append(parts, fmt.Sprintf("level %02d with a long name", i))
	}
	rel := strings.Join(parts, "/") + "/bottom of the well.txt"
	writeFixture(t, src, map[string][]byte{rel: []byte("bottom\n")})

	out := filepath.Join(dir, "deep.img")
	buildTo(t, src, out, 64<<20, "mbr", "DEEP", 0)

	v, err := OpenImage(out)
	if err != nil {
		t.Fatalf("OpenImage: %v", err)
	}
	defer v.Close()

	found := false
	dirs := 0
	err = v.Walk(func(e REntry) error {
		if e.IsDir {
			dirs++
			return nil
		}
		if e.Path == rel {
			found = true
			h, err := v.HashFile(e)
			if err != nil {
				return err
			}
			if h != sha([]byte("bottom\n")) {
				t.Errorf("deep file hash %s", h)
			}
		}
		return nil
	})
	if err != nil {
		t.Fatalf("Walk: %v", err)
	}
	if dirs != depth {
		t.Errorf("found %d directories, want %d", dirs, depth)
	}
	if !found {
		t.Errorf("did not find %q in the image", rel)
	}
}

func TestFileSpanningManyClusters(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")
	// 512-byte clusters on a 64 MiB volume, so this is 4001 clusters with a
	// partially-used last one.
	payload := make([]byte, 512*4000+7)
	rnd := rand.New(rand.NewSource(7))
	rnd.Read(payload)
	writeFixture(t, src, map[string][]byte{"install.wim": payload})

	out := filepath.Join(dir, "big.img")
	g, _ := buildTo(t, src, out, 64<<20, "mbr", "BIG", 0)
	if g.ClusterBytes != 512 {
		t.Fatalf("expected 512-byte clusters, got %d", g.ClusterBytes)
	}

	v, err := OpenImage(out)
	if err != nil {
		t.Fatalf("OpenImage: %v", err)
	}
	defer v.Close()

	var ent REntry
	if err := v.Walk(func(e REntry) error {
		if e.Name == "install.wim" {
			ent = e
		}
		return nil
	}); err != nil {
		t.Fatalf("Walk: %v", err)
	}
	if ent.Name == "" {
		t.Fatal("install.wim not found")
	}
	if int(ent.Size) != len(payload) {
		t.Errorf("size %d, want %d", ent.Size, len(payload))
	}
	if ent.Clusters != 4001 {
		t.Errorf("chain length %d clusters, want 4001", ent.Clusters)
	}
	chain, err := v.Chain(ent.Cluster)
	if err != nil {
		t.Fatalf("Chain: %v", err)
	}
	for i := 1; i < len(chain); i++ {
		if chain[i] != chain[i-1]+1 {
			t.Fatalf("chain is not contiguous at index %d: %d then %d", i, chain[i-1], chain[i])
		}
	}
	h, err := v.HashFile(ent)
	if err != nil {
		t.Fatalf("HashFile: %v", err)
	}
	if h != sha(payload) {
		t.Errorf("content hash %s, want %s", h, sha(payload))
	}
}

// ---------------------------------------------------------------------------
// Cluster chain integrity and FAT copy equality, read straight from the bytes
// ---------------------------------------------------------------------------

func TestFATChainsAndCopiesOnDisk(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")
	rnd := rand.New(rand.NewSource(99))
	files := map[string][]byte{}
	for i := 0; i < 20; i++ {
		b := make([]byte, 1+rnd.Intn(9000))
		rnd.Read(b)
		files[fmt.Sprintf("dir%02d/file %02d.bin", i%4, i)] = b
	}
	writeFixture(t, src, files)

	out := filepath.Join(dir, "chains.img")
	g, res := buildTo(t, src, out, 64<<20, "mbr", "CHAINS", 0)

	raw, err := os.ReadFile(out)
	if err != nil {
		t.Fatalf("read image: %v", err)
	}
	fatOff := (g.PartStartLBA + g.ReservedSectors) * sectorSize
	fat1 := raw[fatOff : fatOff+g.FATBytes]
	fat2 := raw[fatOff+g.FATBytes : fatOff+2*g.FATBytes]
	if !bytes.Equal(fat1, fat2) {
		t.Fatal("FAT copy 1 and FAT copy 2 differ")
	}
	if v := binary.LittleEndian.Uint32(fat1[0:4]); v != 0x0FFFFFF8 {
		t.Errorf("FAT[0] = 0x%08X, want 0x0FFFFFF8", v)
	}
	if v := binary.LittleEndian.Uint32(fat1[4:8]); v != 0x0FFFFFFF {
		t.Errorf("FAT[1] = 0x%08X, want 0x0FFFFFFF", v)
	}

	// Every allocated cluster must appear in exactly one chain and every chain
	// must end in an EOC marker.
	entry := func(c uint32) uint32 {
		return binary.LittleEndian.Uint32(fat1[c*4:c*4+4]) & 0x0FFFFFFF
	}
	v, err := OpenImage(out)
	if err != nil {
		t.Fatalf("OpenImage: %v", err)
	}
	defer v.Close()

	owner := map[uint32]string{}
	var check func(cluster uint32, who string)
	check = func(cluster uint32, who string) {
		if cluster < 2 {
			return
		}
		for {
			if prev, dup := owner[cluster]; dup {
				t.Fatalf("cluster %d is claimed by both %q and %q", cluster, prev, who)
			}
			owner[cluster] = who
			n := entry(cluster)
			if n >= 0x0FFFFFF8 {
				return
			}
			if n < 2 {
				t.Fatalf("chain for %q runs into FAT value 0x%08X at cluster %d", who, n, cluster)
			}
			cluster = n
		}
	}
	check(v.BPB.RootCluster, "/")
	if err := v.Walk(func(e REntry) error {
		check(e.Cluster, e.Path)
		return nil
	}); err != nil {
		t.Fatalf("Walk: %v", err)
	}
	if int64(len(owner)) != res.UsedClusters {
		t.Errorf("chains cover %d clusters, the writer allocated %d", len(owner), res.UsedClusters)
	}
	// Every cluster past the allocated run must read as free.
	for c := uint32(res.UsedClusters) + 2; c < uint32(res.UsedClusters)+2+64; c++ {
		if entry(c) != 0 {
			t.Errorf("cluster %d past the allocated run is 0x%08X, want 0 (free)", c, entry(c))
		}
	}
}

// ---------------------------------------------------------------------------
// Capacity refusal and output safety
// ---------------------------------------------------------------------------

func TestOverCapacityIsRejectedWithAnExactByteFigure(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")
	// A 64 MiB volume has 127006 clusters of 512 bytes = 65027072 usable bytes.
	// Two 31 MiB files plus directories cannot fit.
	rnd := rand.New(rand.NewSource(3))
	a := make([]byte, 33<<20)
	b := make([]byte, 33<<20)
	rnd.Read(a[:1024])
	rnd.Read(b[:1024])
	writeFixture(t, src, map[string][]byte{"a.bin": a, "b.bin": b})

	out := filepath.Join(dir, "toosmall.img")
	g := mustGeometry(t, 64<<20, "mbr", "SMALL", 0)
	root, scan, err := ScanTree(src)
	if err != nil {
		t.Fatalf("ScanTree: %v", err)
	}
	_, err = WriteImage(out, g, root, scan, false, time.Now())
	if err == nil {
		t.Fatal("expected an over-capacity refusal")
	}
	if !errors.Is(err, ErrCapacity) {
		t.Errorf("error %v does not wrap ErrCapacity", err)
	}
	if !strings.Contains(err.Error(), "short by exactly") {
		t.Errorf("error should state the exact shortfall in bytes: %v", err)
	}
	// The exact figure must be right: needed minus available, in bytes.
	req := ComputeLayout(root, g.ClusterBytes, true)
	want := fmt.Sprintf("short by exactly %d bytes", (req.TotalCluster-g.ClusterCount)*g.ClusterBytes)
	if !strings.Contains(err.Error(), want) {
		t.Errorf("error %q does not contain %q", err.Error(), want)
	}
	if _, statErr := os.Stat(out); !os.IsNotExist(statErr) {
		t.Error("a refused build must not leave an output file behind")
	}
}

func TestRefusesToOverwriteWithoutForce(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")
	writeFixture(t, src, map[string][]byte{"a.txt": []byte("hello\n")})
	out := filepath.Join(dir, "exists.img")
	if err := os.WriteFile(out, []byte("PRECIOUS"), 0o644); err != nil {
		t.Fatal(err)
	}
	g := mustGeometry(t, 64<<20, "mbr", "L", 0)
	root, scan, err := ScanTree(src)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := WriteImage(out, g, root, scan, false, time.Now()); err == nil {
		t.Fatal("expected a refusal to overwrite")
	}
	got, _ := os.ReadFile(out)
	if string(got) != "PRECIOUS" {
		t.Fatal("the existing file was modified despite the refusal")
	}
	if _, err := WriteImage(out, g, root, scan, true, time.Now()); err != nil {
		t.Fatalf("--force build failed: %v", err)
	}
	fi, err := os.Stat(out)
	if err != nil {
		t.Fatal(err)
	}
	if fi.Size() != g.ImageBytes {
		t.Errorf("forced build produced %d bytes, want %d", fi.Size(), g.ImageBytes)
	}
}

// ---------------------------------------------------------------------------
// Small units
// ---------------------------------------------------------------------------

func TestParseSize(t *testing.T) {
	for _, tc := range []struct {
		in      string
		want    int64
		wantErr bool
	}{
		{"8GB", 8 << 30, false},
		{"8G", 8 << 30, false},
		{"8GiB", 8 << 30, false},
		{"512MB", 512 << 20, false},
		{"1.5G", 1610612736, false},
		{"2TiB", 2 << 40, false},
		{"1_048_576", 1 << 20, false},
		{"1,048,576", 1 << 20, false},
		{"4096", 4096, false},
		{"", 0, true},
		{"banana", 0, true},
		{"-4GB", 0, true},
	} {
		got, err := parseSize(tc.in)
		if tc.wantErr {
			if err == nil {
				t.Errorf("parseSize(%q) should have failed", tc.in)
			}
			continue
		}
		if err != nil {
			t.Errorf("parseSize(%q): %v", tc.in, err)
			continue
		}
		if got != tc.want {
			t.Errorf("parseSize(%q) = %d, want %d", tc.in, got, tc.want)
		}
	}
}

func TestReorderFlags(t *testing.T) {
	got := reorderFlags([]string{"usb.img", "--against", "src", "--json"}, valueFlags)
	want := []string{"--against", "src", "--json", "usb.img"}
	if strings.Join(got, "|") != strings.Join(want, "|") {
		t.Errorf("reorderFlags = %v, want %v", got, want)
	}
}

func TestHumanBytes(t *testing.T) {
	for _, tc := range []struct {
		in   int64
		want string
	}{
		{0, "0 B"},
		{1023, "1023 B"},
		{1024, "1.0 KiB"},
		{65027072, "62.0 MiB"},
		{8 << 30, "8.0 GiB"},
	} {
		if got := humanBytes(tc.in); got != tc.want {
			t.Errorf("humanBytes(%d) = %q, want %q", tc.in, got, tc.want)
		}
	}
}

func TestFATTimePacking(t *testing.T) {
	tm := time.Date(2026, 8, 11, 14, 35, 47, 0, time.Local)
	d, hm, tenth := fatTime(tm)
	if y := 1980 + int(d>>9); y != 2026 {
		t.Errorf("year decoded as %d", y)
	}
	if m := int(d>>5) & 0x0F; m != 8 {
		t.Errorf("month decoded as %d", m)
	}
	if day := int(d) & 0x1F; day != 11 {
		t.Errorf("day decoded as %d", day)
	}
	if h := int(hm >> 11); h != 14 {
		t.Errorf("hour decoded as %d", h)
	}
	if mi := int(hm>>5) & 0x3F; mi != 35 {
		t.Errorf("minute decoded as %d", mi)
	}
	if s := (int(hm) & 0x1F) * 2; s != 46 {
		t.Errorf("second decoded as %d, want 46 (FAT stores two-second units)", s)
	}
	if tenth != 100 {
		t.Errorf("creation tenth = %d, want 100 (the odd second)", tenth)
	}
	// Anything before the FAT epoch clamps to 1980-01-01.
	d2, hm2, _ := fatTime(time.Date(1969, 7, 20, 20, 17, 40, 0, time.UTC))
	if d2 != (1<<5)|1 || hm2 != 0 {
		t.Errorf("pre-epoch time packed as date=0x%04X time=0x%04X", d2, hm2)
	}
}

func TestValidateLabel(t *testing.T) {
	if got, err := ValidateLabel("winpe"); err != nil || got != "WINPE" {
		t.Errorf("ValidateLabel(winpe) = %q, %v", got, err)
	}
	if _, err := ValidateLabel("THIS LABEL IS TOO LONG"); err == nil {
		t.Error("an over-long label must be rejected")
	}
	if _, err := ValidateLabel("BAD*LABEL"); err == nil {
		t.Error("a label with an invalid character must be rejected")
	}
}

func TestVolumeLabelEntryIsWrittenToRoot(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")
	writeFixture(t, src, map[string][]byte{"a.txt": []byte("x")})
	out := filepath.Join(dir, "label.img")
	buildTo(t, src, out, 64<<20, "mbr", "MYLABEL", 0)

	v, err := OpenImage(out)
	if err != nil {
		t.Fatal(err)
	}
	defer v.Close()
	lab, err := v.VolumeLabelEntry()
	if err != nil {
		t.Fatal(err)
	}
	if lab != "MYLABEL" {
		t.Errorf("root label entry = %q, want %q", lab, "MYLABEL")
	}
	if v.BPB.VolumeLabel != "MYLABEL" {
		t.Errorf("BPB label = %q", v.BPB.VolumeLabel)
	}
}

func TestGPTImageRoundTrip(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src")
	writeFixture(t, src, map[string][]byte{
		"EFI/BOOT/BOOTX64.EFI": []byte("efi\n"),
		"a file.txt":           []byte("content\n"),
	})
	out := filepath.Join(dir, "gpt.img")
	buildTo(t, src, out, 512<<20, "gpt", "ESP", 0)

	v, err := OpenImage(out)
	if err != nil {
		t.Fatalf("OpenImage: %v", err)
	}
	defer v.Close()
	if v.Scheme != "gpt" {
		t.Fatalf("scheme read back as %q", v.Scheme)
	}
	if v.GPT == nil {
		t.Fatal("no GPT was parsed")
	}
	if !v.GPT.HeaderCRCOK || !v.GPT.ArrayCRCOK || !v.GPT.BackupCRCOK || !v.GPT.BackupArrayMtch {
		t.Errorf("GPT checksum verdicts: %+v", v.GPT)
	}
	if len(v.Partitions) != 1 || v.Partitions[0].Type != "EFI System Partition" {
		t.Errorf("partitions = %+v", v.Partitions)
	}
	n := 0
	if err := v.Walk(func(e REntry) error {
		if !e.IsDir {
			n++
		}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	if n != 2 {
		t.Errorf("found %d files, want 2", n)
	}
}

func TestClusterSizeOverride(t *testing.T) {
	g := mustGeometry(t, 4<<30, "mbr", "L", 16384)
	if g.ClusterBytes != 16384 || g.SectorsPerCluster != 32 {
		t.Errorf("override gave cluster %d bytes / %d sectors", g.ClusterBytes, g.SectorsPerCluster)
	}
	if g.ClusterFromTable {
		t.Error("an overridden cluster size must not be reported as coming from the MS table")
	}
	if _, err := PlanGeometry(4<<30, "mbr", "L", 3000); err == nil {
		t.Error("a non-power-of-two cluster size must be rejected")
	}
	if _, err := PlanGeometry(4<<30, "mbr", "L", 1<<20); err == nil {
		t.Error("a cluster size above 64 KiB must be rejected")
	}
}