package main

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"hash/crc32"
)

// ---------------------------------------------------------------------------
// Partition tables: MBR, and GPT with a protective MBR.
// ---------------------------------------------------------------------------

const (
	partTypeFAT32LBA  = 0x0C
	partTypeGPTProtec = 0xEE

	// Classic BIOS geometry used for the CHS fields. Anything past the CHS
	// limit is filled with the conventional 1023/254/63 saturation value.
	chsHeads          = 255
	chsSectorsPerTrak = 63
)

// espTypeGUID is the EFI System Partition type, C12A7328-F81F-11D2-BA4B-
// 00A0C93EC93B, already laid out in GPT's mixed-endian order (first three
// fields little-endian, last two big-endian).
var espTypeGUID = [16]byte{
	0x28, 0x73, 0x2A, 0xC1,
	0x1F, 0xF8,
	0xD2, 0x11,
	0xBA, 0x4B,
	0x00, 0xA0, 0xC9, 0x3E, 0xC9, 0x3B,
}

// lbaToCHS converts an LBA to the packed 3-byte CHS triple an MBR entry wants:
// head, then sector in the low 6 bits of the next byte with cylinder bits 8-9
// in its top 2 bits, then cylinder bits 0-7.
func lbaToCHS(lba int64) [3]byte {
	c := lba / (chsHeads * chsSectorsPerTrak)
	h := (lba / chsSectorsPerTrak) % chsHeads
	s := lba%chsSectorsPerTrak + 1
	if c > 1023 {
		return [3]byte{0xFE, 0xFF, 0xFF} // 1023 / 254 / 63
	}
	return [3]byte{
		byte(h),
		byte(s) | byte((c>>2)&0xC0),
		byte(c & 0xFF),
	}
}

// mbrPartitionEntry builds the 16-byte MBR partition record.
func mbrPartitionEntry(bootable bool, ptype byte, startLBA, sectors int64) [16]byte {
	var e [16]byte
	if bootable {
		e[0] = 0x80
	}
	first := lbaToCHS(startLBA)
	last := lbaToCHS(startLBA + sectors - 1)
	copy(e[1:4], first[:])
	e[4] = ptype
	copy(e[5:8], last[:])
	binary.LittleEndian.PutUint32(e[8:12], uint32(clampU32(startLBA)))
	binary.LittleEndian.PutUint32(e[12:16], uint32(clampU32(sectors)))
	return e
}

func clampU32(v int64) int64 {
	if v > 0xFFFFFFFF {
		return 0xFFFFFFFF
	}
	if v < 0 {
		return 0
	}
	return v
}

// BuildMBR returns sector 0 for a plain MBR-partitioned image: no boot code,
// one FAT32 LBA partition, the 0x55AA signature.
func BuildMBR(diskSig uint32, startLBA, sectors int64) []byte {
	s := make([]byte, sectorSize)
	binary.LittleEndian.PutUint32(s[440:444], diskSig)
	e := mbrPartitionEntry(false, partTypeFAT32LBA, startLBA, sectors)
	copy(s[446:462], e[:])
	s[510], s[511] = 0x55, 0xAA
	return s
}

// BuildProtectiveMBR returns sector 0 for a GPT image: a single 0xEE partition
// covering the whole disk so that MBR-only tools see no free space.
func BuildProtectiveMBR(totalSectors int64) []byte {
	s := make([]byte, sectorSize)
	e := mbrPartitionEntry(false, partTypeGPTProtec, 1, totalSectors-1)
	// The protective entry's CHS fields are specified as 0x000200 to
	// 0xFFFFFF, not computed from the LBA.
	e[1], e[2], e[3] = 0x00, 0x02, 0x00
	e[5], e[6], e[7] = 0xFF, 0xFF, 0xFF
	copy(s[446:462], e[:])
	s[510], s[511] = 0x55, 0xAA
	return s
}

// GPTEntry is one 128-byte partition record.
type GPTEntry struct {
	TypeGUID   [16]byte
	UniqueGUID [16]byte
	FirstLBA   uint64
	LastLBA    uint64
	Attributes uint64
	Name       string // UTF-16LE on disk, 36 units max
}

func (e GPTEntry) bytes() []byte {
	b := make([]byte, gptEntryBytes)
	copy(b[0:16], e.TypeGUID[:])
	copy(b[16:32], e.UniqueGUID[:])
	binary.LittleEndian.PutUint64(b[32:40], e.FirstLBA)
	binary.LittleEndian.PutUint64(b[40:48], e.LastLBA)
	binary.LittleEndian.PutUint64(b[48:56], e.Attributes)
	for i, r := range []rune(e.Name) {
		if i >= 35 || r > 0xFFFF {
			break
		}
		binary.LittleEndian.PutUint16(b[56+i*2:58+i*2], uint16(r))
	}
	return b
}

