// Command driverguard is a file-integrity watchdog with real restore-point
// copies for any critical directory (a drivers folder, a system config
// folder, or anything else you point it at).
//
// It snapshots a directory (hash + an actual backup copy of every file's
// current bytes), detects changes against that snapshot later, and can
// restore a single changed or deleted file from the snapshot's backup copy
// on request. See README.txt for the full scope and rationale, and
// ../plan.md for the full DriverGuard product plan (which additionally
// covers real Windows driver-store inventory and OS-level driver rollback
// via privileged SetupAPI access — not implemented in this prototype).
package main

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

const manifestFileName = "manifest.json"
const filesDirName = "files"

// ---------- shared helpers ----------

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...)
}

func humanBytes(n int64) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := int64(unit), 0
	for x := n / unit; x >= unit; x /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(n)/float64(div), "KMGTPE"[exp])
}

func usage() {
	fmt.Fprint(os.Stderr, `driverguard - file-integrity watchdog with real restore-point copies

Usage:
  driverguard snapshot <dir> --store <storedir> [--label mylabel]
  driverguard check <dir> --store <storedir> [--against <label-or-latest>] [--json]
  driverguard rollback <dir> --store <storedir> --file <relative-path> [--against <label-or-latest>] [--apply]
  driverguard list --store <storedir> [--json]

Commands:
  snapshot   Walk <dir> recursively, hash every regular file, and copy each
             file's current bytes into <storedir> as a restorable snapshot.
  check      Compare <dir>'s current state against a stored snapshot and
             report CHANGED / REMOVED / ADDED files.
  rollback   Restore one file at <dir>/<relative-path> from a snapshot's
             backup copy. Dry-run by default; pass --apply to actually write.
  list       List all snapshots present in <storedir>.

Run 'driverguard <command> -h' for command-specific flags.
`)
}

func fatalf(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "driverguard: "+format+"\n", args...)
	os.Exit(1)
}

func isHelp(arg string) bool {
	return arg == "-h" || arg == "--help" || arg == "help"
}

// ---------- manifest types ----------

type fileEntry struct {
	Path   string `json:"path"`
	Size   int64  `json:"size"`
	SHA256 string `json:"sha256"`
}

type manifest struct {
	CreatedAtUTC string      `json:"created_at_utc"`
	SourceDir    string      `json:"source_dir"`
	Label        string      `json:"label"`
	Files        []fileEntry `json:"files"`
}

func snapshotDir(store, label string) string {
	return filepath.Join(store, label)
}

func manifestPath(store, label string) string {
	return filepath.Join(snapshotDir(store, label), manifestFileName)
}

func loadManifest(store, label string) (*manifest, error) {
	data, err := os.ReadFile(manifestPath(store, label))
	if err != nil {
		return nil, err
	}
	var m manifest
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, fmt.Errorf("parsing manifest for %q: %w", label, err)
	}
	return &m, nil
}

// listSnapshotLabels returns every snapshot label found directly under
// store (a directory containing a manifest.json), sorted by CreatedAtUTC
// ascending.
func listSnapshotLabels(store string) ([]string, error) {
	entries, err := os.ReadDir(store)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	type ls struct {
		label   string
		created string
	}
	var found []ls
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		if _, err := os.Stat(manifestPath(store, e.Name())); err != nil {
			continue
		}
		m, err := loadManifest(store, e.Name())
		if err != nil {
			continue
		}
		found = append(found, ls{label: e.Name(), created: m.CreatedAtUTC})
	}
	sort.Slice(found, func(i, j int) bool { return found[i].created < found[j].created })
	labels := make([]string, len(found))
	for i, f := range found {
		labels[i] = f.label
	}
	return labels, nil
}

func resolveLabel(store, against string) (string, error) {
	labels, err := listSnapshotLabels(store)
	if err != nil {
		return "", err
	}
	if len(labels) == 0 {
		return "", fmt.Errorf("no snapshots found in store %q", store)
	}
	if against == "" || against == "latest" {
		return labels[len(labels)-1], nil
	}
	for _, l := range labels {
		if l == against {
			return l, nil
		}
	}
	return "", fmt.Errorf("snapshot %q not found in store %q", against, store)
}

