// Command photorevive is a CLI-native, scriptable batch corrupt-media
// scanner and narrow best-effort JPEG repairer for large offline photo
// libraries that span many directories.
//
// See README.txt for full documentation and ../plan.md for the product
// roadmap.
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io/fs"
	"os"
	"path/filepath"
	"strings"
)

// ---------------------------------------------------------------------
// Shared helpers (established conventions across the PhotoRevive suite)
// ---------------------------------------------------------------------

// reorderFlags works around a quirk of Go's flag package: it stops
// parsing flags at the first positional argument. This reorders args so
// all flags (and, for flags that take a value, their following value)
// come before all positional arguments, regardless of where the user
// typed them on the command line.
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...)
}

// humanBytes renders a byte count as a human-readable size (KiB, MiB, ...).
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])
}

// ---------------------------------------------------------------------
// Image format detection
// ---------------------------------------------------------------------

// candidateExts are the file extensions PhotoRevive inspects.
var candidateExts = map[string]string{
	".jpg":  "jpeg",
	".jpeg": "jpeg",
	".png":  "png",
	".gif":  "gif",
	".bmp":  "bmp",
}

// jpegEOI is the two-byte End-Of-Image marker every valid JPEG file must
// end with.
var jpegEOI = []byte{0xFF, 0xD9}

// headerMatches reports whether the first bytes of a file match the magic
// number expected for the given logical format ("jpeg", "png", "gif",
// "bmp"). It only checks the specific magic for that format, not whether
// the bytes look like some other known format.
func headerMatches(format string, header []byte) bool {
	switch format {
	case "jpeg":
		return len(header) >= 3 && header[0] == 0xFF && header[1] == 0xD8 && header[2] == 0xFF
	case "png":
		return len(header) >= 4 && header[0] == 0x89 && header[1] == 0x50 && header[2] == 0x4E && header[3] == 0x47
	case "gif":
		return len(header) >= 4 && header[0] == 0x47 && header[1] == 0x49 && header[2] == 0x46 && header[3] == 0x38
	case "bmp":
		return len(header) >= 2 && header[0] == 0x42 && header[1] == 0x4D
	}
	return false
}

// FileResult is the classification of a single candidate file.
type FileResult struct {
	Path   string `json:"path"`
	Ext    string `json:"ext"`
	Format string `json:"expected_format"`
	Status string `json:"status"` // OK, CORRUPT, UNKNOWN
	Width  int    `json:"width,omitempty"`
	Height int    `json:"height,omitempty"`
	Size   int64  `json:"size_bytes"`
	Detail string `json:"detail,omitempty"`
}

// classify opens path and determines whether it is OK, CORRUPT, or
// UNKNOWN per PhotoRevive's detection rules: try image.Decode first; on
// failure, fall back to a magic-byte check against the format implied by
// the file extension.
func classify(path string) (FileResult, error) {
	ext := strings.ToLower(filepath.Ext(path))
	expected := candidateExts[ext]

	fi, err := os.Stat(path)
	if err != nil {
		return FileResult{}, err
	}

	f, err := os.Open(path)
	if err != nil {
		return FileResult{}, err
	}
	defer f.Close()

	header := make([]byte, 8)
	n, _ := f.Read(header)
	header = header[:n]

	if _, err := f.Seek(0, 0); err != nil {
		return FileResult{}, err
	}

	res := FileResult{
		Path:   path,
		Ext:    ext,
		Format: expected,
		Size:   fi.Size(),
	}

	img, _, decErr := image.Decode(f)
	if decErr == nil {
		res.Status = "OK"
		b := img.Bounds()
		res.Width = b.Dx()
		res.Height = b.Dy()
		return res, nil
	}

	if headerMatches(expected, header) {
		res.Status = "CORRUPT"
		res.Detail = fmt.Sprintf("%s header present but decode failed: %v", expected, decErr)
	} else {
		res.Status = "UNKNOWN"
		res.Detail = fmt.Sprintf("does not match expected %s header (decode error: %v)", expected, decErr)
	}
	return res, nil
}

// walkCandidates walks root and calls fn for every regular file whose
// extension is one PhotoRevive recognizes.
func walkCandidates(root string, fn func(path string) error) error {
	return filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			// Skip unreadable entries but keep walking the rest of the tree.
			fmt.Fprintf(os.Stderr, "warning: %v\n", err)
			return nil
		}
		if d.IsDir() {
			return nil
		}
		ext := strings.ToLower(filepath.Ext(path))
		if _, ok := candidateExts[ext]; !ok {
			return nil
		}
		return fn(path)
	})
}

