// Command driveforge is a small, honest prototype of a disk-utility CLI.
//
// DriveForge's full product concept (see ../plan.md) covers partition-table
// manipulation, cloning and migration - all of which require raw
// block-device access and OS-privileged operations that cannot be
// implemented portably in a dependency-free Go CLI (no cgo, no external
// deps, no OS-specific syscall build tags). That's the same honest scoping
// constraint as the sibling tool DiskOps, which implements a single-target
// sequential I/O throughput benchmark as the one piece of the concept that
// genuinely works cross-platform with only the stdlib.
//
// DriveForge is the "Pro" tier of that same family: instead of benchmarking
// one target, it runs the identical write/read throughput methodology
// across SEVERAL target directories/mount points in a single invocation and
// produces a ranked comparison report - the "which of my available drives
// should this workload go on" question that a single-target benchmark can't
// answer on its own. Each target is benchmarked independently and a bad
// target (missing, unwritable) is reported as a per-target failure without
// aborting the rest of the comparison.
package main

import (
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
)

const tempFileName = ".driveforge-benchmark-tmp"

func main() {
	if len(os.Args) < 2 {
		// Double-clicked in Explorer rather than run from a prompt: ask the
		// one question the program needs and stay on screen. Printing usage
		// and exiting here is what made the window vanish instantly.
		if interactiveConsole() {
			runGuided()
			return
		}
		usage()
		os.Exit(1)
	}

	switch os.Args[1] {
	case "-h", "--help", "help":
		usage()
		return
	case "bench":
		runBench(os.Args[2:])
	default:
		fmt.Fprintf(os.Stderr, "driveforge: unknown command %q\n\n", os.Args[1])
		usage()
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprint(os.Stderr, `DriveForge - multi-target disk I/O throughput comparison (prototype)

Usage:
  driveforge bench <dir1> [<dir2> ...] [--size 256MB] [--block 1MB] [--json]
  driveforge help

Commands:
  bench    Run the same sequential write/read throughput benchmark against
           each target directory, independently, and print a ranked
           comparison sorted by write MB/s (fastest first).

Flags for bench:
  --size    Total bytes to write/read per target, e.g. 256MB, 1GiB, 512KB
            (default 256MB). Applied identically to every target.
  --block   Chunk size for each write/read call, e.g. 1MB, 256KB
            (default 1MB). Applied identically to every target.
  --json    Emit a JSON report instead of human-readable text

Notes:
  - Same methodology as DiskOps (sequential, single-threaded, file-API
    throughput - not a raw block-device benchmark) run across multiple
    targets in one pass for an apples-to-apples comparison.
  - A target that doesn't exist or isn't writable is reported as FAILED with
    a reason; the other targets still run and are ranked normally.
  - OS and filesystem caching mean small --size values can report
    unrealistic numbers, especially if targets share the same underlying
    filesystem/cache. Use at least a few hundred MB for a meaningful result.

Examples:
  driveforge bench /mnt/ssd /mnt/hdd --size 512MB
  driveforge bench . /tmp /mnt/data --size 1GiB --block 4MB --json
`)
}

// reorderFlags moves all recognized flags (and their values, for flags that
// take one) to the front of args and all remaining positional arguments to
// the back. This works around flag.FlagSet.Parse's stop-at-first-positional
// behavior, since this CLI's path arguments legitimately come before flags,
// e.g. `driveforge bench dir1 dir2 --size 1GB`.
func reorderFlags(args []string, valueFlags map[string]bool) []string {
	var flags, positional []string
	for i := 0; i < len(args); i++ {
		a := args[i]
		name := strings.TrimLeft(a, "-")
		if strings.HasPrefix(a, "-") && valueFlags[name] {
			flags = append(flags, a)
			if i+1 < len(args) {
				i++
				flags = append(flags, args[i])
			}
			continue
		}
		if strings.HasPrefix(a, "-") {
			flags = append(flags, a)
			continue
		}
		positional = append(positional, a)
	}
	return append(flags, positional...)
}

// parseSize parses a human byte size like "256MB", "1GiB", "512KB", "128"
// (bytes) into a count of bytes. It accepts both SI (KB/MB/GB, powers of
// 1000) and binary (KiB/MiB/GiB, powers of 1024) suffixes, case-insensitive.
func parseSize(s string) (int64, error) {
	orig := s
	s = strings.TrimSpace(s)
	if s == "" {
		return 0, fmt.Errorf("empty size")
	}
	upper := strings.ToUpper(s)

	type unit struct {
		suffix string
		mult   int64
	}
	units := []unit{
		{"GIB", 1024 * 1024 * 1024},
		{"MIB", 1024 * 1024},
		{"KIB", 1024},
		{"GB", 1000 * 1000 * 1000},
		{"MB", 1000 * 1000},
		{"KB", 1000},
		{"B", 1},
	}

	for _, u := range units {
		if strings.HasSuffix(upper, u.suffix) {
			numPart := strings.TrimSpace(s[:len(s)-len(u.suffix)])
			if numPart == "" {
				return 0, fmt.Errorf("invalid size %q: missing number", orig)
			}
			f, err := strconv.ParseFloat(numPart, 64)
			if err != nil {
				return 0, fmt.Errorf("invalid size %q: %w", orig, err)
			}
			if f < 0 {
				return 0, fmt.Errorf("invalid size %q: must not be negative", orig)
			}
			return int64(f * float64(u.mult)), nil
		}
	}

	// No recognized suffix: treat as a plain byte count.
	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return 0, fmt.Errorf("invalid size %q: expected a number optionally followed by B/KB/MB/GB/KiB/MiB/GiB", orig)
	}
	if f < 0 {
		return 0, fmt.Errorf("invalid size %q: must not be negative", orig)
	}
	return int64(f), nil
}

