Add configuration struct for command-line flags and refactor main logic to use it
This commit is contained in:
54
main.go
54
main.go
@@ -15,6 +15,13 @@ import (
|
||||
"github.com/corona10/goimagehash"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Threshold int
|
||||
Workers int
|
||||
NotifyInterval int
|
||||
NoGroupByFolder bool
|
||||
}
|
||||
|
||||
type HashResult struct {
|
||||
file string
|
||||
hash *goimagehash.ExtImageHash
|
||||
@@ -26,14 +33,20 @@ type GroupResult struct {
|
||||
}
|
||||
|
||||
func main() {
|
||||
thresh := flag.Int("thresh", 10, "Threshold for distance")
|
||||
workers := flag.Int("workers", 100, "Number of workers")
|
||||
notifyInterval := flag.Int("notify", 1000, "Notify interval")
|
||||
cfg := Config{
|
||||
Threshold: *flag.Int("thresh", 10, "Threshold for distance"),
|
||||
Workers: *flag.Int("workers", 100, "Number of workers"),
|
||||
NotifyInterval: *flag.Int("notify", 1000, "Notify interval"),
|
||||
NoGroupByFolder: *flag.Bool("nogroup", false, "Don't group by folder"),
|
||||
}
|
||||
flag.Parse()
|
||||
logger.InitFlag()
|
||||
logger.Info("Starting")
|
||||
logger.Info("Threshold: %v", *thresh)
|
||||
logger.Info("Threshold: %v", cfg.Threshold)
|
||||
logger.Info("Patterns: %d", len(flag.Args()))
|
||||
logger.Info("Workers: %v", cfg.Workers)
|
||||
logger.Info("Notify interval: %v", cfg.NotifyInterval)
|
||||
logger.Info("Group by folder: %v", !cfg.NoGroupByFolder)
|
||||
|
||||
// Collect files
|
||||
files := make([]string, 0)
|
||||
@@ -54,15 +67,38 @@ func main() {
|
||||
}
|
||||
logger.Info("Patterns expanded to %d files", len(files))
|
||||
|
||||
if !cfg.NoGroupByFolder {
|
||||
// Group files by folder
|
||||
folderMap := make(map[string][]string)
|
||||
for _, file := range files {
|
||||
dir := filepath.Dir(file)
|
||||
folderMap[dir] = append(folderMap[dir], file)
|
||||
}
|
||||
logger.Info("Processing %d folders", len(folderMap))
|
||||
|
||||
// Process each folder separately
|
||||
for dir, group := range folderMap {
|
||||
logger.Info("Processing folder: %s (%d files)", dir, len(group))
|
||||
processFiles(group, cfg)
|
||||
}
|
||||
} else {
|
||||
// Process all files together
|
||||
processFiles(files, cfg)
|
||||
}
|
||||
|
||||
logger.Info("Done")
|
||||
}
|
||||
|
||||
func processFiles(files []string, cfg Config) {
|
||||
// Phase 1: Hashing
|
||||
hashChan := make(chan string, len(files))
|
||||
hashResults := make(chan HashResult, len(files))
|
||||
var hashWg sync.WaitGroup
|
||||
|
||||
// Start hash workers
|
||||
for i := 0; i < *workers; i++ {
|
||||
for i := 0; i < cfg.Workers; i++ {
|
||||
hashWg.Add(1)
|
||||
go hashWorker(hashChan, hashResults, &hashWg, *notifyInterval, len(files))
|
||||
go hashWorker(hashChan, hashResults, &hashWg, cfg.NotifyInterval, len(files))
|
||||
}
|
||||
|
||||
// Send files to hash workers
|
||||
@@ -88,9 +124,9 @@ func main() {
|
||||
var groupWg sync.WaitGroup
|
||||
|
||||
// Start group workers
|
||||
for i := 0; i < *workers; i++ {
|
||||
for i := 0; i < cfg.Workers; i++ {
|
||||
groupWg.Add(1)
|
||||
go groupWorker(groupChan, groupResults, &groupWg, *thresh, *notifyInterval, len(allHashes), allHashes)
|
||||
go groupWorker(groupChan, groupResults, &groupWg, cfg.Threshold, cfg.NotifyInterval, len(allHashes), allHashes)
|
||||
}
|
||||
|
||||
// Send hash results to group workers
|
||||
@@ -138,8 +174,6 @@ func main() {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.Info("Done")
|
||||
}
|
||||
|
||||
func hashWorker(in <-chan string, out chan<- HashResult, wg *sync.WaitGroup, notifyInterval int, totalFiles int) {
|
||||
|
Reference in New Issue
Block a user