// Command diskops is a small, honest prototype of a disk-utility CLI.
//
// DiskOps's full product concept (see ../plan.md) covers partitioning,
// cloning, migration and SMART awareness. Those all 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). This prototype instead implements the one piece of the
// concept that genuinely works cross-platform with only the stdlib: a
// sequential disk I/O throughput benchmark, in the spirit of the sequential
// numbers reported by tools like CrystalDiskMark, measured through the
// ordinary file API rather than a raw device.
package main

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

const tempFileName = ".diskops-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, "diskops: unknown command %q\n\n", os.Args[1])
		usage()
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprint(os.Stderr, `DiskOps - sequential disk I/O throughput benchmark (prototype)

Usage:
  diskops bench <dir> [--size 256MB] [--block 1MB] [--json]
  diskops help

Commands:
  bench    Write and read a temp file in <dir>, sequentially, single-threaded,
           and report throughput in MB/s.

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

Notes:
  - This measures sequential, single-threaded throughput through the OS file
    API against the target filesystem/directory. It is NOT a raw block-device
    benchmark (no partitioning, cloning, migration or SMART data here - those
    need OS-privileged/raw-device access and are on the DiskOps roadmap, see
    ../plan.md).
  - OS and filesystem caching mean small --size values can report unrealistic
    numbers. Use at least a few hundred MB for a meaningful result; smaller
    sizes are still fine for a quick sanity check.

Examples:
  diskops bench . --size 512MB
  diskops bench /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. `diskops bench . --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])
}

type benchReport struct {
	Dir           string  `json:"dir"`
	SizeBytes     int64   `json:"size_bytes"`
	BlockBytes    int64   `json:"block_bytes"`
	WriteSeconds  float64 `json:"write_seconds"`
	WriteMBPerSec float64 `json:"write_mb_per_sec"`
	ReadSeconds   float64 `json:"read_seconds"`
	ReadMBPerSec  float64 `json:"read_mb_per_sec"`
}

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 (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)

	positional := fs.Args()
	if len(positional) < 1 {
		fmt.Fprintln(os.Stderr, "diskops bench: missing <dir> argument")
		printBenchUsage()
		os.Exit(1)
	}
	dir := positional[0]

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

	report, err := bench(dir, size, block)
	if err != nil {
		fmt.Fprintf(os.Stderr, "diskops bench: %v\n", err)
		os.Exit(1)
	}

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

	printReport(report)
}

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

Runs a sequential write-then-read throughput benchmark against a temp file
created inside <dir>, and reports MB/s for each phase.

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

Caveats:
  - Sequential, single-threaded, file-API throughput only - not a raw
    block-device benchmark.
  - OS/filesystem caching can inflate results for small --size values. Use at
    least a few hundred MB for a realistic number; smaller sizes are still
    useful for a quick sanity check on constrained environments.
`)
}

func printReport(r *benchReport) {
	fmt.Printf("DiskOps sequential benchmark\n")
	fmt.Printf("  target dir:   %s\n", r.Dir)
	fmt.Printf("  total size:   %s (%d bytes)\n", formatBytes(r.SizeBytes), r.SizeBytes)
	fmt.Printf("  block size:   %s (%d bytes)\n", formatBytes(r.BlockBytes), r.BlockBytes)
	fmt.Printf("  write:        %8.2f MB/s  (%.3fs, includes fsync)\n", r.WriteMBPerSec, r.WriteSeconds)
	fmt.Printf("  read:         %8.2f MB/s  (%.3fs)\n", r.ReadMBPerSec, r.ReadSeconds)
	fmt.Printf("\n")
	fmt.Printf("Note: sequential, single-threaded, file-API throughput against the target\n")
	fmt.Printf("filesystem - not a raw block-device benchmark. Small --size values can be\n")
	fmt.Printf("inflated by OS/filesystem caching; use a few hundred MB or more for a\n")
	fmt.Printf("meaningful number.\n")
}

// bench performs the actual write/read benchmark. It always removes the temp
// file it creates, including on error paths, via defer.
func bench(dir string, size, block int64) (*benchReport, error) {
	info, err := os.Stat(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, fmt.Errorf("directory %q does not exist", dir)
		}
		return nil, fmt.Errorf("cannot access directory %q: %w", dir, err)
	}
	if !info.IsDir() {
		return nil, fmt.Errorf("%q is not a directory", dir)
	}

	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 {
		return nil, fmt.Errorf("cannot write to directory %q: %w", dir, err)
	}
	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. Generating a fresh random block per chunk would be correct too,
	// but crypto/rand generation is slow enough to dominate the timed write
	// loop and artificially cap the reported throughput - so we pay that
	// cost once, up front, outside the timer, and write the same block
	// repeatedly instead.
	blockBuf := make([]byte, block)
	if _, err := rand.Read(blockBuf); err != nil {
		return nil, fmt.Errorf("failed to generate random data: %w", err)
	}

	report := &benchReport{
		Dir:        dir,
		SizeBytes:  size,
		BlockBytes: block,
	}

	// 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 {
			return nil, fmt.Errorf("write failed after %d bytes: %w", written, err)
		}
		written += n
	}
	if err := f.Sync(); err != nil {
		return nil, fmt.Errorf("fsync failed: %w", err)
	}
	writeDur := time.Since(writeStart)
	report.WriteSeconds = writeDur.Seconds()
	report.WriteMBPerSec = mbPerSec(written, writeDur)

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

	// 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 {
		return nil, fmt.Errorf("failed to reopen file for read: %w", err)
	}
	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 {
			return nil, fmt.Errorf("read failed after %d bytes: %w", readBytes, err)
		}
	}
	readDur := time.Since(readStart)
	report.ReadSeconds = readDur.Seconds()
	report.ReadMBPerSec = mbPerSec(readBytes, readDur)

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

	return report, nil
}

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