package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

const defaultLabel = "BOOTBUILDER"

func emitJSON(v any) {
	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	if err := enc.Encode(v); err != nil {
		fail("%v", err)
	}
}

// ---------------------------------------------------------------------------
// plan
// ---------------------------------------------------------------------------

// PlanReport is the JSON shape of `plan`.
type PlanReport struct {
	Source      string      `json:"source"`
	Geometry    Geometry    `json:"geometry"`
	Scan        *ScanReport `json:"source_tree"`
	Requirement Requirement `json:"requirement"`
	Fits        bool        `json:"fits"`
	ShortBy     int64       `json:"short_by_bytes"`
	FreeAfter   int64       `json:"free_clusters_after"`
	SlackNote   string      `json:"slack_note"`
}

func cmdPlan(argv []string) {
	fs := newFlagSet("plan")
	src := fs.String("src", "", "source directory tree")
	fs.StringVar(src, "s", "", "shorthand for --src")
	sizeStr := fs.String("size", "", "image size, e.g. 8GB")
	scheme := fs.String("scheme", "mbr", "partition scheme: mbr or gpt")
	label := fs.String("label", defaultLabel, "FAT volume label")
	clus := fs.String("cluster-size", "", "cluster size override in bytes")
	asJSON := fs.Bool("json", false, "JSON output")
	if err := fs.Parse(reorderFlags(argv, valueFlags)); err != nil {
		os.Exit(1)
	}
	if *src == "" && fs.NArg() > 0 {
		*src = fs.Arg(0)
	}
	if *src == "" {
		usageErr("plan needs --src <dir>")
	}
	if *sizeStr == "" {
		usageErr("plan needs --size <size>, e.g. --size 8GB")
	}
	g, root, scan := prepare(*src, *sizeStr, *scheme, *label, *clus)

	req := ComputeLayout(root, g.ClusterBytes, g.Label != "")
	rep := PlanReport{
		Source: *src, Geometry: g, Scan: scan, Requirement: req,
		Fits:      req.TotalCluster <= g.ClusterCount,
		FreeAfter: g.ClusterCount - req.TotalCluster,
	}
	if !rep.Fits {
		rep.ShortBy = (req.TotalCluster - g.ClusterCount) * g.ClusterBytes
	}
	rep.SlackNote = fmt.Sprintf("%d sectors (%s) at the end of the partition are not covered by any cluster",
		g.SlackSectors, humanBytes(g.SlackBytes))

	if *asJSON {
		emitJSON(rep)
		if !rep.Fits {
			os.Exit(1)
		}
		return
	}

	fmt.Printf("BootBuilder layout plan (nothing was written)\n")
	fmt.Printf("source     : %s\n", *src)
	printGeometry(g)
	fmt.Println()
	fmt.Printf("SOURCE TREE\n")
	fmt.Printf("  files             : %d\n", scan.Files)
	fmt.Printf("  directories       : %d\n", scan.Dirs)
	fmt.Printf("  content           : %s (%d bytes)\n", humanBytes(scan.DataBytes), scan.DataBytes)
	if len(scan.Skipped) > 0 {
		fmt.Printf("  skipped           : %d non-regular entries\n", len(scan.Skipped))
		for _, s := range scan.Skipped {
			fmt.Printf("      %s\n", s)
		}
	}
	fmt.Println()
	fmt.Printf("ALLOCATION\n")
	fmt.Printf("  file clusters     : %d\n", req.FileClusters)
	fmt.Printf("  directory clusters: %d\n", req.DirClusters)
	fmt.Printf("  total clusters    : %d of %d\n", req.TotalCluster, g.ClusterCount)
	fmt.Printf("  bytes on volume   : %s (%d bytes, including cluster slack)\n", humanBytes(req.BytesOnDisk), req.BytesOnDisk)
	fmt.Printf("  root dir entries  : %d (%d bytes)\n", req.RootEntries, req.RootEntries*dirEntryBytes)
	fmt.Println()
	if rep.Fits {
		fmt.Printf("verdict    : FITS - %d clusters (%s) would remain free\n",
			rep.FreeAfter, humanBytes(rep.FreeAfter*g.ClusterBytes))
		return
	}
	fmt.Printf("verdict    : DOES NOT FIT - short by exactly %d bytes (%s)\n", rep.ShortBy, humanBytes(rep.ShortBy))
	os.Exit(1)
}

