package main

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

// ---------------------------------------------------------------------------
// The FAT32 writer.
//
// Everything below emits bytes into an ordinary FILE. There is deliberately no
// code path here, or anywhere else in this program, that opens a block device.
// ---------------------------------------------------------------------------

const (
	fatEOC        = 0x0FFFFFFF
	fatEntryMask  = 0x0FFFFFFF
	maxFAT32File  = 0xFFFFFFFF // 4 GiB - 1: the largest size a FAT32 entry records
	oemName       = "MSWIN4.1"
	fsTypeLabel   = "FAT32   "
	backupBootSec = 6
	fsInfoSector  = 1
)

// SrcNode is one entry of the scanned source tree.
type SrcNode struct {
	Name     string
	IsDir    bool
	Size     int64
	Mod      time.Time
	SrcPath  string
	Children []*SrcNode

	Short    [11]byte
	NeedLFN  bool
	Start    uint32
	Clusters int64
	DirBytes int64
}

// ScanReport summarises what the source walk found.
type ScanReport struct {
	Files     int64    `json:"files"`
	Dirs      int64    `json:"dirs"`
	DataBytes int64    `json:"data_bytes"`
	Skipped   []string `json:"skipped,omitempty"`
}

// ScanTree reads root recursively and assigns short names. Symbolic links,
// devices, sockets and FIFOs are skipped and reported; nothing is followed out
// of the tree and nothing is written.
func ScanTree(root string) (*SrcNode, *ScanReport, error) {
	info, err := os.Stat(root)
	if err != nil {
		return nil, nil, fmt.Errorf("cannot read source %s: %w", root, err)
	}
	if !info.IsDir() {
		return nil, nil, fmt.Errorf("source %s is not a directory", root)
	}
	rep := &ScanReport{}
	n := &SrcNode{Name: "", IsDir: true, SrcPath: root, Mod: info.ModTime()}
	if err := scanInto(n, root, "", rep); err != nil {
		return nil, nil, err
	}
	return n, rep, nil
}

func scanInto(parent *SrcNode, dir, rel string, rep *ScanReport) error {
	ents, err := os.ReadDir(dir)
	if err != nil {
		return fmt.Errorf("cannot read directory %s: %w", dir, err)
	}
	sort.Slice(ents, func(i, j int) bool {
		a, b := ents[i].Name(), ents[j].Name()
		la, lb := strings.ToUpper(a), strings.ToUpper(b)
		if la != lb {
			return la < lb
		}
		return a < b
	})
	taken := map[string]bool{}
	for _, de := range ents {
		full := filepath.Join(dir, de.Name())
		childRel := path.Join(rel, de.Name())
		if !de.IsDir() && !de.Type().IsRegular() {
			rep.Skipped = append(rep.Skipped, childRel+" ("+de.Type().String()+")")
			continue
		}
		fi, err := de.Info()
		if err != nil {
			return fmt.Errorf("cannot stat %s: %w", full, err)
		}
		short, needLFN, err := GenerateShortName(de.Name(), taken)
		if err != nil {
			return fmt.Errorf("%s: %w", childRel, err)
		}
		n := &SrcNode{
			Name:    de.Name(),
			IsDir:   de.IsDir(),
			Mod:     fi.ModTime(),
			SrcPath: full,
			Short:   short,
			NeedLFN: needLFN,
		}
		if de.IsDir() {
			rep.Dirs++
			if err := scanInto(n, full, childRel, rep); err != nil {
				return err
			}
		} else {
			if fi.Size() > maxFAT32File {
				return fmt.Errorf("%s is %s; FAT32 cannot store a file larger than %s",
					childRel, humanBytes(fi.Size()), humanBytes(maxFAT32File))
			}
			n.Size = fi.Size()
			rep.Files++
			rep.DataBytes += fi.Size()
		}
		parent.Children = append(parent.Children, n)
	}
	return nil
}

// slotsFor returns how many 32-byte directory entries a child consumes.
func slotsFor(n *SrcNode) int64 {
	if !n.NeedLFN {
		return 1
	}
	units := utf16Len(n.Name)
	lfn := int64((units + lfnCharsPerSlot - 1) / lfnCharsPerSlot)
	if lfn == 0 {
		lfn = 1
	}
	return lfn + 1
}

