minerva-pick/main.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/anacrolix/torrent"
)

const usageText = `minerva-pick - pobieranie z magneta
		Uzycie:
		minerva-pick [opcje] ['magnet:?xt=urn:btih:...']

		Przyklady:
		minerva-pick 'magnet:?xt=urn:btih:...'
		minerva-pick
		minerva-pick -out ./downloads
		minerva-pick -list
		minerva-pick -forget-magnet

		Podany magnet jest zapisywany jako ostatnio uzyty.
		Gdy uruchomisz program bez magneta, uzyje zapamietanego.

		Jesli fzf jest zainstalowany, zostanie uzyty automatycznie.
		Bez fzf wybierz pliki wpisujac np.: 1,3,7-10 albo all.

		Elo.
`

type options struct {
	outDir          string
	listOnly        bool
	noFZF           bool
	forgetMagnet    bool
	metadataTimeout time.Duration
}

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	if err := run(ctx); err != nil {
		if errors.Is(err, context.Canceled) {
			fmt.Fprintln(os.Stderr, "\nPrzerwano.")
			os.Exit(130)
		}
		fmt.Fprintln(os.Stderr, "blad:", err)
		os.Exit(1)
	}
}

func run(ctx context.Context) error {
	var opt options
	flag.StringVar(&opt.outDir, "out", "./downloads", "katalog docelowy")
	flag.BoolVar(&opt.listOnly, "list", false, "tylko pokaz liste plikow")
	flag.BoolVar(&opt.noFZF, "no-fzf", false, "nie uzywaj fzf; wybor tekstowy")
	flag.BoolVar(&opt.forgetMagnet, "forget-magnet", false, "usun zapamietany magnet i zakoncz")
	flag.DurationVar(&opt.metadataTimeout, "metadata-timeout", 5*time.Minute, "timeout pobierania metadanych")
	flag.Usage = func() {
		fmt.Fprint(flag.CommandLine.Output(), usageText)
		flag.PrintDefaults()
	}
	flag.Parse()

	if flag.NArg() > 1 {
		flag.Usage()
		return fmt.Errorf("podaj najwyzej jeden magnet link")
	}

	magnetPath, err := lastMagnetPath()
	if err != nil {
		return err
	}

	if opt.forgetMagnet {
		if flag.NArg() != 0 {
			return fmt.Errorf("-forget-magnet nie przyjmuje magnet linku")
		}
		if err := forgetLastMagnet(magnetPath); err != nil {
			return err
		}
		fmt.Fprintln(os.Stderr, "Zapamietany magnet usuniety.")
		return nil
	}

	var magnet string
	if flag.NArg() == 1 {
		magnet = flag.Arg(0)
	} else {
		magnet, err = loadLastMagnet(magnetPath)
		if err != nil {
			if errors.Is(err, os.ErrNotExist) {
				flag.Usage()
				return fmt.Errorf("brak zapamietanego magneta; podaj magnet link przy pierwszym uruchomieniu")
			}
			return err
		}
		fmt.Fprintln(os.Stderr, "Uzywam zapamietanego magneta.")
	}

	if !strings.HasPrefix(strings.ToLower(magnet), "magnet:?") {
		return fmt.Errorf("argument nie wyglada jak magnet link")
	}

	outDir, err := filepath.Abs(opt.outDir)
	if err != nil {
		return fmt.Errorf("ustalenie katalogu docelowego %q: %w", opt.outDir, err)
	}
	if err := os.MkdirAll(filepath.Dir(outDir), 0o755); err != nil {
		return fmt.Errorf("utworzenie katalogu nadrzednego dla %q: %w", outDir, err)
	}

	stageDir, err := os.MkdirTemp(filepath.Dir(outDir), ".minerva-pick-")
	if err != nil {
		return fmt.Errorf("utworzenie katalogu roboczego: %w", err)
	}
	defer os.RemoveAll(stageDir)

	cfg := torrent.NewDefaultClientConfig()
	cfg.DisableIPv6 = true
	cfg.DataDir = stageDir
	cfg.Seed = false

	client, err := torrent.NewClient(cfg)
	if err != nil {
		return fmt.Errorf("uruchomienie klienta torrent: %w", err)
	}
	clientClosed := false
	defer func() {
		if !clientClosed {
			client.Close()
		}
	}()

	t, err := client.AddMagnet(magnet)
	if err != nil {
		return fmt.Errorf("dodanie magneta: %w", err)
	}

	fmt.Fprintln(os.Stderr, "Pobieram metadane torrenta...")
	if err := waitForMetadata(ctx, t, opt.metadataTimeout); err != nil {
		return err
	}

	files := t.Files()
	if len(files) == 0 {
		return fmt.Errorf("torrent nie zawiera plikow")
	}

	if flag.NArg() == 1 {
		if err := saveLastMagnet(magnetPath, magnet); err != nil {
			fmt.Fprintln(os.Stderr, "ostrzezenie: nie udalo sie zapamietac magneta:", err)
		}
	}

	info := t.Info()
	if info != nil && info.Name != "" {
		fmt.Fprintf(os.Stderr, "Torrent: %s (%d plikow)\n", info.Name, len(files))
	} else {
		fmt.Fprintf(os.Stderr, "Metadane gotowe: %d plikow\n", len(files))
	}

	if opt.listOnly {
		printFiles(files)
		return nil
	}

	selected, err := chooseFiles(ctx, files, opt.noFZF)
	if err != nil {
		return err
	}
	if len(selected) == 0 {
		fmt.Fprintln(os.Stderr, "Nic nie wybrano.")
		return nil
	}

	var total int64
	fmt.Fprintln(os.Stderr, "Wybrane pliki:")
	for _, idx := range selected {
		f := files[idx]
		total += f.Length()
		fmt.Fprintf(os.Stderr, "  [%d] %8s  %s\n", idx+1, humanBytes(f.Length()), f.Path())
		f.Download()
	}
	fmt.Fprintf(os.Stderr, "Razem: %s\n", humanBytes(total))

	if err := waitForSelected(ctx, files, selected, total); err != nil {
		return err
	}

	if errs := client.Close(); len(errs) > 0 {
		fmt.Fprintf(os.Stderr, "ostrzezenie: zamykanie klienta torrent zwrocilo %d bledow\n", len(errs))
	}
	clientClosed = true

	if err := publishSelected(stageDir, outDir, files, selected); err != nil {
		return err
	}
	fmt.Fprintf(os.Stderr, "Gotowe. Zapisano %d wybranych plikow w %s\n", len(selected), outDir)
	return nil
}