func printGeometry(g Geometry) {
	fmt.Printf("image size : %s (%d bytes, %d sectors of %d)\n", g.ImageHuman, g.ImageBytes, g.TotalSectors, g.SectorSize)
	fmt.Printf("scheme     : %s\n", g.Scheme)
	fmt.Println()
	fmt.Printf("GEOMETRY\n")
	fmt.Printf("  partition start   : LBA %d (%s in)\n", g.PartStartLBA, humanBytes(g.PartStartLBA*g.SectorSize))
	fmt.Printf("  partition size    : %d sectors (%s)\n", g.PartSectors, humanBytes(g.PartBytes))
	src := "Microsoft FAT32 table"
	if !g.ClusterFromTable {
		src = "--cluster-size override"
	}
	fmt.Printf("  sectors/cluster   : %d  -> cluster %s (%d bytes, from the %s)\n",
		g.SectorsPerCluster, humanBytes(g.ClusterBytes), g.ClusterBytes, src)
	fmt.Printf("  reserved sectors  : %d\n", g.ReservedSectors)
	fmt.Printf("  FAT copies        : %d x %d sectors (%s each)\n", g.NumFATs, g.FATSectors, humanBytes(g.FATBytes))
	fmt.Printf("  data starts at    : LBA %d\n", g.DataStartLBA)
	fmt.Printf("  data sectors      : %d\n", g.DataSectors)
	fmt.Printf("  clusters          : %d (usable %s)\n", g.ClusterCount, g.UsableHuman)
	fmt.Printf("  slack             : %d sectors (%s) past the last cluster\n", g.SlackSectors, humanBytes(g.SlackBytes))
	fmt.Printf("  root cluster      : %d\n", g.RootCluster)
	fmt.Printf("  volume label      : %q\n", g.Label)
}

// prepare resolves the common flags of plan and build.
func prepare(src, sizeStr, scheme, label, clus string) (Geometry, *SrcNode, *ScanReport) {
	size, err := parseSize(sizeStr)
	if err != nil {
		usageErr("%v", err)
	}
	var clusBytes int64
	if clus != "" {
		clusBytes, err = parseSize(clus)
		if err != nil {
			usageErr("%v", err)
		}
	}
	lab, err := ValidateLabel(label)
	if err != nil {
		usageErr("%v", err)
	}
	g, err := PlanGeometry(size, strings.ToLower(scheme), lab, clusBytes)
	if err != nil {
		fail("%v", err)
	}
	abs, err := filepath.Abs(src)
	if err != nil {
		fail("cannot resolve %q: %v", src, err)
	}
	root, scan, err := ScanTree(abs)
	if err != nil {
		fail("%v", err)
	}
	return g, root, scan
}

// ---------------------------------------------------------------------------
// build
// ---------------------------------------------------------------------------

