// Command photosweep finds near-duplicate / visually-similar photos using a
// perceptual average-hash (aHash) plus Hamming-distance grouping, and can
// safely quarantine the extras. It is the "Pro" companion to the sibling
// tool DupePilot: DupePilot finds byte-identical files via SHA-256 hashing;
// PhotoSweep finds photos that LOOK the same even when their bytes, format,
// or dimensions differ (a re-export, a burst-shot near-duplicate, a resize,
// a re-compression). See README.txt for the full write-up.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"math/bits"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// ---------------------------------------------------------------------------
// CLI plumbing
// ---------------------------------------------------------------------------

func usage() {
	fmt.Fprint(os.Stderr, `PhotoSweep - near-duplicate photo detector (perceptual aHash)

Usage:
  photosweep scan <dir> [<dir> ...] [--threshold 5] [--min-size 10KB] [--json]
  photosweep clean <dir> [<dir> ...] --quarantine <qdir> [--threshold 5] [--min-size 10KB] [--apply]
  photosweep help

scan
  Walks the given directories, computes a 64-bit average-hash (aHash) for
  every recognized image (.jpg, .jpeg, .png, .gif), and groups images whose
  Hamming distance is <= --threshold into near-duplicate sets. Only groups
  with 2+ members are reported.

    --threshold N   max Hamming distance (0-64) to count as near-duplicate (default 5)
    --min-size S    skip files smaller than S, e.g. 10KB, 2MB (default: no minimum)
    --json          emit machine-readable JSON instead of text

clean
  Same grouping as scan. For each near-duplicate group, the FIRST member
  found (by scan order) is KEPT; every other member is a cleanup candidate.
  Without --apply this is a dry run that only prints the plan. With --apply,
  every non-kept file is MOVED (never deleted) into --quarantine, preserving
  its path relative to the scanned root directory it was found under.

    --quarantine D  destination directory for quarantined files (required)
    --threshold N   max Hamming distance to count as near-duplicate (default 5)
    --min-size S    skip files smaller than S, e.g. 10KB, 2MB (default: no minimum)
    --apply         actually move files (default: dry run, no changes made)

Examples:
  photosweep scan ~/Pictures --threshold 8
  photosweep clean ~/Pictures --quarantine ~/PhotoSweepQuarantine --apply
`)
}

// reorderFlags works around a stdlib flag package quirk: flag.Parse stops
// scanning for flags at the first positional argument. Subcommands here take
// positional directory arguments that may appear before or between flags, so
// we reorder argv into "all flags, then all positionals" before handing it
// to flag.FlagSet.Parse.
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])
}

// parseSize parses human-friendly sizes like "10KB", "2MB", "512", "1GiB".
// Units are treated as base-1024 (KB == KiB) to match humanBytes' output.
func parseSize(s string) (int64, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return 0, nil
	}
	upper := strings.ToUpper(s)
	suffixes := []struct {
		suf  string
		mult int64
	}{
		{"KIB", 1024}, {"MIB", 1024 * 1024}, {"GIB", 1024 * 1024 * 1024},
		{"KB", 1024}, {"MB", 1024 * 1024}, {"GB", 1024 * 1024 * 1024},
		{"K", 1024}, {"M", 1024 * 1024}, {"G", 1024 * 1024 * 1024},
		{"B", 1},
	}
	for _, sfx := range suffixes {
		if strings.HasSuffix(upper, sfx.suf) {
			numPart := strings.TrimSpace(upper[:len(upper)-len(sfx.suf)])
			if numPart == "" {
				return 0, fmt.Errorf("invalid size %q", s)
			}
			f, err := strconv.ParseFloat(numPart, 64)
			if err != nil {
				return 0, fmt.Errorf("invalid size %q: %w", s, err)
			}
			return int64(f * float64(sfx.mult)), nil
		}
	}
	n, err := strconv.ParseInt(upper, 10, 64)
	if err != nil {
		return 0, fmt.Errorf("invalid size %q", s)
	}
	return n, nil
}