// ---------- hashing / walking ----------

func hashFile(path string) (string, int64, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", 0, err
	}
	defer f.Close()
	h := sha256.New()
	n, err := io.Copy(h, f)
	if err != nil {
		return "", 0, err
	}
	return hex.EncodeToString(h.Sum(nil)), n, nil
}

// walkFiles returns a sorted list of paths (relative to root, using forward
// slashes) for every regular file under root.
func walkFiles(root string) ([]string, error) {
	var rels []string
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		if !info.Mode().IsRegular() {
			return nil
		}
		rel, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		rels = append(rels, filepath.ToSlash(rel))
		return nil
	})
	if err != nil {
		return nil, err
	}
	sort.Strings(rels)
	return rels, nil
}

func copyFile(src, dst string) error {
	if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
		return err
	}
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	tmp := dst + ".driverguard-tmp"
	out, err := os.Create(tmp)
	if err != nil {
		return err
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		os.Remove(tmp)
		return err
	}
	if err := out.Close(); err != nil {
		os.Remove(tmp)
		return err
	}
	return os.Rename(tmp, dst)
}

// ---------- commands ----------

func cmdSnapshot(args []string) {
	fs := flag.NewFlagSet("snapshot", flag.ExitOnError)
	store := fs.String("store", "", "snapshot store directory (required)")
	label := fs.String("label", "", "snapshot label (default: UTC timestamp)")
	fs.Usage = func() {
		fmt.Fprint(os.Stderr, `Usage: driverguard snapshot <dir> --store <storedir> [--label mylabel]

Walks <dir> recursively (regular files only), computes a SHA-256 for each
file, and copies each file's current bytes into the snapshot store so it
can be restored later even if the original is changed or deleted.
`)
	}
	args = reorderFlags(args, map[string]bool{"store": true, "label": true})
	fs.Parse(args)

	pos := fs.Args()
	if len(pos) != 1 || *store == "" {
		fs.Usage()
		os.Exit(1)
	}
	dir := pos[0]

	info, err := os.Stat(dir)
	if err != nil {
		fatalf("cannot access %q: %v", dir, err)
	}
	if !info.IsDir() {
		fatalf("%q is not a directory", dir)
	}

	lbl := *label
	now := time.Now().UTC()
	if lbl == "" {
		lbl = now.Format("20060102-150405")
	}

	dstDir := snapshotDir(*store, lbl)
	if _, err := os.Stat(dstDir); err == nil {
		fatalf("snapshot %q already exists in %q", lbl, *store)
	}

	rels, err := walkFiles(dir)
	if err != nil {
		fatalf("walking %q: %v", dir, err)
	}

	filesRoot := filepath.Join(dstDir, filesDirName)
	if err := os.MkdirAll(filesRoot, 0o755); err != nil {
		fatalf("creating snapshot store: %v", err)
	}

	var entries []fileEntry
	var totalBytes int64
	for _, rel := range rels {
		srcPath := filepath.Join(dir, filepath.FromSlash(rel))
		sum, size, err := hashFile(srcPath)
		if err != nil {
			fatalf("hashing %q: %v", srcPath, err)
		}
		dstPath := filepath.Join(filesRoot, filepath.FromSlash(rel))
		if err := copyFile(srcPath, dstPath); err != nil {
			fatalf("backing up %q: %v", srcPath, err)
		}
		entries = append(entries, fileEntry{Path: rel, Size: size, SHA256: sum})
		totalBytes += size
	}

	absDir, err := filepath.Abs(dir)
	if err != nil {
		absDir = dir
	}
	m := manifest{
		CreatedAtUTC: now.Format(time.RFC3339),
		SourceDir:    absDir,
		Label:        lbl,
		Files:        entries,
	}
	data, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		fatalf("encoding manifest: %v", err)
	}
	if err := os.WriteFile(manifestPath(*store, lbl), data, 0o644); err != nil {
		fatalf("writing manifest: %v", err)
	}

	fmt.Printf("Snapshot %q created\n", lbl)
	fmt.Printf("  source:  %s\n", absDir)
	fmt.Printf("  store:   %s\n", dstDir)
	fmt.Printf("  files:   %d\n", len(entries))
	fmt.Printf("  copied:  %s\n", humanBytes(totalBytes))
}