func cmdBuild(argv []string) {
	fs := newFlagSet("build")
	src := fs.String("src", "", "source directory tree")
	fs.StringVar(src, "s", "", "shorthand for --src")
	out := fs.String("out", "", "image file to write")
	fs.StringVar(out, "o", "", "shorthand for --out")
	sizeStr := fs.String("size", "", "image size, e.g. 8GB")
	scheme := fs.String("scheme", "mbr", "partition scheme: mbr or gpt")
	label := fs.String("label", defaultLabel, "FAT volume label")
	clus := fs.String("cluster-size", "", "cluster size override in bytes")
	force := fs.Bool("force", false, "overwrite an existing --out")
	asJSON := fs.Bool("json", false, "JSON output")
	if err := fs.Parse(reorderFlags(argv, valueFlags)); err != nil {
		os.Exit(1)
	}
	if *out == "" && fs.NArg() > 0 {
		*out = fs.Arg(0)
	}
	if *src == "" {
		usageErr("build needs --src <dir>")
	}
	if *out == "" {
		usageErr("build needs --out <file>")
	}
	if *sizeStr == "" {
		usageErr("build needs --size <size>, e.g. --size 8GB")
	}
	g, root, scan := prepare(*src, *sizeStr, *scheme, *label, *clus)

	res, err := WriteImage(*out, g, root, scan, *force, time.Now())
	if err != nil {
		fail("%v", err)
	}

	if *asJSON {
		emitJSON(res)
		return
	}
	fmt.Printf("BootBuilder wrote %s\n", res.Out)
	fmt.Printf("source     : %s\n", *src)
	printGeometry(g)
	fmt.Println()
	fmt.Printf("VOLUME\n")
	fmt.Printf("  volume id         : %s\n", res.VolumeID)
	fmt.Printf("  files written     : %d in %d directories\n", scan.Files, scan.Dirs)
	fmt.Printf("  clusters used     : %d of %d (%s)\n", res.UsedClusters, g.ClusterCount, humanBytes(res.UsedClusters*g.ClusterBytes))
	fmt.Printf("  clusters free     : %d (%s)\n", res.FreeClusters, humanBytes(res.FreeClusters*g.ClusterBytes))
	fmt.Printf("  next free cluster : %d\n", res.NextFree)
	if g.Scheme == "gpt" {
		fmt.Printf("  disk GUID         : %s\n", res.DiskGUID)
		fmt.Printf("  partition GUID    : %s\n", res.PartGUID)
		fmt.Printf("  GPT header CRC32  : 0x%s\n", res.GPTHeaderCRC)
		fmt.Printf("  GPT array CRC32   : 0x%s\n", res.GPTArrayCRC)
	}
	if len(scan.Skipped) > 0 {
		fmt.Printf("  skipped           : %d non-regular entries\n", len(scan.Skipped))
	}
	fmt.Println()
	fmt.Printf("No bootloader was installed. See SCOPE in README.txt.\n")
	fmt.Printf("Next: %s inspect %s\n", appName, res.Out)
}

// ---------------------------------------------------------------------------
// inspect
// ---------------------------------------------------------------------------

// InspectReport is the JSON shape of `inspect`.
type InspectReport struct {
	Volume  *RVolume `json:"volume"`
	Label   string   `json:"root_label_entry"`
	Entries []REntry `json:"entries"`
	Files   int      `json:"file_count"`
	Dirs    int      `json:"dir_count"`
	Bytes   int64    `json:"total_file_bytes"`
}

