package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
)

// DiskLayout is the full picture of a source disk that clone-disk needs to
// reproduce a target: the partition table (as parsed from `sfdisk -d`, the
// canonical restorable form for both GPT and MBR), the byte ranges outside
// any partition that still have to be copied verbatim (MBR bootstrap + gap,
// BIOS-boot partition), the per-partition filesystem facts, and which of the
// external tools we shell out to are actually present.
type DiskLayout struct {
	// Whole-disk geometry.
	DiskPath       string `json:"diskPath"`
	DiskSize       int64  `json:"diskSize"`       // bytes
	LogicalSector  int64  `json:"logicalSector"`  // bytes, from the kernel/ioctl
	PhysicalSector int64  `json:"physicalSector"` // bytes
	Scheme         string `json:"scheme"`         // "gpt" | "mbr" | "raw" (no recognisable table)

	// Partition table, parsed from `sfdisk -d`.
	Label      string      `json:"label"`    // "gpt" | "dos"
	LabelID    string      `json:"labelID"`  // GPT disk GUID / MBR 4-byte signature (0x...)
	FirstLBA   int64       `json:"firstLBA"` // sectors; 0 if the dump omitted it
	LastLBA    int64       `json:"lastLBA"`  // sectors; 0 if omitted
	Partitions []Partition `json:"partitions"`
	SfdiskDump string      `json:"sfdiskDump"` // the verbatim `sfdisk -d` output
	SgdiskB64  string      `json:"sgdiskB64"`  // base64 `sgdisk --backup` blob (GPT only), fallback restore path

	// Byte ranges that must be raw-copied and are not covered by a cloned
	// partition: MBR bootstrap (0..0x1BE) + the gap before the first
	// partition, and any BIOS-boot partition.
	BootRegions []Region `json:"bootRegions"`

	Tools    map[string]bool `json:"tools"` // external tool name -> present on PATH
	Warnings []string        `json:"warnings,omitempty"`
}

// Region is a [Offset, Offset+Length) byte window on the whole disk.
type Region struct {
	Offset int64  `json:"offset"`
	Length int64  `json:"length"`
	Note   string `json:"note,omitempty"`
}

// Partition is one row of the partition table plus the filesystem facts the
// source agent gathered for it.
type Partition struct {
	Num   int    `json:"num"`
	Start int64  `json:"start"` // sectors
	Size  int64  `json:"size"`  // sectors
	Type  string `json:"type"`  // GPT type GUID or MBR hex code
	UUID  string `json:"uuid"`  // GPT partition GUID (PARTUUID)
	Name  string `json:"name"`  // GPT partition name
	Attrs string `json:"attrs"` // trailing flags: "bootable", "attrs=\"...\"", ...

	// Filesystem facts (source agent; not from sfdisk).
	DevPath     string `json:"devPath"` // source node, e.g. /dev/sda2
	FSType      string `json:"fsType"`  // ext4 | ntfs | vfat | xfs | btrfs | swap | "" (none/unknown)
	FSLabel     string `json:"fsLabel"`
	FSUUID      string `json:"fsUUID"`
	FSUsedBytes int64  `json:"fsUsedBytes"` // 0 if unknown
	FSMinBytes  int64  `json:"fsMinBytes"`  // smallest the fs can be shrunk to; 0 if unknown/unshrinkable
}

// StartBytes / SizeBytes convert the sector-unit table fields using the
// disk's logical sector size.
func (p Partition) StartBytes(sectorSize int64) int64 { return p.Start * sectorSize }
func (p Partition) SizeBytes(sectorSize int64) int64  { return p.Size * sectorSize }

func (d *DiskLayout) JSON() string {
	b, _ := json.Marshal(d)
	return string(b)
}

func parseDiskLayoutJSON(s string) (*DiskLayout, error) {
	var d DiskLayout
	if err := json.Unmarshal([]byte(s), &d); err != nil {
		return nil, fmt.Errorf("decode disk layout: %w", err)
	}
	return &d, nil
}

// parseSfdiskDump fills the Label/LabelID/FirstLBA/LastLBA/Partitions fields
// of d from the text of `sfdisk -d`. It leaves the raw text in d.SfdiskDump.
func (d *DiskLayout) parseSfdiskDump(dump string) error {
	d.SfdiskDump = dump
	sc := bufio.NewScanner(strings.NewReader(dump))
	sc.Buffer(make([]byte, 64*1024), 1024*1024)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		if i := strings.Index(line, " : "); i >= 0 && strings.Contains(line, "start=") {
			p, err := parseSfdiskPartLine(line)
			if err != nil {
				return err
			}
			d.Partitions = append(d.Partitions, p)
			continue
		}
		key, val, ok := splitKV(line)
		if !ok {
			continue
		}
		switch key {
		case "label":
			d.Label = val
		case "label-id":
			d.LabelID = val
		case "first-lba":
			d.FirstLBA, _ = strconv.ParseInt(val, 10, 64)
		case "last-lba":
			d.LastLBA, _ = strconv.ParseInt(val, 10, 64)
		case "sector-size":
			if n, err := strconv.ParseInt(val, 10, 64); err == nil && d.LogicalSector == 0 {
				d.LogicalSector = n
			}
		}
	}
	if err := sc.Err(); err != nil {
		return fmt.Errorf("scan sfdisk dump: %w", err)
	}
	switch d.Label {
	case "gpt":
		d.Scheme = "gpt"
	case "dos":
		d.Scheme = "mbr"
	}
	return nil
}

