package main

import (
	"errors"
	"fmt"
)

// ---------------------------------------------------------------------------
// Volume geometry
//
// Everything here is arithmetic on sector counts. Nothing in this file touches
// a device; the only consumer is the image-file writer.
// ---------------------------------------------------------------------------

const (
	sectorSize = 512

	// The partition starts one mebibyte in. 2048 sectors is the alignment every
	// modern partitioner uses and it leaves room for a GPT partition array.
	partitionStartLBA = 2048

	// FAT32 fixes 32 reserved sectors: boot sector, FSInfo, backup boot sector
	// at offset 6, backup FSInfo at offset 7, then padding to the FAT.
	reservedSectors = 32

	numFATs = 2

	// FAT32's own definition: a volume with fewer than 65525 clusters is FAT16
	// or FAT12, not FAT32. Microsoft's fatgen103 spec, section "FAT Type
	// Determination".
	minFAT32Clusters = 65525

	// 0x0FFFFFF6 is the last usable cluster number; values above it are the
	// reserved/bad/EOC range.
	maxFAT32Clusters = 0x0FFFFFF5 - 1

	// GPT reserves LBA 1 for the header and 32 sectors for a 128 x 128 byte
	// partition array, mirrored at the end of the disk.
	gptEntryCount   = 128
	gptEntryBytes   = 128
	gptArraySectors = (gptEntryCount * gptEntryBytes) / sectorSize // 32
)

// Geometry is the complete on-disk layout of one image.
type Geometry struct {
	Scheme            string `json:"scheme"`
	ImageBytes        int64  `json:"image_bytes"`
	ImageHuman        string `json:"image_human"`
	SectorSize        int64  `json:"sector_size"`
	TotalSectors      int64  `json:"total_sectors"`
	PartStartLBA      int64  `json:"partition_start_lba"`
	PartSectors       int64  `json:"partition_sectors"`
	PartBytes         int64  `json:"partition_bytes"`
	SectorsPerCluster int64  `json:"sectors_per_cluster"`
	ClusterBytes      int64  `json:"cluster_bytes"`
	ReservedSectors   int64  `json:"reserved_sectors"`
	NumFATs           int64  `json:"num_fats"`
	FATSectors        int64  `json:"fat_sectors_each"`
	FATBytes          int64  `json:"fat_bytes_each"`
	DataStartLBA      int64  `json:"data_start_lba"` // absolute, not partition-relative
	DataSectors       int64  `json:"data_sectors"`
	ClusterCount      int64  `json:"cluster_count"`
	UsableBytes       int64  `json:"usable_bytes"`
	UsableHuman       string `json:"usable_human"`
	SlackSectors      int64  `json:"slack_sectors"`
	SlackBytes        int64  `json:"slack_bytes"`
	RootCluster       int64  `json:"root_cluster"`
	Label             string `json:"volume_label"`
	ClusterFromTable  bool   `json:"cluster_size_from_ms_table"`
}

// msSecPerClusEntry is one row of Microsoft's DSKSZTOSECPERCLUS table for
// FAT32 (fatgen103, "Determining FAT type"): the largest volume size in
// sectors for which the given sectors-per-cluster value is used.
type msSecPerClusEntry struct {
	MaxSectors int64
	SecPerClus int64
}

// msFAT32Table is reproduced verbatim from the specification. A zero
// SecPerClus means "FAT32 is not valid for a volume this small".
var msFAT32Table = []msSecPerClusEntry{
	{66600, 0},       // up to 32.5 MB - FAT32 not permitted
	{532480, 1},      // up to 260 MB  - 512 byte clusters
	{16777216, 8},    // up to 8 GB    - 4 KiB clusters
	{33554432, 16},   // up to 16 GB   - 8 KiB clusters
	{67108864, 32},   // up to 32 GB   - 16 KiB clusters
	{0xFFFFFFFF, 64}, // above 32 GB   - 32 KiB clusters
}

// secPerClusForSectors returns the sectors-per-cluster Microsoft's table
// prescribes for a volume of the given sector count. Zero means FAT32 is not
// valid at that size.
func secPerClusForSectors(volSectors int64) int64 {
	for _, e := range msFAT32Table {
		if volSectors <= e.MaxSectors {
			return e.SecPerClus
		}
	}
	return 64
}

var errTooSmall = errors.New("volume too small for FAT32")