// Requirement is the space the scanned tree needs on the volume.
type Requirement struct {
	FileClusters int64 `json:"file_clusters"`
	DirClusters  int64 `json:"dir_clusters"`
	TotalCluster int64 `json:"total_clusters"`
	BytesOnDisk  int64 `json:"bytes_on_disk"`
	RootEntries  int64 `json:"root_entries"`
}

// ComputeLayout fills in DirBytes and Clusters for every node and returns the
// total requirement. hasLabel adds the volume-label entry to the root.
func ComputeLayout(root *SrcNode, clusterBytes int64, hasLabel bool) Requirement {
	var req Requirement
	var walk func(n *SrcNode, isRoot bool)
	walk = func(n *SrcNode, isRoot bool) {
		var entries int64
		if !isRoot {
			entries += 2 // "." and ".."
		} else if hasLabel {
			entries++ // volume label entry
		}
		for _, c := range n.Children {
			entries += slotsFor(c)
		}
		n.DirBytes = entries * dirEntryBytes
		n.Clusters = (n.DirBytes + clusterBytes - 1) / clusterBytes
		if n.Clusters == 0 {
			n.Clusters = 1
		}
		req.DirClusters += n.Clusters
		if isRoot {
			req.RootEntries = entries
		}
		for _, c := range n.Children {
			if c.IsDir {
				walk(c, false)
				continue
			}
			c.Clusters = (c.Size + clusterBytes - 1) / clusterBytes
			req.FileClusters += c.Clusters
		}
	}
	walk(root, true)
	req.TotalCluster = req.FileClusters + req.DirClusters
	req.BytesOnDisk = req.TotalCluster * clusterBytes
	return req
}

// allocate assigns a contiguous cluster run to every node, depth first, and
// returns the first cluster number that remains free.
func allocate(root *SrcNode) uint32 {
	next := uint32(2)
	var walk func(n *SrcNode)
	walk = func(n *SrcNode) {
		n.Start = next
		next += uint32(n.Clusters)
		for _, c := range n.Children {
			if c.IsDir {
				walk(c)
				continue
			}
			if c.Clusters == 0 {
				c.Start = 0
				continue
			}
			c.Start = next
			next += uint32(c.Clusters)
		}
	}
	walk(root)
	return next
}

// ---------------------------------------------------------------------------
// On-disk structures
// ---------------------------------------------------------------------------

// fatTime packs a time into FAT's 16-bit date and time words plus the
// hundredths-of-a-second byte. FAT epochs at 1980 and stores seconds in units
// of two.
func fatTime(t time.Time) (date, tm uint16, tenth byte) {
	t = t.Local()
	y := t.Year()
	if y < 1980 {
		return (1 << 5) | 1, 0, 0 // 1980-01-01 00:00:00
	}
	if y > 2107 {
		y = 2107
	}
	date = uint16((y-1980)<<9) | uint16(int(t.Month())<<5) | uint16(t.Day())
	tm = uint16(t.Hour()<<11) | uint16(t.Minute()<<5) | uint16(t.Second()/2)
	tenth = byte((t.Second() % 2) * 100)
	return date, tm, tenth
}

// shortDirEntry renders one 32-byte 8.3 directory record.
func shortDirEntry(short [11]byte, attr byte, cluster uint32, size uint32, mod time.Time) []byte {
	e := make([]byte, dirEntryBytes)
	copy(e[0:11], short[:])
	e[11] = attr
	e[12] = 0 // NTRes: no lowercase-name shortcuts, LFN slots carry the case
	d, t, tenth := fatTime(mod)
	e[13] = tenth
	binary.LittleEndian.PutUint16(e[14:16], t) // creation time
	binary.LittleEndian.PutUint16(e[16:18], d) // creation date
	binary.LittleEndian.PutUint16(e[18:20], d) // last access date
	binary.LittleEndian.PutUint16(e[20:22], uint16(cluster>>16))
	binary.LittleEndian.PutUint16(e[22:24], t) // write time
	binary.LittleEndian.PutUint16(e[24:26], d) // write date
	binary.LittleEndian.PutUint16(e[26:28], uint16(cluster&0xFFFF))
	binary.LittleEndian.PutUint32(e[28:32], size)
	return e
}