// ---------------------------------------------------------------------
// scan
// ---------------------------------------------------------------------

type dirScan struct {
	Path    string       `json:"path"`
	Error   string       `json:"error,omitempty"`
	OK      int          `json:"ok"`
	Corrupt int          `json:"corrupt"`
	Unknown int          `json:"unknown"`
	Files   []FileResult `json:"files"`
}

type scanReport struct {
	Directories  []dirScan `json:"directories"`
	TotalOK      int       `json:"total_ok"`
	TotalCorrupt int       `json:"total_corrupt"`
	TotalUnknown int       `json:"total_unknown"`
	TotalFiles   int       `json:"total_files"`
}

func runScan(args []string) {
	fs := flag.NewFlagSet("scan", flag.ExitOnError)
	jsonOut := fs.Bool("json", false, "emit machine-readable JSON instead of a text report")
	fs.Usage = func() { fmt.Fprint(os.Stderr, scanUsage) }
	if err := fs.Parse(reorderFlags(args, map[string]bool{})); err != nil {
		os.Exit(1)
	}
	dirs := fs.Args()
	if len(dirs) == 0 {
		fmt.Fprint(os.Stderr, scanUsage)
		os.Exit(1)
	}

	report := scanReport{}
	for _, dir := range dirs {
		ds := dirScan{Path: dir}
		if info, err := os.Stat(dir); err != nil {
			ds.Error = err.Error()
			report.Directories = append(report.Directories, ds)
			continue
		} else if !info.IsDir() {
			ds.Error = fmt.Sprintf("%s is not a directory", dir)
			report.Directories = append(report.Directories, ds)
			continue
		}

		walkErr := walkCandidates(dir, func(path string) error {
			res, err := classify(path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "warning: %s: %v\n", path, err)
				return nil
			}
			ds.Files = append(ds.Files, res)
			switch res.Status {
			case "OK":
				ds.OK++
			case "CORRUPT":
				ds.Corrupt++
			case "UNKNOWN":
				ds.Unknown++
			}
			return nil
		})
		if walkErr != nil {
			ds.Error = walkErr.Error()
		}
		report.Directories = append(report.Directories, ds)
		report.TotalOK += ds.OK
		report.TotalCorrupt += ds.Corrupt
		report.TotalUnknown += ds.Unknown
		report.TotalFiles += ds.OK + ds.Corrupt + ds.Unknown
	}

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

	printScanReport(report)
}

func printScanReport(report scanReport) {
	fmt.Println("PhotoRevive - Batch Corrupt Media Scanner")
	fmt.Println(strings.Repeat("=", 42))

	for _, ds := range report.Directories {
		fmt.Printf("\nDirectory: %s\n", ds.Path)
		if ds.Error != "" {
			fmt.Printf("  ERROR: %s\n", ds.Error)
			continue
		}
		fmt.Printf("  OK:      %d\n", ds.OK)
		fmt.Printf("  CORRUPT: %d\n", ds.Corrupt)
		fmt.Printf("  UNKNOWN: %d\n", ds.Unknown)

		if ds.Corrupt > 0 || ds.Unknown > 0 {
			fmt.Println("  --- CORRUPT / UNKNOWN files ---")
			for _, f := range ds.Files {
				if f.Status == "OK" {
					continue
				}
				fmt.Printf("  %-8s %s  (%s)  %s\n", f.Status, f.Path, humanBytes(f.Size), f.Detail)
			}
		}
	}

	fmt.Println()
	fmt.Println(strings.Repeat("=", 42))
	fmt.Printf("GRAND TOTAL across %d director(y/ies):\n", len(report.Directories))
	fmt.Printf("  OK:      %d\n", report.TotalOK)
	fmt.Printf("  CORRUPT: %d\n", report.TotalCorrupt)
	fmt.Printf("  UNKNOWN: %d\n", report.TotalUnknown)
	fmt.Printf("  TOTAL:   %d\n", report.TotalFiles)
}

// ---------------------------------------------------------------------
// repair
// ---------------------------------------------------------------------