// ---------------------------------------------------------------------------
// Perceptual average-hash (aHash)
// ---------------------------------------------------------------------------
//
// aHash is a simple, well-known, genuinely effective perceptual hash:
//
//  1. Decode the image (image.Decode auto-detects JPEG/PNG/GIF because we
//     blank-import their decoders below, which register themselves).
//  2. Downsample to a fixed 8x8 = 64 cell grid. We hand-roll a box-average
//     downsample: the source image is divided into 64 evenly-spaced
//     rectangular blocks (bounds computed independently per axis, so this
//     works for any source resolution/aspect ratio) and every pixel in a
//     block is averaged together. This is more robust to resizing/
//     recompression noise than plain nearest-neighbor sampling.
//  3. Convert each block's averaged RGB to grayscale luminance using the
//     standard broadcast/ITU-R BT.601 formula 0.299R + 0.587G + 0.114B,
//     applied by hand (rather than color.GrayModel.Convert per source
//     pixel) so we can average in RGB space first and convert once per
//     block - mathematically equivalent for this purpose, cheaper.
//  4. Compute the mean brightness across all 64 cells.
//  5. Emit a 64-bit hash: bit i = 1 if cell i's brightness >= mean, else 0.
//
// Two images are "near-duplicates" if the Hamming distance between their
// 64-bit hashes is small (default threshold: 5 of 64 bits). aHash is good
// at catching resizes, re-compressions, format conversions, and minor
// edits/crops - exactly DupePilot's blind spot, since those all change the
// file's bytes (and therefore its SHA-256) while leaving it visually almost
// identical. It is NOT as robust as DCT-based perceptual hashes (pHash/
// dHash) or ML embedding similarity, which handle rotation, heavier crops,
// and watermarks much better - those are a documented roadmap item, not
// implemented here.

const gridSize = 8 // 8x8 = 64-bit hash

func computeAHash(img image.Image) uint64 {
	bounds := img.Bounds()
	w := bounds.Dx()
	h := bounds.Dy()
	if w <= 0 || h <= 0 {
		return 0
	}

	var cellLuminance [gridSize * gridSize]float64

	for gy := 0; gy < gridSize; gy++ {
		y0 := bounds.Min.Y + (gy*h)/gridSize
		y1 := bounds.Min.Y + ((gy+1)*h)/gridSize
		if y1 <= y0 {
			y1 = y0 + 1
		}
		if y1 > bounds.Max.Y {
			y1 = bounds.Max.Y
		}
		for gx := 0; gx < gridSize; gx++ {
			x0 := bounds.Min.X + (gx*w)/gridSize
			x1 := bounds.Min.X + ((gx+1)*w)/gridSize
			if x1 <= x0 {
				x1 = x0 + 1
			}
			if x1 > bounds.Max.X {
				x1 = bounds.Max.X
			}

			var rSum, gSum, bSum float64
			var count int64
			for y := y0; y < y1; y++ {
				for x := x0; x < x1; x++ {
					r, g, b, _ := img.At(x, y).RGBA() // 16-bit channel values
					rSum += float64(r)
					gSum += float64(g)
					bSum += float64(b)
					count++
				}
			}
			if count == 0 {
				count = 1
			}
			rAvg := rSum / float64(count)
			gAvg := gSum / float64(count)
			bAvg := bSum / float64(count)
			// Standard luminance formula (BT.601). The absolute scale
			// doesn't matter here (channels are 16-bit), only the
			// relative brightness across cells does.
			cellLuminance[gy*gridSize+gx] = 0.299*rAvg + 0.587*gAvg + 0.114*bAvg
		}
	}

	var mean float64
	for _, v := range cellLuminance {
		mean += v
	}
	mean /= float64(len(cellLuminance))

	var hash uint64
	for i, v := range cellLuminance {
		if v >= mean {
			hash |= 1 << uint(i)
		}
	}
	return hash
}

func hammingDistance(a, b uint64) int {
	return bits.OnesCount64(a ^ b)
}

// ---------------------------------------------------------------------------
// Scanning
// ---------------------------------------------------------------------------

var imageExts = map[string]bool{
	".jpg":  true,
	".jpeg": true,
	".png":  true,
	".gif":  true,
}

type imgEntry struct {
	path   string
	root   string // the scanned root directory this file was found under
	hash   uint64
	width  int
	height int
	size   int64
}