// serializeDir builds the full byte content of one directory's clusters.
func serializeDir(n *SrcNode, parentStart uint32, isRoot bool, label string, clusterBytes int64) []byte {
	buf := make([]byte, 0, n.Clusters*clusterBytes)
	if isRoot {
		if label != "" {
			var l [11]byte
			for i := range l {
				l[i] = ' '
			}
			copy(l[:], label)
			buf = append(buf, shortDirEntry(l, attrVolumeID, 0, 0, n.Mod)...)
		}
	} else {
		var dot, dotdot [11]byte
		for i := range dot {
			dot[i], dotdot[i] = ' ', ' '
		}
		dot[0] = '.'
		dotdot[0], dotdot[1] = '.', '.'
		buf = append(buf, shortDirEntry(dot, attrDirectory, n.Start, 0, n.Mod)...)
		buf = append(buf, shortDirEntry(dotdot, attrDirectory, parentStart, 0, n.Mod)...)
	}
	for _, c := range n.Children {
		if c.NeedLFN {
			buf = append(buf, BuildLFNSlots(c.Name, c.Short)...)
		}
		attr := byte(attrArchive)
		size := uint32(c.Size)
		if c.IsDir {
			attr = attrDirectory
			size = 0
		}
		buf = append(buf, shortDirEntry(c.Short, attr, c.Start, size, c.Mod)...)
	}
	// Pad to whole clusters; the trailing zeros are the end-of-directory marker.
	total := n.Clusters * clusterBytes
	if int64(len(buf)) < total {
		buf = append(buf, make([]byte, total-int64(len(buf)))...)
	}
	return buf
}

// BuildBootSector lays out the FAT32 boot sector and BPB.
func BuildBootSector(g Geometry, volID uint32) []byte {
	s := make([]byte, sectorSize)

	// Jump instruction: JMP SHORT 0x5A / NOP, the value every FAT32 formatter
	// writes. Drivers reject a volume whose first byte is not 0xEB or 0xE9.
	s[0], s[1], s[2] = 0xEB, 0x58, 0x90
	copy(s[3:11], []byte(oemName))

	binary.LittleEndian.PutUint16(s[11:13], uint16(g.SectorSize))      // BPB_BytsPerSec
	s[13] = byte(g.SectorsPerCluster)                                  // BPB_SecPerClus
	binary.LittleEndian.PutUint16(s[14:16], uint16(g.ReservedSectors)) // BPB_RsvdSecCnt
	s[16] = byte(g.NumFATs)                                            // BPB_NumFATs
	binary.LittleEndian.PutUint16(s[17:19], 0)                         // BPB_RootEntCnt, 0 on FAT32
	binary.LittleEndian.PutUint16(s[19:21], 0)                         // BPB_TotSec16, 0 on FAT32
	s[21] = 0xF8                                                       // BPB_Media, fixed disk
	binary.LittleEndian.PutUint16(s[22:24], 0)                         // BPB_FATSz16, 0 on FAT32
	binary.LittleEndian.PutUint16(s[24:26], chsSectorsPerTrak)         // BPB_SecPerTrk
	binary.LittleEndian.PutUint16(s[26:28], chsHeads)                  // BPB_NumHeads
	binary.LittleEndian.PutUint32(s[28:32], uint32(g.PartStartLBA))    // BPB_HiddSec
	binary.LittleEndian.PutUint32(s[32:36], uint32(g.PartSectors))     // BPB_TotSec32
	binary.LittleEndian.PutUint32(s[36:40], uint32(g.FATSectors))      // BPB_FATSz32
	binary.LittleEndian.PutUint16(s[40:42], 0)                         // BPB_ExtFlags: mirrored FATs
	binary.LittleEndian.PutUint16(s[42:44], 0)                         // BPB_FSVer
	binary.LittleEndian.PutUint32(s[44:48], uint32(g.RootCluster))     // BPB_RootClus
	binary.LittleEndian.PutUint16(s[48:50], fsInfoSector)              // BPB_FSInfo
	binary.LittleEndian.PutUint16(s[50:52], backupBootSec)             // BPB_BkBootSec
	// s[52:64] BPB_Reserved stays zero.
	s[64] = 0x80 // BS_DrvNum
	s[65] = 0    // BS_Reserved1
	s[66] = 0x29 // BS_BootSig, says the next three fields are present
	binary.LittleEndian.PutUint32(s[67:71], volID)
	lab := g.Label
	for len(lab) < 11 {
		lab += " "
	}
	copy(s[71:82], lab[:11])
	copy(s[82:90], []byte(fsTypeLabel))
	// s[90:510] is where boot code would live. BootBuilder installs no
	// bootloader, so it stays zero - see the SCOPE section of README.txt.
	s[510], s[511] = 0x55, 0xAA
	return s
}