// PlanGeometry computes the layout for an image of imageBytes total size.
// clusterOverride, when non-zero, is a cluster size in BYTES that replaces the
// value from Microsoft's table.
func PlanGeometry(imageBytes int64, scheme, label string, clusterOverride int64) (Geometry, error) {
	var g Geometry
	switch scheme {
	case "mbr", "gpt":
	default:
		return g, fmt.Errorf("unknown partition scheme %q (want mbr or gpt)", scheme)
	}
	if imageBytes%sectorSize != 0 {
		imageBytes -= imageBytes % sectorSize
	}
	total := imageBytes / sectorSize
	if total <= partitionStartLBA+gptArraySectors+2 {
		return g, fmt.Errorf("%w: image of %s leaves no room for a partition", errTooSmall, humanBytes(imageBytes))
	}

	partStart := int64(partitionStartLBA)
	partEndExclusive := total
	if scheme == "gpt" {
		// Backup partition array (32 sectors) plus the backup header (1) sit at
		// the very end of the image and must not overlap the partition.
		partEndExclusive = total - gptArraySectors - 1
	}
	partSectors := partEndExclusive - partStart
	if partSectors <= 0 {
		return g, fmt.Errorf("%w: image of %s leaves no room for a partition", errTooSmall, humanBytes(imageBytes))
	}

	spc := secPerClusForSectors(partSectors)
	fromTable := true
	if clusterOverride != 0 {
		if clusterOverride < sectorSize || clusterOverride > 65536 || clusterOverride&(clusterOverride-1) != 0 {
			return g, fmt.Errorf("cluster size %d is not a power of two between %d and 65536", clusterOverride, sectorSize)
		}
		spc = clusterOverride / sectorSize
		fromTable = false
	}
	if spc == 0 {
		return g, fmt.Errorf("%w: a FAT32 volume needs more than %d sectors (%s); this partition has %d sectors (%s)",
			errTooSmall, msFAT32Table[0].MaxSectors, humanBytes(msFAT32Table[0].MaxSectors*sectorSize),
			partSectors, humanBytes(partSectors*sectorSize))
	}

	fatSectors, clusters, err := solveFATSize(partSectors, spc)
	if err != nil {
		return g, err
	}
	if clusters < minFAT32Clusters {
		return g, fmt.Errorf("%w: layout yields %d clusters of %s, but FAT32 requires at least %d; use a larger --size or a smaller --cluster-size",
			errTooSmall, clusters, humanBytes(spc*sectorSize), minFAT32Clusters)
	}
	if clusters > maxFAT32Clusters {
		return g, fmt.Errorf("layout yields %d clusters, above the FAT32 maximum of %d; use a larger --cluster-size", clusters, maxFAT32Clusters)
	}

	dataStart := partStart + reservedSectors + numFATs*fatSectors
	dataSectors := clusters * spc
	used := reservedSectors + numFATs*fatSectors + dataSectors

	g = Geometry{
		Scheme:            scheme,
		ImageBytes:        imageBytes,
		ImageHuman:        humanBytes(imageBytes),
		SectorSize:        sectorSize,
		TotalSectors:      total,
		PartStartLBA:      partStart,
		PartSectors:       partSectors,
		PartBytes:         partSectors * sectorSize,
		SectorsPerCluster: spc,
		ClusterBytes:      spc * sectorSize,
		ReservedSectors:   reservedSectors,
		NumFATs:           numFATs,
		FATSectors:        fatSectors,
		FATBytes:          fatSectors * sectorSize,
		DataStartLBA:      dataStart,
		DataSectors:       dataSectors,
		ClusterCount:      clusters,
		UsableBytes:       clusters * spc * sectorSize,
		SlackSectors:      partSectors - used,
		SlackBytes:        (partSectors - used) * sectorSize,
		RootCluster:       2,
		Label:             label,
		ClusterFromTable:  fromTable,
	}
	g.UsableHuman = humanBytes(g.UsableBytes)
	return g, nil
}

// solveFATSize finds the smallest FAT size, in sectors, that can address every
// data cluster that fits in the space left over after the FATs themselves.
//
// Each FAT32 entry is 4 bytes, so one FAT sector addresses 128 clusters. The
// two quantities are mutually dependent - a bigger FAT leaves fewer clusters,
// which needs a smaller FAT - so this iterates to a fixed point. It always
// terminates because fatSectors only ever increases and is bounded.
func solveFATSize(partSectors, spc int64) (fatSectors, clusters int64, err error) {
	const entriesPerSector = sectorSize / 4 // 128

	avail := partSectors - reservedSectors
	if avail <= 0 {
		return 0, 0, fmt.Errorf("%w: %d sectors cannot hold %d reserved sectors", errTooSmall, partSectors, int64(reservedSectors))
	}

	// Microsoft's own first approximation from fatgen103, then refine.
	tmp2 := (256*spc + numFATs) / 2
	fatSectors = (avail + tmp2 - 1) / tmp2
	if fatSectors < 1 {
		fatSectors = 1
	}

	for i := 0; i < 64; i++ {
		dataSectors := avail - numFATs*fatSectors
		if dataSectors <= 0 {
			return 0, 0, fmt.Errorf("%w: no data sectors remain after two %d-sector FATs", errTooSmall, fatSectors)
		}
		clusters = dataSectors / spc
		// Clusters 0 and 1 are reserved entries in the FAT, so the table must
		// hold clusters+2 entries.
		need := (clusters + 2 + entriesPerSector - 1) / entriesPerSector
		if need > fatSectors {
			fatSectors = need
			continue
		}
		// Shrink the FAT if it is larger than necessary, but never below what
		// the resulting cluster count needs.
		for fatSectors > need && fatSectors > 1 {
			trial := fatSectors - 1
			d := avail - numFATs*trial
			c := d / spc
			n := (c + 2 + entriesPerSector - 1) / entriesPerSector
			if n > trial {
				break
			}
			fatSectors = trial
			clusters = c
			need = n
		}
		return fatSectors, clusters, nil
	}
	return 0, 0, errors.New("FAT size did not converge")
}