func cmdInspect(argv []string) {
	fs := newFlagSet("inspect")
	asJSON := fs.Bool("json", false, "JSON output")
	if err := fs.Parse(reorderFlags(argv, valueFlags)); err != nil {
		os.Exit(1)
	}
	if fs.NArg() != 1 {
		usageErr("inspect needs exactly one image file")
	}
	v, err := OpenImage(fs.Arg(0))
	if err != nil {
		fail("%v", err)
	}
	defer v.Close()

	rep := InspectReport{Volume: v}
	rep.Label, err = v.VolumeLabelEntry()
	if err != nil {
		fail("%v", err)
	}
	err = v.Walk(func(e REntry) error {
		rep.Entries = append(rep.Entries, e)
		if e.IsDir {
			rep.Dirs++
		} else {
			rep.Files++
			rep.Bytes += int64(e.Size)
		}
		return nil
	})
	if err != nil {
		fail("%v", err)
	}

	if *asJSON {
		emitJSON(rep)
		return
	}

	b := v.BPB
	fmt.Printf("BootBuilder inspect: %s\n", v.Path)
	fmt.Printf("  image size        : %s (%d bytes)\n", v.ImageHuman, v.ImageBytes)
	fmt.Printf("  scheme            : %s\n", v.Scheme)
	fmt.Println()
	fmt.Printf("PARTITION TABLE\n")
	for _, p := range v.Partitions {
		fmt.Printf("  #%d %-22s %s\n", p.Index, p.Type, p.TypeRaw)
		fmt.Printf("      LBA %d..%d  %d sectors  %s\n", p.FirstLBA, p.LastLBA, p.Sectors, p.Human)
		if p.FirstCHS != "" {
			fmt.Printf("      CHS %s .. %s%s\n", p.FirstCHS, p.LastCHS, bootFlag(p.Bootable))
		}
		if p.Name != "" {
			fmt.Printf("      name %q  guid %s\n", p.Name, p.UniqueGUID)
		}
	}
	if v.GPT != nil {
		g := v.GPT
		fmt.Println()
		fmt.Printf("GPT\n")
		fmt.Printf("  signature         : %q  revision %s  header %d bytes\n", g.Signature, g.Revision, g.HeaderSize)
		fmt.Printf("  header CRC32      : 0x%s  %s\n", g.HeaderCRC, okText(g.HeaderCRCOK))
		fmt.Printf("  array CRC32       : 0x%s  %s\n", g.ArrayCRC, okText(g.ArrayCRCOK))
		fmt.Printf("  my/alternate LBA  : %d / %d\n", g.MyLBA, g.AlternateLBA)
		fmt.Printf("  usable LBA range  : %d..%d\n", g.FirstUsableLBA, g.LastUsableLBA)
		fmt.Printf("  disk GUID         : %s\n", g.DiskGUID)
		fmt.Printf("  entry array       : LBA %d, %d x %d bytes\n", g.EntryLBA, g.EntryCount, g.EntrySize)
		fmt.Printf("  backup header     : present=%v at LBA %d, CRC32 %s\n", g.BackupPresent, g.BackupMyLBA, okText(g.BackupCRCOK))
		fmt.Printf("  backup array      : matches primary %v\n", g.BackupArrayMtch)
	}
	fmt.Println()
	fmt.Printf("BOOT SECTOR / BPB (LBA %d)\n", v.FSPartLBA)
	fmt.Printf("  jump / OEM        : %s  %q\n", b.JumpInstruction, b.OEMName)
	fmt.Printf("  bytes/sector      : %d\n", b.BytesPerSector)
	fmt.Printf("  sectors/cluster   : %d (cluster %s)\n", b.SectorsPerCluster, humanBytes(int64(b.ClusterBytes)))
	fmt.Printf("  reserved sectors  : %d\n", b.ReservedSectors)
	fmt.Printf("  FATs              : %d x %d sectors\n", b.NumFATs, b.FATSize32)
	fmt.Printf("  media descriptor  : %s\n", b.MediaDescriptor)
	fmt.Printf("  hidden sectors    : %d\n", b.HiddenSectors)
	fmt.Printf("  total sectors 32  : %d\n", b.TotalSectors32)
	fmt.Printf("  root cluster      : %d\n", b.RootCluster)
	fmt.Printf("  FSInfo / backup   : sector %d / sector %d\n", b.FSInfoSector, b.BackupBootSector)
	fmt.Printf("  volume id / label : %s  %q\n", b.VolumeID, b.VolumeLabel)
	fmt.Printf("  fs type field     : %q\n", b.FileSystemType)
	fmt.Printf("  0x55AA signature  : %v\n", b.Signature55AA)
	fmt.Printf("  backup boot sector: identical=%v\n", b.BackupMatches)
	fmt.Printf("  data starts at    : LBA %d\n", b.DataStartLBA)
	fmt.Printf("  cluster count     : %d\n", b.ClusterCount)
	fmt.Println()
	fmt.Printf("FSINFO\n")
	fmt.Printf("  signatures        : lead=%v struct=%v trail=%v\n", v.FSInfo.LeadSigOK, v.FSInfo.StructSig, v.FSInfo.TrailSigOK)
	fmt.Printf("  free clusters     : %d (%s)\n", v.FSInfo.FreeCount, humanBytes(int64(v.FSInfo.FreeCount)*int64(b.ClusterBytes)))
	fmt.Printf("  next free cluster : %d\n", v.FSInfo.NextFree)
	fmt.Printf("  FAT copies equal  : %v\n", v.FATsEqual)
	fmt.Println()
	fmt.Printf("ROOT LABEL ENTRY    : %q\n", rep.Label)
	fmt.Println()
	fmt.Printf("DIRECTORY LISTING (%d files, %d directories, %s of file data)\n", rep.Files, rep.Dirs, humanBytes(rep.Bytes))
	for _, e := range rep.Entries {
		depth := strings.Count(e.Path, "/")
		indent := strings.Repeat("  ", depth)
		kind := "     "
		if e.IsDir {
			kind = "<dir>"
		}
		fmt.Printf("  %s %10d  clus %-8d lfn %d  %s%s   [%s]\n",
			kind, e.Size, e.Cluster, e.LFNSlots, indent, e.Name, e.Short)
	}
}