type dupGroup struct {
	members []imgEntry
}

// walkImages finds every file under roots whose extension looks like an
// image, decodes it, and returns one imgEntry per successfully-decoded
// image. Files that fail to decode (corrupt, or a non-image with a matching
// extension) are skipped with a warning on stderr rather than aborting the
// whole scan.
func walkImages(roots []string, minSize int64) (entries []imgEntry, totalCandidates int, warnings []string) {
	for _, root := range roots {
		absRoot, err := filepath.Abs(root)
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("skipping %s: %v", root, err))
			continue
		}
		err = filepath.WalkDir(absRoot, func(path string, d os.DirEntry, err error) error {
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("%s: %v", path, err))
				return nil
			}
			if d.IsDir() {
				return nil
			}
			ext := strings.ToLower(filepath.Ext(path))
			if !imageExts[ext] {
				return nil
			}
			totalCandidates++

			info, err := d.Info()
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("%s: %v", path, err))
				return nil
			}
			if info.Size() < minSize {
				return nil
			}

			f, err := os.Open(path)
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("%s: %v", path, err))
				return nil
			}
			img, _, err := image.Decode(f)
			f.Close()
			if err != nil {
				warnings = append(warnings, fmt.Sprintf("%s: could not decode as an image (%v) - skipped", path, err))
				return nil
			}

			bounds := img.Bounds()
			entries = append(entries, imgEntry{
				path:   path,
				root:   absRoot,
				hash:   computeAHash(img),
				width:  bounds.Dx(),
				height: bounds.Dy(),
				size:   info.Size(),
			})
			return nil
		})
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("%s: %v", root, err))
		}
	}
	return entries, totalCandidates, warnings
}

// groupNearDuplicates buckets entries into near-duplicate groups using
// simple pairwise Hamming-distance comparison: for each new image, compare
// its hash against every hash in every existing group so far; if it is
// within threshold of ANY member of a group, it joins that group, otherwise
// it starts a new one. This is O(n^2) in the number of images, which is
// fine for a typical photo folder (hundreds to low thousands of files) but
// would not scale to a huge library - a real product would want a proper
// nearest-neighbor index (e.g. an LSH / BK-tree over Hamming space) instead.
// That is a roadmap item, not implemented here.
func groupNearDuplicates(entries []imgEntry, threshold int) []*dupGroup {
	var groups []*dupGroup
	for _, e := range entries {
		var joined *dupGroup
		for _, g := range groups {
			for _, m := range g.members {
				if hammingDistance(e.hash, m.hash) <= threshold {
					joined = g
					break
				}
			}
			if joined != nil {
				break
			}
		}
		if joined != nil {
			joined.members = append(joined.members, e)
		} else {
			groups = append(groups, &dupGroup{members: []imgEntry{e}})
		}
	}
	return groups
}

func reportableGroups(groups []*dupGroup) []*dupGroup {
	var out []*dupGroup
	for _, g := range groups {
		if len(g.members) >= 2 {
			out = append(out, g)
		}
	}
	return out
}

// ---------------------------------------------------------------------------
// scan command
// ---------------------------------------------------------------------------

type jsonGroupMember struct {
	Path     string `json:"path"`
	Distance int    `json:"hamming_distance_from_reference"`
	Width    int    `json:"width"`
	Height   int    `json:"height"`
	Size     int64  `json:"size_bytes"`
}

type jsonGroup struct {
	Reference string            `json:"reference"`
	Members   []jsonGroupMember `json:"members"`
}

type jsonScanResult struct {
	TotalScanned    int         `json:"total_scanned"`
	TotalCandidates int         `json:"total_candidate_files"`
	Threshold       int         `json:"threshold"`
	GroupsFound     int         `json:"groups_found"`
	TotalForCleanup int         `json:"total_candidates_for_cleanup"`
	Groups          []jsonGroup `json:"groups"`
	Warnings        []string    `json:"warnings,omitempty"`
}