// GPTImage carries every GPT structure for one image.
type GPTImage struct {
	ProtectiveMBR []byte
	PrimaryHeader []byte
	BackupHeader  []byte
	EntryArray    []byte
	PrimaryArrLBA int64
	BackupArrLBA  int64
	BackupHdrLBA  int64
	DiskGUID      [16]byte
	PartGUID      [16]byte
	HeaderCRC     uint32
	BackupHdrCRC  uint32
	ArrayCRC      uint32
}

// newGUID returns a random RFC 4122 version-4 UUID in GPT's mixed-endian
// on-disk order.
func newGUID() ([16]byte, error) {
	var g [16]byte
	if _, err := rand.Read(g[:]); err != nil {
		return g, fmt.Errorf("cannot read random bytes for a GUID: %w", err)
	}
	// Version 4 and RFC 4122 variant, applied to the big-endian logical form;
	// the mixed-endian layout puts those bytes at 7 and 8.
	g[7] = (g[7] & 0x0F) | 0x40
	g[8] = (g[8] & 0x3F) | 0x80
	return g, nil
}

// BuildGPT constructs the protective MBR, both headers and the partition array
// for a single partition spanning [startLBA, startLBA+sectors).
//
// Both CRC32 values are the standard IEEE polynomial: the partition-array CRC
// covers the whole 128 x 128 byte array, and the header CRC covers the first
// 92 bytes of the header with its own CRC field zeroed.
func BuildGPT(totalSectors, startLBA, sectors int64, name string, diskGUID, partGUID [16]byte) *GPTImage {
	entries := make([]byte, gptEntryCount*gptEntryBytes)
	e := GPTEntry{
		TypeGUID:   espTypeGUID,
		UniqueGUID: partGUID,
		FirstLBA:   uint64(startLBA),
		LastLBA:    uint64(startLBA + sectors - 1),
		Name:       name,
	}
	copy(entries[0:gptEntryBytes], e.bytes())
	arrayCRC := crc32.ChecksumIEEE(entries)

	backupHdrLBA := totalSectors - 1
	backupArrLBA := totalSectors - 1 - gptArraySectors
	firstUsable := int64(2 + gptArraySectors)
	lastUsable := backupArrLBA - 1

	hdr := func(myLBA, altLBA, entryLBA int64) ([]byte, uint32) {
		h := make([]byte, sectorSize)
		copy(h[0:8], []byte("EFI PART"))
		binary.LittleEndian.PutUint32(h[8:12], 0x00010000) // revision 1.0
		binary.LittleEndian.PutUint32(h[12:16], 92)        // header size
		binary.LittleEndian.PutUint32(h[16:20], 0)         // CRC placeholder
		binary.LittleEndian.PutUint32(h[20:24], 0)         // reserved
		binary.LittleEndian.PutUint64(h[24:32], uint64(myLBA))
		binary.LittleEndian.PutUint64(h[32:40], uint64(altLBA))
		binary.LittleEndian.PutUint64(h[40:48], uint64(firstUsable))
		binary.LittleEndian.PutUint64(h[48:56], uint64(lastUsable))
		copy(h[56:72], diskGUID[:])
		binary.LittleEndian.PutUint64(h[72:80], uint64(entryLBA))
		binary.LittleEndian.PutUint32(h[80:84], gptEntryCount)
		binary.LittleEndian.PutUint32(h[84:88], gptEntryBytes)
		binary.LittleEndian.PutUint32(h[88:92], arrayCRC)
		c := crc32.ChecksumIEEE(h[0:92])
		binary.LittleEndian.PutUint32(h[16:20], c)
		return h, c
	}

	primary, pc := hdr(1, backupHdrLBA, 2)
	backup, bc := hdr(backupHdrLBA, 1, backupArrLBA)

	return &GPTImage{
		ProtectiveMBR: BuildProtectiveMBR(totalSectors),
		PrimaryHeader: primary,
		BackupHeader:  backup,
		EntryArray:    entries,
		PrimaryArrLBA: 2,
		BackupArrLBA:  backupArrLBA,
		BackupHdrLBA:  backupHdrLBA,
		DiskGUID:      diskGUID,
		PartGUID:      partGUID,
		HeaderCRC:     pc,
		BackupHdrCRC:  bc,
		ArrayCRC:      arrayCRC,
	}
}

// formatGUID renders a mixed-endian on-disk GUID in the canonical text form.
func formatGUID(g [16]byte) string {
	return fmt.Sprintf("%02X%02X%02X%02X-%02X%02X-%02X%02X-%02X%02X-%02X%02X%02X%02X%02X%02X",
		g[3], g[2], g[1], g[0],
		g[5], g[4],
		g[7], g[6],
		g[8], g[9],
		g[10], g[11], g[12], g[13], g[14], g[15])
}