func bootFlag(b bool) string {
	if b {
		return "  (bootable flag set)"
	}
	return ""
}

func okText(ok bool) string {
	if ok {
		return "VALID"
	}
	return "MISMATCH"
}

// ---------------------------------------------------------------------------
// verify
// ---------------------------------------------------------------------------

// VerifyReport is the JSON shape of `verify`.
type VerifyReport struct {
	Image     string   `json:"image"`
	Against   string   `json:"against"`
	Checked   int      `json:"files_checked"`
	Dirs      int      `json:"dirs_checked"`
	Bytes     int64    `json:"bytes_hashed"`
	Missing   []string `json:"missing_from_image,omitempty"`
	Extra     []string `json:"extra_in_image,omitempty"`
	Mismatch  []string `json:"content_mismatch,omitempty"`
	NameIssue []string `json:"name_mismatch,omitempty"`
	OK        bool     `json:"ok"`
}

func cmdVerify(argv []string) {
	fs := newFlagSet("verify")
	against := fs.String("against", "", "source tree to compare with")
	fs.StringVar(against, "a", "", "shorthand for --against")
	asJSON := fs.Bool("json", false, "JSON output")
	if err := fs.Parse(reorderFlags(argv, valueFlags)); err != nil {
		os.Exit(1)
	}
	if fs.NArg() != 1 {
		usageErr("verify needs exactly one image file")
	}
	if *against == "" {
		usageErr("verify needs --against <dir>")
	}
	img := fs.Arg(0)

	srcFiles, srcDirs, err := hashSourceTree(*against)
	if err != nil {
		fail("%v", err)
	}
	v, err := OpenImage(img)
	if err != nil {
		fail("%v", err)
	}
	defer v.Close()

	rep := VerifyReport{Image: img, Against: *against}
	seen := map[string]bool{}
	seenDirs := map[string]bool{}
	err = v.Walk(func(e REntry) error {
		if e.IsDir {
			seenDirs[e.Path] = true
			if !srcDirs[e.Path] {
				rep.Extra = append(rep.Extra, e.Path+"/")
			}
			rep.Dirs++
			return nil
		}
		seen[e.Path] = true
		want, ok := srcFiles[e.Path]
		if !ok {
			rep.Extra = append(rep.Extra, e.Path)
			return nil
		}
		got, err := v.HashFile(e)
		if err != nil {
			return fmt.Errorf("%s: %w", e.Path, err)
		}
		rep.Checked++
		rep.Bytes += int64(e.Size)
		if got != want.sha {
			rep.Mismatch = append(rep.Mismatch,
				fmt.Sprintf("%s: image sha256 %s, source sha256 %s", e.Path, got, want.sha))
			return nil
		}
		if int64(e.Size) != want.size {
			rep.Mismatch = append(rep.Mismatch,
				fmt.Sprintf("%s: image size %d, source size %d", e.Path, e.Size, want.size))
			return nil
		}
		if path.Base(e.Path) != path.Base(want.rel) {
			rep.NameIssue = append(rep.NameIssue, e.Path)
		}
		return nil
	})
	if err != nil {
		fail("%v", err)
	}
	for p := range srcFiles {
		if !seen[p] {
			rep.Missing = append(rep.Missing, p)
		}
	}
	for p := range srcDirs {
		if !seenDirs[p] {
			rep.Missing = append(rep.Missing, p+"/")
		}
	}
	sort.Strings(rep.Missing)
	sort.Strings(rep.Extra)
	sort.Strings(rep.Mismatch)
	rep.OK = len(rep.Missing) == 0 && len(rep.Extra) == 0 && len(rep.Mismatch) == 0 && len(rep.NameIssue) == 0

	if *asJSON {
		emitJSON(rep)
		if !rep.OK {
			os.Exit(2)
		}
		return
	}
	fmt.Printf("BootBuilder verify\n")
	fmt.Printf("  image   : %s\n", img)
	fmt.Printf("  against : %s\n", *against)
	fmt.Printf("  checked : %d files (%s hashed), %d directories\n", rep.Checked, humanBytes(rep.Bytes), rep.Dirs)
	for _, m := range rep.Missing {
		fmt.Printf("  MISSING  %s\n", m)
	}
	for _, m := range rep.Extra {
		fmt.Printf("  EXTRA    %s\n", m)
	}
	for _, m := range rep.Mismatch {
		fmt.Printf("  DIFFERS  %s\n", m)
	}
	for _, m := range rep.NameIssue {
		fmt.Printf("  NAME     %s\n", m)
	}
	if rep.OK {
		fmt.Printf("  result  : OK - every file matches by SHA-256 and every long name survived\n")
		return
	}
	fmt.Printf("  result  : FAILED\n")
	os.Exit(2)
}