func cmdScan(args []string) {
	fs := flag.NewFlagSet("scan", flag.ExitOnError)
	threshold := fs.Int("threshold", 5, "max Hamming distance to count as near-duplicate (0-64)")
	minSizeStr := fs.String("min-size", "", "skip files smaller than this, e.g. 10KB, 2MB")
	jsonOut := fs.Bool("json", false, "emit JSON instead of text")

	valueFlags := map[string]bool{"threshold": true, "min-size": true}
	if err := fs.Parse(reorderFlags(args, valueFlags)); err != nil {
		os.Exit(1)
	}
	dirs := fs.Args()
	if len(dirs) == 0 {
		fmt.Fprintln(os.Stderr, "error: scan requires at least one <dir>")
		usage()
		os.Exit(1)
	}
	minSize, err := parseSize(*minSizeStr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}

	entries, totalCandidates, warnings := walkImages(dirs, minSize)
	for _, w := range warnings {
		fmt.Fprintf(os.Stderr, "warning: %s\n", w)
	}

	groups := groupNearDuplicates(entries, *threshold)
	reportable := reportableGroups(groups)

	totalForCleanup := 0
	for _, g := range reportable {
		totalForCleanup += len(g.members) - 1
	}

	if *jsonOut {
		result := jsonScanResult{
			TotalScanned:    len(entries),
			TotalCandidates: totalCandidates,
			Threshold:       *threshold,
			GroupsFound:     len(reportable),
			TotalForCleanup: totalForCleanup,
			Warnings:        warnings,
		}
		for _, g := range reportable {
			ref := g.members[0]
			jg := jsonGroup{Reference: ref.path}
			for _, m := range g.members {
				jg.Members = append(jg.Members, jsonGroupMember{
					Path:     m.path,
					Distance: hammingDistance(m.hash, ref.hash),
					Width:    m.width,
					Height:   m.height,
					Size:     m.size,
				})
			}
			result.Groups = append(result.Groups, jg)
		}
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		enc.Encode(result)
		return
	}

	for i, g := range reportable {
		ref := g.members[0]
		fmt.Printf("Group %d (%d files, reference: %s)\n", i+1, len(g.members), ref.path)
		for _, m := range g.members {
			d := hammingDistance(m.hash, ref.hash)
			fmt.Printf("  [dist %2d] %s  (%dx%d, %s)\n", d, m.path, m.width, m.height, humanBytes(m.size))
		}
		fmt.Println()
	}

	fmt.Println("--- Summary ---")
	fmt.Printf("Files scanned (decoded as images): %d\n", len(entries))
	fmt.Printf("Candidate files considered (by extension): %d\n", totalCandidates)
	fmt.Printf("Threshold: %d\n", *threshold)
	fmt.Printf("Near-duplicate groups found: %d\n", len(reportable))
	fmt.Printf("Total candidates for cleanup (group size - 1, summed): %d\n", totalForCleanup)
	if len(warnings) > 0 {
		fmt.Printf("Warnings: %d (see stderr)\n", len(warnings))
	}
}

// ---------------------------------------------------------------------------
// clean command
// ---------------------------------------------------------------------------