type checkResult struct {
	Snapshot string   `json:"snapshot"`
	SourceOK bool     `json:"source_ok"`
	Changed  []string `json:"changed"`
	Removed  []string `json:"removed"`
	Added    []string `json:"added"`
}

func cmdCheck(args []string) {
	if runCheck(args) {
		os.Exit(1)
	}
}

// runCheck does the whole of "check" and reports whether the directory has
// diverged from the snapshot, instead of exiting on the spot.
//
// The exit code is the caller's business, and one caller cannot afford it: the
// guided session that runs when the program is double-clicked has to print
// "Press Enter to close this window" after the report, and os.Exit here would
// close the window before the reader saw a word of it. From the command line
// the behaviour is unchanged — cmdCheck still exits 1 on divergence.
func runCheck(args []string) bool {
	fs := flag.NewFlagSet("check", flag.ExitOnError)
	store := fs.String("store", "", "snapshot store directory (required)")
	against := fs.String("against", "latest", "snapshot label to compare against, or 'latest'")
	asJSON := fs.Bool("json", false, "output JSON")
	fs.Usage = func() {
		fmt.Fprint(os.Stderr, `Usage: driverguard check <dir> --store <storedir> [--against <label-or-latest>] [--json]

Compares <dir>'s current state against a stored snapshot: what changed
since you last trusted this directory. Reports CHANGED, REMOVED, and ADDED
files. Exits non-zero if there are any CHANGED or REMOVED files.
`)
	}
	args = reorderFlags(args, map[string]bool{"store": true, "against": true})
	fs.Parse(args)

	pos := fs.Args()
	if len(pos) != 1 || *store == "" {
		fs.Usage()
		os.Exit(1)
	}
	dir := pos[0]

	label, err := resolveLabel(*store, *against)
	if err != nil {
		fatalf("%v", err)
	}
	m, err := loadManifest(*store, label)
	if err != nil {
		fatalf("loading snapshot %q: %v", label, err)
	}

	snapshotIndex := make(map[string]fileEntry, len(m.Files))
	for _, fe := range m.Files {
		snapshotIndex[fe.Path] = fe
	}

	currentRels, err := walkFiles(dir)
	if err != nil {
		fatalf("walking %q: %v", dir, err)
	}
	currentSet := make(map[string]bool, len(currentRels))
	for _, r := range currentRels {
		currentSet[r] = true
	}

	var changed, removed, added []string

	for _, rel := range currentRels {
		if fe, ok := snapshotIndex[rel]; ok {
			sum, _, err := hashFile(filepath.Join(dir, filepath.FromSlash(rel)))
			if err != nil {
				fatalf("hashing %q: %v", rel, err)
			}
			if sum != fe.SHA256 {
				changed = append(changed, rel)
			}
		} else {
			added = append(added, rel)
		}
	}
	for path := range snapshotIndex {
		if !currentSet[path] {
			removed = append(removed, path)
		}
	}
	sort.Strings(changed)
	sort.Strings(removed)
	sort.Strings(added)

	if *asJSON {
		res := checkResult{
			Snapshot: label,
			SourceOK: len(changed) == 0 && len(removed) == 0,
			Changed:  nonNil(changed),
			Removed:  nonNil(removed),
			Added:    nonNil(added),
		}
		data, _ := json.MarshalIndent(res, "", "  ")
		fmt.Println(string(data))
	} else {
		fmt.Printf("Checked %q against snapshot %q\n", dir, label)
		fmt.Printf("  CHANGED (%d):\n", len(changed))
		for _, p := range changed {
			fmt.Printf("    ~ %s\n", p)
		}
		fmt.Printf("  REMOVED (%d):\n", len(removed))
		for _, p := range removed {
			fmt.Printf("    - %s\n", p)
		}
		fmt.Printf("  ADDED (%d):\n", len(added))
		for _, p := range added {
			fmt.Printf("    + %s\n", p)
		}
		if len(changed) == 0 && len(removed) == 0 {
			fmt.Println("Result: directory matches the trusted snapshot (no changed or removed files).")
		} else {
			fmt.Println("Result: directory has diverged from the trusted snapshot.")
		}
	}

	return len(changed) > 0 || len(removed) > 0
}