func lastMagnetPath() (string, error) {
	configDir, err := os.UserConfigDir()
	if err != nil {
		return "", fmt.Errorf("ustalenie katalogu konfiguracji: %w", err)
	}
	return filepath.Join(configDir, "minerva-pick", "last-magnet"), nil
}

func saveLastMagnet(path, magnet string) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return fmt.Errorf("utworzenie katalogu konfiguracji: %w", err)
	}
	if err := os.WriteFile(path, []byte(strings.TrimSpace(magnet)+"\n"), 0o600); err != nil {
		return fmt.Errorf("zapis magneta: %w", err)
	}
	return nil
}

func loadLastMagnet(path string) (string, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	magnet := strings.TrimSpace(string(b))
	if magnet == "" {
		return "", fmt.Errorf("zapamietany magnet jest pusty")
	}
	return magnet, nil
}

func forgetLastMagnet(path string) error {
	err := os.Remove(path)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("usuniecie zapamietanego magneta: %w", err)
	}
	return nil
}

func waitForMetadata(ctx context.Context, t *torrent.Torrent, timeout time.Duration) error {
	timer := time.NewTimer(timeout)
	defer timer.Stop()

	select {
	case <-t.GotInfo():
		return nil
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return fmt.Errorf("timeout po %s podczas pobierania metadanych", timeout)
	}
}