// BuildFSInfo lays out the FAT32 FSInfo sector.
func BuildFSInfo(freeClusters, nextFree uint32) []byte {
	s := make([]byte, sectorSize)
	binary.LittleEndian.PutUint32(s[0:4], 0x41615252)     // "RRaA" lead signature
	binary.LittleEndian.PutUint32(s[484:488], 0x61417272) // "rrAa" struct signature
	binary.LittleEndian.PutUint32(s[488:492], freeClusters)
	binary.LittleEndian.PutUint32(s[492:496], nextFree)
	binary.LittleEndian.PutUint32(s[508:512], 0xAA550000) // trail signature
	return s
}

// BuildFAT returns the populated prefix of one FAT, as bytes. Entries beyond
// the returned slice are free (zero) and are left as holes in the image.
func BuildFAT(root *SrcNode, lastUsed uint32) []byte {
	fat := make([]uint32, lastUsed+1)
	fat[0] = 0x0FFFFFF8 // media descriptor in the low byte
	if len(fat) > 1 {
		fat[1] = fatEOC
	}
	var chain func(start uint32, count int64)
	chain = func(start uint32, count int64) {
		for i := int64(0); i < count; i++ {
			c := start + uint32(i)
			if i == count-1 {
				fat[c] = fatEOC
			} else {
				fat[c] = c + 1
			}
		}
	}
	var walk func(n *SrcNode)
	walk = func(n *SrcNode) {
		chain(n.Start, n.Clusters)
		for _, c := range n.Children {
			if c.IsDir {
				walk(c)
				continue
			}
			if c.Clusters > 0 {
				chain(c.Start, c.Clusters)
			}
		}
	}
	walk(root)

	out := make([]byte, len(fat)*4)
	for i, v := range fat {
		binary.LittleEndian.PutUint32(out[i*4:i*4+4], v&fatEntryMask)
	}
	return out
}

// ---------------------------------------------------------------------------
// Image writing
// ---------------------------------------------------------------------------

// BuildResult reports what a build actually wrote.
type BuildResult struct {
	Out          string      `json:"out"`
	Geometry     Geometry    `json:"geometry"`
	Scan         *ScanReport `json:"source"`
	Requirement  Requirement `json:"requirement"`
	UsedClusters int64       `json:"used_clusters"`
	FreeClusters int64       `json:"free_clusters"`
	NextFree     uint32      `json:"next_free_cluster"`
	VolumeID     string      `json:"volume_id"`
	DiskGUID     string      `json:"disk_guid,omitempty"`
	PartGUID     string      `json:"partition_guid,omitempty"`
	GPTHeaderCRC string      `json:"gpt_header_crc32,omitempty"`
	GPTArrayCRC  string      `json:"gpt_array_crc32,omitempty"`
	BytesWritten int64       `json:"image_bytes"`
}

// ErrCapacity is returned when the source tree cannot fit the volume.
var ErrCapacity = errors.New("source tree does not fit")