func cmdClean(args []string) {
	fs := flag.NewFlagSet("clean", flag.ExitOnError)
	threshold := fs.Int("threshold", 5, "max Hamming distance to count as near-duplicate (0-64)")
	minSizeStr := fs.String("min-size", "", "skip files smaller than this, e.g. 10KB, 2MB")
	quarantine := fs.String("quarantine", "", "destination directory for quarantined files (required)")
	apply := fs.Bool("apply", false, "actually move files (default: dry run)")

	valueFlags := map[string]bool{"threshold": true, "min-size": true, "quarantine": true}
	if err := fs.Parse(reorderFlags(args, valueFlags)); err != nil {
		os.Exit(1)
	}
	dirs := fs.Args()
	if len(dirs) == 0 {
		fmt.Fprintln(os.Stderr, "error: clean requires at least one <dir>")
		usage()
		os.Exit(1)
	}
	if *quarantine == "" {
		fmt.Fprintln(os.Stderr, "error: clean requires --quarantine <dir>")
		usage()
		os.Exit(1)
	}
	minSize, err := parseSize(*minSizeStr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
	absQuarantine, err := filepath.Abs(*quarantine)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}

	entries, totalCandidates, warnings := walkImages(dirs, minSize)
	for _, w := range warnings {
		fmt.Fprintf(os.Stderr, "warning: %s\n", w)
	}

	groups := groupNearDuplicates(entries, *threshold)
	reportable := reportableGroups(groups)

	if len(reportable) == 0 {
		fmt.Println("No near-duplicate groups found. Nothing to do.")
		fmt.Printf("Files scanned (decoded as images): %d\n", len(entries))
		fmt.Printf("Candidate files considered (by extension): %d\n", totalCandidates)
		return
	}

	mode := "DRY RUN (no files will be moved - pass --apply to execute)"
	if *apply {
		mode = "APPLYING (files WILL be moved)"
	}
	fmt.Printf("PhotoSweep clean - %s\n\n", mode)

	var movedCount int
	var reclaimedBytes int64
	var moveErrors int

	for i, g := range reportable {
		ref := g.members[0]
		fmt.Printf("Group %d (%d files, reference: %s)\n", i+1, len(g.members), ref.path)
		for mi, m := range g.members {
			d := hammingDistance(m.hash, ref.hash)
			if mi == 0 {
				fmt.Printf("  [dist %2d] KEEP        %s  (%dx%d, %s)\n", d, m.path, m.width, m.height, humanBytes(m.size))
				continue
			}
			if !*apply {
				fmt.Printf("  [dist %2d] QUARANTINE  %s  (%dx%d, %s)\n", d, m.path, m.width, m.height, humanBytes(m.size))
				continue
			}
			dest, err := quarantineDest(m, absQuarantine)
			if err != nil {
				fmt.Printf("  [dist %2d] ERROR       %s  (%v)\n", d, m.path, err)
				moveErrors++
				continue
			}
			if err := moveFile(m.path, dest); err != nil {
				fmt.Printf("  [dist %2d] ERROR       %s -> %s (%v)\n", d, m.path, dest, err)
				moveErrors++
				continue
			}
			fmt.Printf("  [dist %2d] QUARANTINED %s -> %s  (%dx%d, %s)\n", d, m.path, dest, m.width, m.height, humanBytes(m.size))
			movedCount++
			reclaimedBytes += m.size
		}
		fmt.Println()
	}

	totalForCleanup := 0
	for _, g := range reportable {
		totalForCleanup += len(g.members) - 1
	}

	fmt.Println("--- Summary ---")
	fmt.Printf("Files scanned (decoded as images): %d\n", len(entries))
	fmt.Printf("Near-duplicate groups found: %d\n", len(reportable))
	fmt.Printf("Total cleanup candidates: %d\n", totalForCleanup)
	if *apply {
		fmt.Printf("Files moved to quarantine: %d\n", movedCount)
		fmt.Printf("Bytes reclaimed from originals' locations: %s\n", humanBytes(reclaimedBytes))
		if moveErrors > 0 {
			fmt.Printf("Errors while moving: %d\n", moveErrors)
		}
	} else {
		fmt.Println("(dry run: no files were moved; re-run with --apply to execute this plan)")
	}
}

// quarantineDest computes the destination path for a quarantined file,
// preserving its path relative to the scanned root directory it was found
// under.
func quarantineDest(e imgEntry, quarantineDir string) (string, error) {
	rel, err := filepath.Rel(e.root, e.path)
	if err != nil {
		return "", err
	}
	return filepath.Join(quarantineDir, rel), nil
}

// moveFile moves src to dest, creating parent directories as needed. It
// falls back to copy+remove if the destination is on a different filesystem
// (os.Rename returns an EXDEV-style error in that case).
func moveFile(src, dest string) error {
	if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
		return err
	}
	if err := os.Rename(src, dest); err == nil {
		return nil
	}
	// Fallback: copy then remove (handles cross-device moves).
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		return err
	}
	if err := out.Close(); err != nil {
		return err
	}
	return os.Remove(src)
}

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

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 "scan":
		cmdScan(os.Args[2:])
	case "clean":
		cmdClean(os.Args[2:])
	default:
		fmt.Fprintf(os.Stderr, "error: unknown command %q\n", os.Args[1])
		usage()
		os.Exit(1)
	}
}