// splitKV parses a "key: value" header line.
func splitKV(line string) (string, string, bool) {
	i := strings.Index(line, ":")
	if i < 0 {
		return "", "", false
	}
	return strings.TrimSpace(line[:i]), strings.TrimSpace(line[i+1:]), true
}

// parseSfdiskPartLine parses a body line like
//
//	/dev/sda2 : start=     1050624, size=  1951xxxxxx, type=0FC6..., uuid=..., name="root", attrs="..."
//	/dev/sda1 : start=        2048, size=   204800, type=83, bootable
func parseSfdiskPartLine(line string) (Partition, error) {
	var p Partition
	i := strings.Index(line, " : ")
	dev := strings.TrimSpace(line[:i])
	p.DevPath = dev
	p.Num = trailingInt(dev)
	rest := line[i+3:]
	for _, fld := range splitTopLevel(rest, ',') {
		fld = strings.TrimSpace(fld)
		if fld == "" {
			continue
		}
		k, v, ok := splitEq(fld)
		if !ok {
			// bare flag, e.g. "bootable"
			if p.Attrs == "" {
				p.Attrs = fld
			} else {
				p.Attrs += ", " + fld
			}
			continue
		}
		v = strings.Trim(strings.TrimSpace(v), `"`)
		switch k {
		case "start":
			p.Start, _ = strconv.ParseInt(strings.TrimSpace(v), 10, 64)
		case "size":
			p.Size, _ = strconv.ParseInt(strings.TrimSpace(v), 10, 64)
		case "type":
			p.Type = v
		case "uuid":
			p.UUID = v
		case "name":
			p.Name = v
		case "attrs":
			if p.Attrs == "" {
				p.Attrs = `attrs="` + v + `"`
			} else {
				p.Attrs += `, attrs="` + v + `"`
			}
		}
	}
	if p.Start == 0 && p.Size == 0 {
		return p, fmt.Errorf("sfdisk line without start/size: %q", line)
	}
	return p, nil
}

// sfdiskRestoreScript renders a device-independent script that recreates this
// table on any target with `sfdisk`. It intentionally drops `device:` and the
// per-line device prefixes (sfdisk numbers the partitions in order) and
// `last-lba` (let sfdisk size the secondary GPT for the target). label-id and
// per-partition uuid are dropped when newIDs is set so a replacement disk can
// coexist with the original; otherwise they are preserved so existing
// BCD/fstab/GRUB references still resolve. sizeOverride maps a 1-based
// partition number to a new size in sectors (used by the shrink path).
func (d *DiskLayout) sfdiskRestoreScript(newIDs bool, sizeOverride map[int]int64) string {
	var b strings.Builder
	fmt.Fprintf(&b, "label: %s\n", d.Label)
	if d.LabelID != "" && !newIDs {
		fmt.Fprintf(&b, "label-id: %s\n", d.LabelID)
	}
	if d.FirstLBA > 0 {
		fmt.Fprintf(&b, "first-lba: %d\n", d.FirstLBA)
	}
	if d.LogicalSector > 0 {
		fmt.Fprintf(&b, "sector-size: %d\n", d.LogicalSector)
	}
	b.WriteString("unit: sectors\n\n")
	for _, p := range d.Partitions {
		size := p.Size
		if ov, ok := sizeOverride[p.Num]; ok && ov > 0 {
			size = ov
		}
		fmt.Fprintf(&b, "start=%d, size=%d", p.Start, size)
		if p.Type != "" {
			fmt.Fprintf(&b, ", type=%s", p.Type)
		}
		if p.UUID != "" && !newIDs {
			fmt.Fprintf(&b, ", uuid=%s", p.UUID)
		}
		if p.Name != "" {
			fmt.Fprintf(&b, ", name=\"%s\"", p.Name)
		}
		if p.Attrs != "" {
			fmt.Fprintf(&b, ", %s", p.Attrs)
		}
		b.WriteByte('\n')
	}
	return b.String()
}

// biosBootType / espType are the GPT type GUIDs clone-disk special-cases.
const (
	gptBIOSBoot = "21686148-6449-6E6F-744E-656564454649" // EF02 BIOS boot partition
	gptESP      = "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" // EF00 EFI system partition
)

func (p Partition) isBIOSBoot() bool {
	return strings.EqualFold(p.Type, gptBIOSBoot) || p.Type == "21" /* mbr */
}
func (p Partition) isESP() bool {
	return strings.EqualFold(p.Type, gptESP) || p.Type == "ef" || p.Type == "EF"
}

// --- small string helpers -------------------------------------------------

func trailingInt(s string) int {
	j := len(s)
	for j > 0 && s[j-1] >= '0' && s[j-1] <= '9' {
		j--
	}
	if j == len(s) {
		return 0
	}
	n, _ := strconv.Atoi(s[j:])
	return n
}

func splitEq(s string) (string, string, bool) {
	i := strings.Index(s, "=")
	if i < 0 {
		return "", "", false
	}
	return strings.TrimSpace(s[:i]), s[i+1:], true
}

// splitTopLevel splits on sep but not inside double quotes.
func splitTopLevel(s string, sep byte) []string {
	var out []string
	var cur strings.Builder
	inQ := false
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch {
		case c == '"':
			inQ = !inQ
			cur.WriteByte(c)
		case c == sep && !inQ:
			out = append(out, cur.String())
			cur.Reset()
		default:
			cur.WriteByte(c)
		}
	}
	out = append(out, cur.String())
	return out
}