type srcFile struct {
	rel  string
	sha  string
	size int64
}

func hashSourceTree(root string) (map[string]srcFile, map[string]bool, error) {
	abs, err := filepath.Abs(root)
	if err != nil {
		return nil, nil, fmt.Errorf("cannot resolve %q: %w", root, err)
	}
	fi, err := os.Stat(abs)
	if err != nil {
		return nil, nil, fmt.Errorf("cannot read %s: %w", root, err)
	}
	if !fi.IsDir() {
		return nil, nil, fmt.Errorf("%s is not a directory", root)
	}
	files := map[string]srcFile{}
	dirs := map[string]bool{}
	var walk func(dir, rel string) error
	walk = func(dir, rel string) error {
		ents, err := os.ReadDir(dir)
		if err != nil {
			return fmt.Errorf("cannot read %s: %w", dir, err)
		}
		for _, de := range ents {
			full := filepath.Join(dir, de.Name())
			r := path.Join(rel, de.Name())
			if de.IsDir() {
				dirs[r] = true
				if err := walk(full, r); err != nil {
					return err
				}
				continue
			}
			if !de.Type().IsRegular() {
				continue
			}
			f, err := os.Open(full)
			if err != nil {
				return fmt.Errorf("cannot read %s: %w", full, err)
			}
			h := sha256.New()
			n, err := io.Copy(h, f)
			f.Close()
			if err != nil {
				return fmt.Errorf("cannot read %s: %w", full, err)
			}
			files[r] = srcFile{rel: r, sha: hex.EncodeToString(h.Sum(nil)), size: n}
		}
		return nil
	}
	if err := walk(abs, ""); err != nil {
		return nil, nil, err
	}
	return files, dirs, nil
}