func nonNil(s []string) []string {
	if s == nil {
		return []string{}
	}
	return s
}

func cmdRollback(args []string) {
	fs := flag.NewFlagSet("rollback", flag.ExitOnError)
	store := fs.String("store", "", "snapshot store directory (required)")
	against := fs.String("against", "latest", "snapshot label to restore from, or 'latest'")
	file := fs.String("file", "", "relative path (within <dir>) of the file to restore (required)")
	apply := fs.Bool("apply", false, "actually write the restored file (default is a dry run)")
	fs.Usage = func() {
		fmt.Fprint(os.Stderr, `Usage: driverguard rollback <dir> --store <storedir> --file <relative-path> [--against <label-or-latest>] [--apply]

Restores one file at <dir>/<relative-path> from a snapshot's backup copy.
Without --apply this is a dry run that reports the plan only. With --apply
the file is actually overwritten (or recreated) from the snapshot's stored
bytes.
`)
	}
	args = reorderFlags(args, map[string]bool{"store": true, "against": true, "file": true})
	fs.Parse(args)

	pos := fs.Args()
	if len(pos) != 1 || *store == "" || *file == "" {
		fs.Usage()
		os.Exit(1)
	}
	dir := pos[0]
	rel := filepath.ToSlash(*file)

	label, err := resolveLabel(*store, *against)
	if err != nil {
		fatalf("%v", err)
	}
	m, err := loadManifest(*store, label)
	if err != nil {
		fatalf("loading snapshot %q: %v", label, err)
	}

	var target *fileEntry
	for i := range m.Files {
		if m.Files[i].Path == rel {
			target = &m.Files[i]
			break
		}
	}
	if target == nil {
		fatalf("file %q is not present in snapshot %q — nothing to roll back", rel, label)
	}

	backupPath := filepath.Join(snapshotDir(*store, label), filesDirName, filepath.FromSlash(rel))
	currentPath := filepath.Join(dir, filepath.FromSlash(rel))

	status := "UNCHANGED"
	if curInfo, err := os.Stat(currentPath); err != nil {
		if os.IsNotExist(err) {
			status = "REMOVED"
		} else {
			fatalf("checking %q: %v", currentPath, err)
		}
	} else if curInfo.Mode().IsRegular() {
		sum, _, err := hashFile(currentPath)
		if err != nil {
			fatalf("hashing %q: %v", currentPath, err)
		}
		if sum != target.SHA256 {
			status = "CHANGED"
		}
	}

	if !*apply {
		fmt.Printf("Dry run — no files were modified.\n")
		fmt.Printf("  snapshot:      %s\n", label)
		fmt.Printf("  file:          %s\n", rel)
		fmt.Printf("  current state: %s\n", status)
		switch status {
		case "REMOVED":
			fmt.Printf("  plan:          would recreate REMOVED file at %s (%s, sha256 %s)\n", currentPath, humanBytes(target.Size), target.SHA256)
		case "CHANGED":
			fmt.Printf("  plan:          would overwrite CHANGED file at %s with snapshot version (%s, sha256 %s)\n", currentPath, humanBytes(target.Size), target.SHA256)
		default:
			fmt.Printf("  plan:          would overwrite %s at %s with snapshot version (%s, sha256 %s) — no functional change expected\n", status, currentPath, humanBytes(target.Size), target.SHA256)
		}
		fmt.Println("Re-run with --apply to perform the restore.")
		return
	}

	if err := copyFile(backupPath, currentPath); err != nil {
		fatalf("restoring %q: %v", currentPath, err)
	}
	fmt.Printf("Restored %q from snapshot %q (was %s)\n", currentPath, label, status)
	fmt.Printf("  size:   %s\n", humanBytes(target.Size))
	fmt.Printf("  sha256: %s\n", target.SHA256)
}