// formatBytes renders a byte count in human-friendly binary units, for
// display purposes only (parseSize is the inverse-ish operation on input).
func formatBytes(n int64) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := int64(unit), 0
	for m := n / unit; m >= unit; m /= unit {
		div *= unit
		exp++
	}
	suffixes := []string{"KiB", "MiB", "GiB", "TiB", "PiB"}
	return fmt.Sprintf("%.2f %s", float64(n)/float64(div), suffixes[exp])
}

// targetResult holds the outcome of benchmarking a single target directory,
// either a successful measurement or a failure reason. Exactly one of the
// two is meaningful, distinguished by OK.
type targetResult struct {
	Dir           string  `json:"dir"`
	OK            bool    `json:"ok"`
	Error         string  `json:"error,omitempty"`
	SizeBytes     int64   `json:"size_bytes,omitempty"`
	BlockBytes    int64   `json:"block_bytes,omitempty"`
	WrittenBytes  int64   `json:"written_bytes,omitempty"`
	ReadBytes     int64   `json:"read_bytes,omitempty"`
	WriteSeconds  float64 `json:"write_seconds,omitempty"`
	WriteMBPerSec float64 `json:"write_mb_per_sec,omitempty"`
	ReadSeconds   float64 `json:"read_seconds,omitempty"`
	ReadMBPerSec  float64 `json:"read_mb_per_sec,omitempty"`
}

// compareReport is the full multi-target comparison, in the exact order the
// targets were benchmarked (input order) for Targets, plus a Ranked view
// sorted by write MB/s descending among successful targets, with failures
// appended after.
type compareReport struct {
	SizeBytes  int64          `json:"size_bytes"`
	BlockBytes int64          `json:"block_bytes"`
	Targets    []targetResult `json:"targets"`
	Ranked     []string       `json:"ranked_dirs"`
}

func runBench(args []string) {
	for _, a := range args {
		if a == "-h" || a == "--help" || a == "help" {
			printBenchUsage()
			return
		}
	}

	valueFlags := map[string]bool{"size": true, "block": true}
	args = reorderFlags(args, valueFlags)

	fs := flag.NewFlagSet("bench", flag.ExitOnError)
	sizeStr := fs.String("size", "256MB", "total bytes to write/read per target (e.g. 256MB, 1GiB)")
	blockStr := fs.String("block", "1MB", "chunk size for each write/read call (e.g. 1MB, 256KB)")
	jsonOut := fs.Bool("json", false, "emit JSON report")
	fs.Usage = printBenchUsage
	fs.Parse(args)

	dirs := fs.Args()
	if len(dirs) < 1 {
		fmt.Fprintln(os.Stderr, "driveforge bench: missing <dir> argument(s) - at least one target directory is required")
		printBenchUsage()
		os.Exit(1)
	}

	size, err := parseSize(*sizeStr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "driveforge bench: --size: %v\n", err)
		os.Exit(1)
	}
	block, err := parseSize(*blockStr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "driveforge bench: --block: %v\n", err)
		os.Exit(1)
	}
	if size <= 0 {
		fmt.Fprintln(os.Stderr, "driveforge bench: --size must be greater than 0")
		os.Exit(1)
	}
	if block <= 0 {
		fmt.Fprintln(os.Stderr, "driveforge bench: --block must be greater than 0")
		os.Exit(1)
	}

	report := &compareReport{SizeBytes: size, BlockBytes: block}

	// Benchmark every target independently. A failure on one target is
	// recorded and benchmarking continues with the rest - one bad target
	// must never abort the whole comparison.
	for _, dir := range dirs {
		res := benchOne(dir, size, block)
		report.Targets = append(report.Targets, res)
	}

	report.Ranked = rankedDirs(report.Targets)

	successCount := 0
	for _, t := range report.Targets {
		if t.OK {
			successCount++
		}
	}

	if *jsonOut {
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		if err := enc.Encode(report); err != nil {
			fmt.Fprintf(os.Stderr, "driveforge bench: failed to encode JSON: %v\n", err)
			os.Exit(1)
		}
	} else {
		printReport(report)
	}

	if successCount == 0 {
		os.Exit(1)
	}
}