// WriteImage constructs the whole image file. out must not be a device; the
// function creates or truncates a regular file and never opens anything else.
func WriteImage(out string, g Geometry, root *SrcNode, scan *ScanReport, force bool, now time.Time) (*BuildResult, error) {
	req := ComputeLayout(root, g.ClusterBytes, g.Label != "")
	if req.TotalCluster > g.ClusterCount {
		short := (req.TotalCluster - g.ClusterCount) * g.ClusterBytes
		return nil, fmt.Errorf("%w: it needs %d clusters (%s of %s allocation units) but the volume has only %d (%s); short by exactly %d bytes",
			ErrCapacity, req.TotalCluster, humanBytes(req.BytesOnDisk), humanBytes(g.ClusterBytes),
			g.ClusterCount, humanBytes(g.UsableBytes), short)
	}

	if fi, err := os.Stat(out); err == nil {
		if !fi.Mode().IsRegular() {
			return nil, fmt.Errorf("refusing to write %s: it is not a regular file (mode %s). BootBuilder writes image files only", out, fi.Mode())
		}
		if !force {
			return nil, fmt.Errorf("refusing to overwrite existing file %s; pass --force to replace it", out)
		}
	} else if !os.IsNotExist(err) {
		return nil, fmt.Errorf("cannot stat %s: %w", out, err)
	}

	next := allocate(root)
	lastUsed := next - 1
	used := int64(next) - 2
	free := g.ClusterCount - used
	nextFree := next
	if used >= g.ClusterCount {
		nextFree = 0xFFFFFFFF
	}

	volID := uint32(now.Unix())&0xFFFF0000 | uint32(now.Nanosecond()/1000)&0x0000FFFF
	if volID == 0 {
		volID = 0x1337C0DE
	}

	f, err := os.OpenFile(out, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		return nil, fmt.Errorf("cannot create %s: %w", out, err)
	}
	ok := false
	defer func() {
		if !ok {
			f.Close()
		}
	}()
	if err := f.Truncate(g.ImageBytes); err != nil {
		return nil, fmt.Errorf("cannot size %s to %d bytes: %w", out, g.ImageBytes, err)
	}

	res := &BuildResult{
		Out: out, Geometry: g, Scan: scan, Requirement: req,
		UsedClusters: used, FreeClusters: free, NextFree: nextFree,
		VolumeID:     fmt.Sprintf("%04X-%04X", volID>>16, volID&0xFFFF),
		BytesWritten: g.ImageBytes,
	}

	at := func(lba int64, b []byte) error {
		if _, err := f.WriteAt(b, lba*sectorSize); err != nil {
			return fmt.Errorf("cannot write %s at LBA %d: %w", out, lba, err)
		}
		return nil
	}

	// --- partition table -----------------------------------------------------
	switch g.Scheme {
	case "mbr":
		if err := at(0, BuildMBR(volID, g.PartStartLBA, g.PartSectors)); err != nil {
			return nil, err
		}
	case "gpt":
		diskGUID, err := newGUID()
		if err != nil {
			return nil, err
		}
		partGUID, err := newGUID()
		if err != nil {
			return nil, err
		}
		gpt := BuildGPT(g.TotalSectors, g.PartStartLBA, g.PartSectors, g.Label, diskGUID, partGUID)
		if err := at(0, gpt.ProtectiveMBR); err != nil {
			return nil, err
		}
		if err := at(1, gpt.PrimaryHeader); err != nil {
			return nil, err
		}
		if err := at(gpt.PrimaryArrLBA, gpt.EntryArray); err != nil {
			return nil, err
		}
		if err := at(gpt.BackupArrLBA, gpt.EntryArray); err != nil {
			return nil, err
		}
		if err := at(gpt.BackupHdrLBA, gpt.BackupHeader); err != nil {
			return nil, err
		}
		res.DiskGUID = formatGUID(diskGUID)
		res.PartGUID = formatGUID(partGUID)
		res.GPTHeaderCRC = fmt.Sprintf("%08X", gpt.HeaderCRC)
		res.GPTArrayCRC = fmt.Sprintf("%08X", gpt.ArrayCRC)
	}

	// --- reserved region -----------------------------------------------------
	boot := BuildBootSector(g, volID)
	fsi := BuildFSInfo(uint32(free), nextFree)
	if err := at(g.PartStartLBA, boot); err != nil {
		return nil, err
	}
	if err := at(g.PartStartLBA+fsInfoSector, fsi); err != nil {
		return nil, err
	}
	if err := at(g.PartStartLBA+backupBootSec, boot); err != nil {
		return nil, err
	}
	if err := at(g.PartStartLBA+backupBootSec+fsInfoSector, fsi); err != nil {
		return nil, err
	}

	// --- both FAT copies, byte-identical -------------------------------------
	fat := BuildFAT(root, lastUsed)
	if int64(len(fat)) > g.FATBytes {
		return nil, fmt.Errorf("internal error: FAT of %d bytes exceeds the %d reserved for it", len(fat), g.FATBytes)
	}
	fatLBA := g.PartStartLBA + g.ReservedSectors
	for i := int64(0); i < g.NumFATs; i++ {
		if err := at(fatLBA+i*g.FATSectors, fat); err != nil {
			return nil, err
		}
	}

	// --- directories and file data -------------------------------------------
	clusterOffset := func(c uint32) int64 {
		return (g.DataStartLBA + (int64(c)-2)*g.SectorsPerCluster) * sectorSize
	}
	var emit func(n *SrcNode, parentStart uint32, isRoot bool) error
	emit = func(n *SrcNode, parentStart uint32, isRoot bool) error {
		label := ""
		if isRoot {
			label = g.Label
		}
		content := serializeDir(n, parentStart, isRoot, label, g.ClusterBytes)
		if _, err := f.WriteAt(content, clusterOffset(n.Start)); err != nil {
			return fmt.Errorf("cannot write directory at cluster %d: %w", n.Start, err)
		}
		for _, c := range n.Children {
			if c.IsDir {
				ps := n.Start
				if isRoot {
					ps = 0 // ".." in a first-level directory points at the root
				}
				if err := emit(c, ps, false); err != nil {
					return err
				}
				continue
			}
			if c.Size == 0 {
				continue
			}
			if err := copyFileInto(f, c.SrcPath, c.Size, clusterOffset(c.Start), g.ClusterBytes); err != nil {
				return err
			}
		}
		return nil
	}
	if err := emit(root, 0, true); err != nil {
		return nil, err
	}

	if err := f.Sync(); err != nil {
		return nil, fmt.Errorf("cannot flush %s: %w", out, err)
	}
	if err := f.Close(); err != nil {
		return nil, fmt.Errorf("cannot close %s: %w", out, err)
	}
	ok = true
	return res, nil
}