type snapshotSummary struct {
	Label      string `json:"label"`
	CreatedUTC string `json:"created_at_utc"`
	SourceDir  string `json:"source_dir"`
	FileCount  int    `json:"file_count"`
	TotalBytes int64  `json:"total_bytes"`
}

func cmdList(args []string) {
	fs := flag.NewFlagSet("list", flag.ExitOnError)
	store := fs.String("store", "", "snapshot store directory (required)")
	asJSON := fs.Bool("json", false, "output JSON")
	fs.Usage = func() {
		fmt.Fprint(os.Stderr, `Usage: driverguard list --store <storedir> [--json]

Lists all snapshots present in <storedir> with label, creation time, file
count, and total backed-up bytes.
`)
	}
	args = reorderFlags(args, map[string]bool{"store": true})
	fs.Parse(args)

	if fs.NArg() != 0 || *store == "" {
		fs.Usage()
		os.Exit(1)
	}

	labels, err := listSnapshotLabels(*store)
	if err != nil {
		fatalf("reading store %q: %v", *store, err)
	}

	var summaries []snapshotSummary
	for _, lbl := range labels {
		m, err := loadManifest(*store, lbl)
		if err != nil {
			continue
		}
		var total int64
		for _, fe := range m.Files {
			total += fe.Size
		}
		summaries = append(summaries, snapshotSummary{
			Label:      lbl,
			CreatedUTC: m.CreatedAtUTC,
			SourceDir:  m.SourceDir,
			FileCount:  len(m.Files),
			TotalBytes: total,
		})
	}

	if *asJSON {
		data, _ := json.MarshalIndent(summaries, "", "  ")
		fmt.Println(string(data))
		return
	}

	if len(summaries) == 0 {
		fmt.Printf("No snapshots found in %q\n", *store)
		return
	}
	fmt.Printf("%-20s %-25s %-8s %-10s %s\n", "LABEL", "CREATED (UTC)", "FILES", "SIZE", "SOURCE")
	for _, s := range summaries {
		fmt.Printf("%-20s %-25s %-8d %-10s %s\n", s.Label, s.CreatedUTC, s.FileCount, humanBytes(s.TotalBytes), s.SourceDir)
	}
}

// ---------- main ----------

func main() {
	if len(os.Args) < 2 {
		// Double-clicked in Explorer rather than run from a prompt: ask the
		// questions 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)
	}

	cmd := os.Args[1]
	if isHelp(cmd) {
		usage()
		return
	}

	rest := os.Args[2:]
	switch cmd {
	case "snapshot":
		if len(rest) > 0 && isHelp(rest[0]) {
			cmdSnapshot([]string{"-h"})
			return
		}
		cmdSnapshot(rest)
	case "check":
		if len(rest) > 0 && isHelp(rest[0]) {
			cmdCheck([]string{"-h"})
			return
		}
		cmdCheck(rest)
	case "rollback":
		if len(rest) > 0 && isHelp(rest[0]) {
			cmdRollback([]string{"-h"})
			return
		}
		cmdRollback(rest)
	case "list":
		if len(rest) > 0 && isHelp(rest[0]) {
			cmdList([]string{"-h"})
			return
		}
		cmdList(rest)
	default:
		fmt.Fprintf(os.Stderr, "driverguard: unknown command %q\n\n", cmd)
		usage()
		os.Exit(1)
	}
}