// rankedDirs returns the target directories sorted by write MB/s descending
// among the successful ones, followed by the failed ones (in their original
// order).
func rankedDirs(targets []targetResult) []string {
	var ok, failed []targetResult
	for _, t := range targets {
		if t.OK {
			ok = append(ok, t)
		} else {
			failed = append(failed, t)
		}
	}
	sort.SliceStable(ok, func(i, j int) bool {
		return ok[i].WriteMBPerSec > ok[j].WriteMBPerSec
	})
	var out []string
	for _, t := range ok {
		out = append(out, t.Dir)
	}
	for _, t := range failed {
		out = append(out, t.Dir)
	}
	return out
}

func printBenchUsage() {
	fmt.Fprint(os.Stderr, `Usage: driveforge bench <dir1> [<dir2> ...] [--size 256MB] [--block 1MB] [--json]

Runs the same sequential write-then-read throughput benchmark against a temp
file created inside each target directory, independently, then prints a
ranked comparison table sorted by write MB/s (fastest first).

Flags:
  --size    Total bytes to write/read per target, e.g. 256MB, 1GiB, 512KB
            (default 256MB). The same value is used for every target.
  --block   Chunk size for each write/read call, e.g. 1MB, 256KB
            (default 1MB). The same value is used for every target.
  --json    Emit a JSON report instead of human-readable text

Fault isolation:
  A target that doesn't exist or isn't writable is reported as FAILED with a
  reason and excluded from the ranking; the other targets still run to
  completion. The command only exits non-zero if EVERY target failed.

Caveats:
  - Sequential, single-threaded, file-API throughput only - not a raw
    block-device benchmark.
  - OS/filesystem caching can inflate results for small --size values,
    especially when multiple targets share the same underlying filesystem
    or page cache. Use at least a few hundred MB for a realistic number;
    smaller sizes are still fine for a quick sanity check.

Examples:
  driveforge bench /mnt/ssd /mnt/hdd --size 512MB
  driveforge bench . /tmp /mnt/data --size 1GiB --block 4MB --json
`)
}

func printReport(r *compareReport) {
	fmt.Printf("DriveForge multi-target benchmark comparison\n")
	fmt.Printf("  targets:      %d\n", len(r.Targets))
	fmt.Printf("  total size:   %s (%d bytes) per target\n", formatBytes(r.SizeBytes), r.SizeBytes)
	fmt.Printf("  block size:   %s (%d bytes) per target\n", formatBytes(r.BlockBytes), r.BlockBytes)
	fmt.Printf("\n")

	byDir := make(map[string]targetResult, len(r.Targets))
	for _, t := range r.Targets {
		byDir[t.Dir] = t
	}

	fmt.Printf("Rank  Write MB/s  Read MB/s   Target\n")
	fmt.Printf("----  ----------  ----------  ------\n")
	rank := 0
	for _, dir := range r.Ranked {
		t := byDir[dir]
		if !t.OK {
			continue
		}
		rank++
		fmt.Printf("%-4d  %10.2f  %10.2f  %s\n", rank, t.WriteMBPerSec, t.ReadMBPerSec, t.Dir)
	}
	failedAny := false
	for _, dir := range r.Ranked {
		t := byDir[dir]
		if t.OK {
			continue
		}
		failedAny = true
		fmt.Printf("--    %10s  %10s  %s  (FAILED: %s)\n", "-", "-", t.Dir, t.Error)
	}

	fmt.Printf("\n")
	if rank > 0 {
		fmt.Printf("Ranked by write MB/s descending (fastest write first).\n")
	}
	if failedAny {
		fmt.Printf("One or more targets failed and were excluded from the ranking; see reasons above.\n")
	}
	fmt.Printf("\n")
	fmt.Printf("Note: sequential, single-threaded, file-API throughput against each target\n")
	fmt.Printf("filesystem - not a raw block-device benchmark. The same --size/--block was\n")
	fmt.Printf("used for every target for an apples-to-apples comparison. Small --size values\n")
	fmt.Printf("can be inflated by OS/filesystem caching, especially when targets share the\n")
	fmt.Printf("same underlying filesystem or page cache.\n")
}