// copyFileInto streams one source file into the image at off, refusing to write
// more than the size recorded at scan time.
func copyFileInto(dst *os.File, src string, size, off, clusterBytes int64) error {
	in, err := os.Open(src)
	if err != nil {
		return fmt.Errorf("cannot read %s: %w", src, err)
	}
	defer in.Close()

	buf := make([]byte, 1<<20)
	if clusterBytes > int64(len(buf)) {
		buf = make([]byte, clusterBytes)
	}
	var done int64
	for done < size {
		want := int64(len(buf))
		if size-done < want {
			want = size - done
		}
		n, err := io.ReadFull(in, buf[:want])
		if n > 0 {
			if _, werr := dst.WriteAt(buf[:n], off+done); werr != nil {
				return fmt.Errorf("cannot write %s into the image: %w", src, werr)
			}
			done += int64(n)
		}
		if err != nil {
			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
				return fmt.Errorf("%s shrank while the image was being written (%d of %d bytes read)", src, done, size)
			}
			return fmt.Errorf("cannot read %s: %w", src, err)
		}
	}
	return nil
}

// utf16Len counts UTF-16 code units, surrogate pairs included.
func utf16Len(s string) int {
	n := 0
	for _, r := range s {
		if r > 0xFFFF {
			n += 2
		} else {
			n++
		}
	}
	return n
}

// ValidateLabel normalises a FAT volume label.
func ValidateLabel(s string) (string, error) {
	if s == "" {
		return "", nil
	}
	u := strings.ToUpper(s)
	if len(u) > 11 {
		return "", fmt.Errorf("volume label %q is %d characters; FAT allows at most 11", s, len(u))
	}
	for i := 0; i < len(u); i++ {
		if u[i] != ' ' && shortNameInvalid(u[i]) {
			return "", fmt.Errorf("volume label %q contains the character %q, which FAT does not allow", s, string(u[i]))
		}
	}
	return u, nil
}