func chooseFiles(ctx context.Context, files []*torrent.File, noFZF bool) ([]int, error) {
	if !noFZF {
		if path, err := exec.LookPath("fzf"); err == nil {
			selected, err := chooseWithFZF(ctx, path, files)
			if err == nil {
				return selected, nil
			}
			if errors.Is(err, context.Canceled) {
				return nil, err
			}
			fmt.Fprintln(os.Stderr, "fzf nie zadzialal, przechodze do wyboru tekstowego:", err)
		}
	}
	return chooseFromPrompt(files)
}

func chooseWithFZF(ctx context.Context, fzfPath string, files []*torrent.File) ([]int, error) {
	var input strings.Builder
	for i, f := range files {
		fmt.Fprintf(&input, "%d\t%8s\t%s\n", i+1, humanBytes(f.Length()), f.Path())
	}

	cmd := exec.CommandContext(ctx, fzfPath,
		"--multi",
		"--height=90%",
		"--border",
		"--delimiter=\t",
		"--with-nth=2..",
		"--prompt=pliki> ",
		"--bind=ctrl-a:select-all,ctrl-d:deselect-all,ctrl-t:toggle-all",
		"--header=TAB: zaznacz/odznacz, ENTER: pobierz, ESC: anuluj",
	)
	cmd.Stdin = strings.NewReader(input.String())
	cmd.Stderr = os.Stderr

	var out bytes.Buffer
	cmd.Stdout = &out

	if err := cmd.Run(); err != nil {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) && (exitErr.ExitCode() == 1 || exitErr.ExitCode() == 130) {
			return nil, nil
		}
		return nil, err
	}

	var selected []int
	scanner := bufio.NewScanner(&out)
	for scanner.Scan() {
		line := scanner.Text()
		first, _, ok := strings.Cut(line, "\t")
		if !ok {
			continue
		}
		n, err := strconv.Atoi(strings.TrimSpace(first))
		if err == nil && n >= 1 && n <= len(files) {
			selected = append(selected, n-1)
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return uniqueSorted(selected), nil
}

func chooseFromPrompt(files []*torrent.File) ([]int, error) {
	printFiles(files)
	fmt.Fprint(os.Stderr, "Wybierz pliki (np. 1,3,7-10 albo all): ")

	line, err := bufio.NewReader(os.Stdin).ReadString('\n')
	if err != nil && len(line) == 0 {
		return nil, err
	}
	return parseSelection(strings.TrimSpace(line), len(files))
}

func parseSelection(s string, max int) ([]int, error) {
	if strings.EqualFold(strings.TrimSpace(s), "all") {
		ret := make([]int, max)
		for i := range ret {
			ret[i] = i
		}
		return ret, nil
	}

	if strings.TrimSpace(s) == "" {
		return nil, nil
	}

	var result []int
	for _, part := range strings.Split(s, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}

		if strings.Contains(part, "-") {
			loS, hiS, ok := strings.Cut(part, "-")
			if !ok {
				return nil, fmt.Errorf("niepoprawny zakres %q", part)
			}
			lo, err1 := strconv.Atoi(strings.TrimSpace(loS))
			hi, err2 := strconv.Atoi(strings.TrimSpace(hiS))
			if err1 != nil || err2 != nil || lo < 1 || hi < lo || hi > max {
				return nil, fmt.Errorf("niepoprawny zakres %q (dozwolone 1-%d)", part, max)
			}
			for n := lo; n <= hi; n++ {
				result = append(result, n-1)
			}
			continue
		}

		n, err := strconv.Atoi(part)
		if err != nil || n < 1 || n > max {
			return nil, fmt.Errorf("niepoprawny numer %q (dozwolone 1-%d)", part, max)
		}
		result = append(result, n-1)
	}

	return uniqueSorted(result), nil
}

func uniqueSorted(in []int) []int {
	if len(in) == 0 {
		return nil
	}
	sort.Ints(in)
	out := in[:0]
	last := -1
	for _, n := range in {
		if n != last {
			out = append(out, n)
			last = n
		}
	}
	return out
}