// benchOne runs the write/read benchmark against a single target directory
// and never panics or aborts the caller - all failure modes are captured in
// the returned targetResult with OK=false and a human-readable Error.
func benchOne(dir string, size, block int64) targetResult {
	res := targetResult{Dir: dir, SizeBytes: size, BlockBytes: block}

	info, err := os.Stat(dir)
	if err != nil {
		if os.IsNotExist(err) {
			res.Error = fmt.Sprintf("directory does not exist: %v", err)
		} else {
			res.Error = fmt.Sprintf("cannot access directory: %v", err)
		}
		return res
	}
	if !info.IsDir() {
		res.Error = fmt.Sprintf("%q is not a directory", dir)
		return res
	}

	path := filepath.Join(dir, tempFileName)

	// Verify the directory is writable by attempting to create the temp
	// file up front, before we start timing anything.
	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
	if err != nil {
		res.Error = fmt.Sprintf("cannot write to directory: %v", err)
		return res
	}
	// Always clean up the temp file, regardless of success or failure past
	// this point.
	defer func() {
		f.Close()
		os.Remove(path)
	}()

	// Pre-generate ONE random block and reuse it for every chunk written.
	// crypto/rand is used (not math/rand) so the data is genuinely
	// incompressible and can't be skewed by filesystem-level compression or
	// dedup. The random-generation cost is paid once, up front, outside the
	// timer, and the same block is written repeatedly, so it doesn't
	// artificially cap the reported throughput.
	blockBuf := make([]byte, block)
	if _, err := rand.Read(blockBuf); err != nil {
		res.Error = fmt.Sprintf("failed to generate random data: %v", err)
		return res
	}

	// Write phase: write `size` bytes total in `block`-sized chunks, using
	// the same pre-generated buffer each time. Timing includes f.Sync() at
	// the end, since without it we'd only be measuring how fast data lands
	// in the page cache, not how fast it reaches stable storage.
	writeStart := time.Now()
	var written int64
	for written < size {
		n := block
		if remaining := size - written; remaining < n {
			n = remaining
		}
		if _, err := f.Write(blockBuf[:n]); err != nil {
			res.Error = fmt.Sprintf("write failed after %d bytes: %v", written, err)
			return res
		}
		written += n
	}
	if err := f.Sync(); err != nil {
		res.Error = fmt.Sprintf("fsync failed: %v", err)
		return res
	}
	writeDur := time.Since(writeStart)
	res.WriteSeconds = writeDur.Seconds()
	res.WriteMBPerSec = mbPerSec(written, writeDur)
	res.WrittenBytes = written

	if err := f.Close(); err != nil {
		res.Error = fmt.Sprintf("failed to close file after write: %v", err)
		return res
	}

	// Read phase: reopen the file and read it back sequentially in
	// `block`-sized chunks, discarding the data into a reused buffer.
	rf, err := os.Open(path)
	if err != nil {
		res.Error = fmt.Sprintf("failed to reopen file for read: %v", err)
		return res
	}
	defer rf.Close()

	readBuf := make([]byte, block)
	readStart := time.Now()
	var readBytes int64
	for {
		n, err := rf.Read(readBuf)
		readBytes += int64(n)
		if err == io.EOF {
			break
		}
		if err != nil {
			res.Error = fmt.Sprintf("read failed after %d bytes: %v", readBytes, err)
			return res
		}
	}
	readDur := time.Since(readStart)
	res.ReadSeconds = readDur.Seconds()
	res.ReadMBPerSec = mbPerSec(readBytes, readDur)
	res.ReadBytes = readBytes

	if readBytes != written {
		res.Error = fmt.Sprintf("read back %d bytes but wrote %d bytes", readBytes, written)
		return res
	}

	res.OK = true
	return res
}

func mbPerSec(bytes int64, d time.Duration) float64 {
	if d <= 0 {
		return 0
	}
	mb := float64(bytes) / (1000 * 1000)
	return mb / d.Seconds()
}