// attemptEOIRepair tries the narrow, best-effort repair: append the
// missing FF D9 End-Of-Image marker to a copy of the file's bytes, held
// entirely in memory, and verify the result actually decodes before
// anything is written to disk. It never touches the original file.
//
// Returns (repairedBytes, ok, reason).
func attemptEOIRepair(path string) ([]byte, bool, string) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, false, fmt.Sprintf("could not read source: %v", err)
	}
	if len(data) >= 2 && bytes.Equal(data[len(data)-2:], jpegEOI) {
		return nil, false, "file already ends with FF D9 (EOI present); corruption is not a missing-EOI truncation"
	}

	repaired := make([]byte, len(data)+2)
	copy(repaired, data)
	repaired[len(data)] = jpegEOI[0]
	repaired[len(data)+1] = jpegEOI[1]

	if _, _, err := image.Decode(bytes.NewReader(repaired)); err != nil {
		return nil, false, fmt.Sprintf("appending FF D9 was not enough to decode: %v", err)
	}
	return repaired, true, "missing EOI marker appended; image now decodes"
}

type repairOutcome struct {
	Path      string `json:"path"`
	SourceDir string `json:"source_dir"`
	Status    string `json:"status"` // REPAIRED, REPAIR-FAILED, SKIP, WOULD-REPAIR, WOULD-FAIL
	OutPath   string `json:"out_path,omitempty"`
	Width     int    `json:"width,omitempty"`
	Height    int    `json:"height,omitempty"`
	Detail    string `json:"detail,omitempty"`
}

func runRepair(args []string) {
	fs := flag.NewFlagSet("repair", flag.ExitOnError)
	outDir := fs.String("out-dir", "", "directory to write repaired files into (required)")
	apply := fs.Bool("apply", false, "actually write repaired files (default is a dry run)")
	fs.Usage = func() { fmt.Fprint(os.Stderr, repairUsage) }
	if err := fs.Parse(reorderFlags(args, map[string]bool{"out-dir": true})); err != nil {
		os.Exit(1)
	}
	dirs := fs.Args()
	if len(dirs) == 0 || *outDir == "" {
		fmt.Fprint(os.Stderr, repairUsage)
		os.Exit(1)
	}

	mode := "DRY RUN"
	if *apply {
		mode = "APPLY"
	}
	fmt.Println("PhotoRevive - JPEG Missing-EOI Repair")
	fmt.Println(strings.Repeat("=", 42))
	fmt.Printf("Mode: %s\n", mode)
	if !*apply {
		fmt.Println("(no files will be written; pass --apply to write repaired output)")
	}

	var repaired, failed, skipped, examined int

	for _, dir := range dirs {
		info, err := os.Stat(dir)
		if err != nil {
			fmt.Printf("\nDirectory: %s\n  ERROR: %v\n", dir, err)
			continue
		}
		if !info.IsDir() {
			fmt.Printf("\nDirectory: %s\n  ERROR: not a directory\n", dir)
			continue
		}
		fmt.Printf("\nDirectory: %s\n", dir)

		var outcomes []repairOutcome
		walkErr := walkCandidates(dir, func(path string) error {
			res, err := classify(path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "warning: %s: %v\n", path, err)
				return nil
			}
			if res.Status != "CORRUPT" {
				return nil
			}
			examined++

			outcome := repairOutcome{Path: path, SourceDir: dir}

			if res.Format != "jpeg" {
				outcome.Status = "SKIP"
				outcome.Detail = "SKIP (not a JPEG-EOI-repairable case): not a JPEG file"
				skipped++
				outcomes = append(outcomes, outcome)
				return nil
			}

			repairedBytes, ok, reason := attemptEOIRepair(path)
			if !ok {
				outcome.Detail = reason
				if strings.Contains(reason, "already ends with FF D9") {
					outcome.Status = "SKIP"
					outcome.Detail = "SKIP (not a JPEG-EOI-repairable case): " + reason
					skipped++
				} else if *apply {
					outcome.Status = "REPAIR-FAILED"
					failed++
				} else {
					outcome.Status = "WOULD-FAIL"
					failed++
				}
				outcomes = append(outcomes, outcome)
				return nil
			}

			rel, relErr := filepath.Rel(dir, path)
			if relErr != nil {
				rel = filepath.Base(path)
			}
			outPath := filepath.Join(*outDir, filepath.Base(filepath.Clean(dir)), rel)

			if *apply {
				if err := os.MkdirAll(filepath.Dir(outPath), 0o755); err != nil {
					outcome.Status = "REPAIR-FAILED"
					outcome.Detail = fmt.Sprintf("could not create output dir: %v", err)
					failed++
					outcomes = append(outcomes, outcome)
					return nil
				}
				if err := os.WriteFile(outPath, repairedBytes, 0o644); err != nil {
					outcome.Status = "REPAIR-FAILED"
					outcome.Detail = fmt.Sprintf("could not write repaired file: %v", err)
					failed++
					outcomes = append(outcomes, outcome)
					return nil
				}
				outcome.Status = "REPAIRED"
				outcome.OutPath = outPath
				repaired++
			} else {
				outcome.Status = "WOULD-REPAIR"
				outcome.OutPath = outPath
				repaired++
			}
			if img, _, err := image.Decode(bytes.NewReader(repairedBytes)); err == nil {
				b := img.Bounds()
				outcome.Width, outcome.Height = b.Dx(), b.Dy()
			}
			outcome.Detail = reason
			outcomes = append(outcomes, outcome)
			return nil
		})
		if walkErr != nil {
			fmt.Printf("  ERROR: %v\n", walkErr)
		}

		if len(outcomes) == 0 {
			fmt.Println("  (no CORRUPT files found in this directory)")
		}
		for _, o := range outcomes {
			if o.OutPath != "" {
				dims := ""
				if o.Width > 0 {
					dims = fmt.Sprintf(" [%dx%d]", o.Width, o.Height)
				}
				fmt.Printf("  %-14s %s -> %s%s\n", o.Status, o.Path, o.OutPath, dims)
			} else {
				fmt.Printf("  %-14s %s\n", o.Status, o.Path)
			}
			if o.Detail != "" {
				fmt.Printf("                 %s\n", o.Detail)
			}
		}
	}

	fmt.Println()
	fmt.Println(strings.Repeat("=", 42))
	verb := "REPAIRED"
	failVerb := "REPAIR-FAILED"
	if !*apply {
		verb = "WOULD-REPAIR"
		failVerb = "WOULD-FAIL"
	}
	fmt.Printf("SUMMARY (%s):\n", mode)
	fmt.Printf("  CORRUPT files examined: %d\n", examined)
	fmt.Printf("  %s: %d\n", verb, repaired)
	fmt.Printf("  %s: %d\n", failVerb, failed)
	fmt.Printf("  SKIP: %d\n", skipped)
}