func printFiles(files []*torrent.File) {
	for i, f := range files {
		fmt.Printf("%6d  %8s  %s\n", i+1, humanBytes(f.Length()), f.Path())
	}
}

func waitForSelected(ctx context.Context, files []*torrent.File, selected []int, total int64) error {
	if total == 0 {
		return nil
	}

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	lastDone := selectedBytesCompleted(files, selected)
	lastAt := time.Now()
	printProgress(lastDone, total, 0)

	for {
		select {
		case <-ctx.Done():
			fmt.Fprintln(os.Stderr)
			return ctx.Err()
		case now := <-ticker.C:
			done := selectedBytesCompleted(files, selected)
			elapsed := now.Sub(lastAt).Seconds()
			rate := float64(done-lastDone) / elapsed
			if rate < 0 {
				rate = 0
			}
			printProgress(done, total, rate)

			if done >= total {
				fmt.Fprintln(os.Stderr)
				return nil
			}

			lastDone = done
			lastAt = now
		}
	}
}

func publishSelected(stageDir, outDir string, files []*torrent.File, selected []int) error {
	for _, idx := range selected {
		f := files[idx]
		rel, err := safeTorrentPath(f.Path())
		if err != nil {
			return fmt.Errorf("niebezpieczna sciezka w torrencie %q: %w", f.Path(), err)
		}

		src := filepath.Join(stageDir, rel)
		dst := filepath.Join(outDir, rel)
		if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
			return fmt.Errorf("utworzenie katalogu dla %q: %w", dst, err)
		}

		if f.Length() == 0 {
			if _, err := os.Stat(src); errors.Is(err, os.ErrNotExist) {
				if err := os.WriteFile(dst, nil, 0o644); err != nil {
					return fmt.Errorf("utworzenie pustego pliku %q: %w", dst, err)
				}
				continue
			}
		}

		if err := moveFile(src, dst); err != nil {
			return fmt.Errorf("zapis wybranego pliku %q: %w", f.Path(), err)
		}
	}
	return nil
}

func safeTorrentPath(path string) (string, error) {
	rel := filepath.Clean(filepath.FromSlash(path))
	if rel == "." || rel == "" || filepath.IsAbs(rel) || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("sciezka wychodzi poza katalog docelowy")
	}
	return rel, nil
}

func moveFile(src, dst string) error {
	if err := os.Rename(src, dst); err == nil {
		return nil
	}

	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	info, err := in.Stat()
	if err != nil {
		return err
	}

	tmp, err := os.CreateTemp(filepath.Dir(dst), ".minerva-pick-copy-")
	if err != nil {
		return err
	}
	tmpName := tmp.Name()
	cleanup := true
	defer func() {
		if cleanup {
			os.Remove(tmpName)
		}
	}()

	if err := tmp.Chmod(info.Mode().Perm()); err != nil {
		tmp.Close()
		return err
	}
	if _, err := io.Copy(tmp, in); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if err := os.Rename(tmpName, dst); err != nil {
		return err
	}
	cleanup = false
	if err := os.Remove(src); err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}
	return nil
}

func selectedBytesCompleted(files []*torrent.File, selected []int) int64 {
	var done int64
	for _, idx := range selected {
		f := files[idx]
		n := f.BytesCompleted()
		if n > f.Length() {
			n = f.Length()
		}
		done += n
	}
	return done
}

func printProgress(done, total int64, bytesPerSecond float64) {
	pct := float64(done) / float64(total) * 100
	fmt.Fprintf(os.Stderr, "\r%6.2f%%  %s / %s  %s/s", pct, humanBytes(done), humanBytes(total), humanBytes(int64(bytesPerSecond)))
}

func humanBytes(n int64) string {
	const unit = int64(1024)
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := unit, 0
	for value := n / unit; value >= unit && exp < 5; value /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(n)/float64(div), "KMGTPE"[exp])
}