// main.go
package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/glaslos/ssdeep"
)

const blockSize = 1024 * 1024

// main is the entry point of the program.
func main() {
	// 1. Get and validate command-line arguments
	if len(os.Args) != 3 {
		usage := `
Binwally (Go Version): Binary and Directory tree comparison tool
                       using the Fuzzy Hashing concept (ssdeep)

Usage: go run . <dir1/file1> <dir2/file2>
`
		fmt.Printf(usage)
		os.Exit(1)
	}
	path1 := os.Args[1]
	path2 := os.Args[2]

	info1, err1 := os.Stat(path1)
	if err1 != nil {
		log.Fatalf("Invalid path: %s", path1)
	}
	info2, err2 := os.Stat(path2)
	if err2 != nil {
		log.Fatalf("Invalid path: %s", path2)
	}

	// 2. Route to appropriate comparison logic
	if info1.IsDir() && info2.IsDir() {
		// Both are directories
		var diffs []int
		fmt.Println("\nSCORE RESULT  PATH")
		compareTrees(path1, path2, &diffs)

		if len(diffs) == 0 {
			fmt.Println("No diffs found\n")
		} else {
			totalScore := 0
			for _, score := range diffs {
				totalScore += score
			}
			averageScore := totalScore / len(diffs)
			fmt.Println("\nTotal files compared:", len(diffs))
			fmt.Printf("Overall match score:  %d%%\n\n", averageScore)
		}
	} else if !info1.IsDir() && !info2.IsDir() {
		// Both are files
		hash1, err := ssdeep.HashFromFile(path1) // CORRECTED
		if err != nil {
			log.Fatalf("Failed to hash file %s: %v", path1, err)
		}
		hash2, err := ssdeep.HashFromFile(path2) // CORRECTED
		if err != nil {
			log.Fatalf("Failed to hash file %s: %v", path2, err)
		}
		score, err := ssdeep.Distance(hash1, hash2) // CORRECTED
		if err != nil {
			log.Fatalf("Failed to compare hashes: %v", err)
		}
		fmt.Printf("Overall match score:  %d%%\n\n", score)
	} else {
		log.Fatal("Error: Both arguments must be files or both must be directories.")
	}
}

// compareTrees recursively compares two directory trees.
func compareTrees(dir1, dir2 string, diffs *[]int) {
	// 1. Read contents of both directories
	entries1, err := os.ReadDir(dir1)
	if err != nil {
		log.Printf("Error reading directory %s: %v", dir1, err)
		return
	}
	entries2, err := os.ReadDir(dir2)
	if err != nil {
		log.Printf("Error reading directory %s: %v", dir2, err)
		return
	}

	// Create maps for quick lookup by name
	map1 := make(map[string]os.DirEntry)
	for _, e := range entries1 {
		map1[e.Name()] = e
	}
	map2 := make(map[string]os.DirEntry)
	for _, e := range entries2 {
		map2[e.Name()] = e
	}

	// 2. Find and report unique files/dirs
	reportUniques(map1, map2, dir1, " <<< unique ", diffs) // In dir1 only
	reportUniques(map2, map1, dir2, " >>> unique ", diffs) // In dir2 only

	// 3. Process common items
	for name, e1 := range map1 {
		if e2, ok := map2[name]; ok {
			// Name exists in both directories
			path1 := filepath.Join(dir1, name)
			path2 := filepath.Join(dir2, name)

			isDir1 := e1.IsDir()
			isDir2 := e2.IsDir()

			if isDir1 && isDir2 {
				// Both are directories, recurse
				compareTrees(path1, path2, diffs)
			} else if !isDir1 && !isDir2 {
				// Both are files, compare them
				compareFiles(path1, path2, diffs)
			} else {
				// Mismatched types (file vs dir) or symlinks
				fmt.Printf("    - ignored %s (symlink or type mismatch)\n", path1)
				*diffs = append(*diffs, 100) // Original script counts this as 100
			}
		}
	}
}

// reportUniques finds items in map1 not in map2 and reports them.
func reportUniques(map1, map2 map[string]os.DirEntry, baseDir, prefix string, diffs *[]int) {
	for name, entry := range map1 {
		if _, exists := map2[name]; !exists {
			fullPath := filepath.Join(baseDir, name)
			if !entry.IsDir() {
				// It's a unique file
				fmt.Println(prefix, fullPath)
				*diffs = append(*diffs, 0)
			} else {
				// It's a unique directory, walk it and report all files within
				filepath.WalkDir(fullPath, func(path string, d os.DirEntry, err error) error {
					if err != nil {
						return err
					}
					if !d.IsDir() {
						fmt.Println(prefix, path)
						*diffs = append(*diffs, 0)
					}
					return nil
				})
			}
		}
	}
}

// compareFiles compares two files, first byte-by-byte, then with ssdeep if different.
func compareFiles(path1, path2 string, diffs *[]int) {
	f1, err := os.Open(path1)
	if err != nil {
		log.Printf("Error opening file %s: %v", path1, err)
		return
	}
	defer f1.Close()

	f2, err := os.Open(path2)
	if err != nil {
		log.Printf("Error opening file %s: %v", path2, err)
		return
	}
	defer f2.Close()

	// Compare files chunk by chunk
	for {
		b1 := make([]byte, blockSize)
		_, err1 := f1.Read(b1)

		b2 := make([]byte, blockSize)
		_, err2 := f2.Read(b2)

		if err1 == io.EOF && err2 == io.EOF {
			// Files are identical
			fmt.Printf("%5s matches %s\n", "100", getRelativePath(path1))
			*diffs = append(*diffs, 100)
			return
		}

		if err1 != nil && err1 != io.EOF || err2 != nil && err2 != io.EOF {
			log.Printf("Error reading from files: %v, %v", err1, err2)
			return
		}

		if !bytes.Equal(b1, b2) {
			// Files differ, use ssdeep
			hash1, _ := ssdeep.HashFromFile(path1)    // CORRECTED
			hash2, _ := ssdeep.HashFromFile(path2)    // CORRECTED
			score, _ := ssdeep.Distance(hash1, hash2) // CORRECTED
			fmt.Printf("%5s differs %s\n", strconv.Itoa(score), getRelativePath(path1))
			*diffs = append(*diffs, score)
			return
		}
	}
}

// getRelativePath tries to mimic the Python script's path trimming.
func getRelativePath(path string) string {
	parts := strings.Split(filepath.ToSlash(path), "/")
	if len(parts) > 1 {
		return strings.Join(parts[1:], "/")
	}
	return path
}