// ---------------------------------------------------------------------
// usage / main
// ---------------------------------------------------------------------

const topUsage = `PhotoRevive - batch corrupt-media scanner and narrow JPEG repair tool

Usage:
  photorevive scan <dir> [<dir> ...] [--json]
  photorevive repair <dir> [<dir> ...] --out-dir <dir> [--apply]
  photorevive help

Run "photorevive <command> -h" for details on a specific command.
`

const scanUsage = `Usage: photorevive scan <dir> [<dir> ...] [--json]

Recursively scans one or more directories for .jpg/.jpeg/.png/.gif/.bmp
files, classifying each as OK, CORRUPT, or UNKNOWN, and prints a
consolidated per-directory and grand-total report. Read-only: no files
are modified.

Flags:
  --json    emit a machine-readable JSON report instead of text
`

const repairUsage = `Usage: photorevive repair <dir> [<dir> ...] --out-dir <dir> [--apply]

Scans one or more directories like "scan" does, then for every file
classified CORRUPT and identified as a JPEG missing its trailing FF D9
End-Of-Image marker, attempts a narrow best-effort repair: append FF D9
to an in-memory copy and verify it decodes before writing anything.

Without --apply this is a dry run: the same repair-and-verify logic runs
in memory and an accurate WOULD-REPAIR / WOULD-FAIL preview is printed,
but --out-dir is never created or written to.

With --apply, successfully repaired files are written under --out-dir,
preserving each file's path relative to its source directory.

Flags:
  --out-dir <dir>   directory to write repaired files into (required)
  --apply           actually write output (default is dry run)
`

func usage() {
	fmt.Fprint(os.Stderr, topUsage)
}

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

	cmd := os.Args[1]
	switch cmd {
	case "-h", "--help", "help":
		fmt.Fprint(os.Stdout, topUsage)
		return
	case "scan":
		runScan(os.Args[2:])
	case "repair":
		runRepair(os.Args[2:])
	default:
		fmt.Fprintf(os.Stderr, "photorevive: unknown command %q\n\n", cmd)
		usage()
		os.Exit(1)
	}
}