package main

import (
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"lukechampine.com/blake3"
)

func main() {
	// Allow dist directory to be overridden by environment variable or command line arg
	distDir := os.Getenv("DIST_DIR")
	if distDir == "" {
		if len(os.Args) > 1 {
			distDir = os.Args[1]
		} else {
			distDir = "dist"
		}
	}
	outputFile := filepath.Join(distDir, "BLAKE3SUMS.txt")

	if _, err := os.Stat(distDir); os.IsNotExist(err) {
		fmt.Fprintf(os.Stderr, "Error: dist directory does not exist\n")
		os.Exit(1)
	}

	files, err := filepath.Glob(filepath.Join(distDir, "*.tar.gz"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Failed to find tar.gz files: %v\n", err)
		os.Exit(1)
	}

	zipFiles, err := filepath.Glob(filepath.Join(distDir, "*.zip"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Failed to find zip files: %v\n", err)
		os.Exit(1)
	}

	files = append(files, zipFiles...)

	if len(files) == 0 {
		fmt.Fprintf(os.Stderr, "Error: No release files found in dist/\n")
		os.Exit(1)
	}

	var checksums []string

	for _, file := range files {
		hash, err := computeBLAKE3(file)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: Failed to compute hash for %s: %v\n", filepath.Base(file), err)
			continue
		}
		filename := filepath.Base(file)
		checksums = append(checksums, fmt.Sprintf("%s  %s", hash, filename))
		fmt.Printf("Computed: %s -> %s\n", hash, filename)
	}

	if len(checksums) == 0 {
		fmt.Fprintf(os.Stderr, "Error: No checksums computed\n")
		os.Exit(1)
	}

	content := ""
	for _, checksum := range checksums {
		content += checksum + "\n"
	}

	if err := os.WriteFile(outputFile, []byte(content), 0644); err != nil {
		fmt.Fprintf(os.Stderr, "Error: Failed to write BLAKE3SUMS.txt: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("\nBLAKE3 checksums written to %s\n", outputFile)
}

func computeBLAKE3(filePath string) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	hash := blake3.New(32, nil)
	if _, err := io.Copy(hash, file); err != nil {
		return "", err
	}

	return "blake3:" + hex.EncodeToString(hash.Sum(nil)), nil
}