main.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	runMatchmakerFlag := flag.Bool("run-matchmaker", false, "Start self-hosted matchmaker server (signalling + STUN)")
	matchmakerURL := flag.String("matchmaker", YeetMatchmakerServer, "Custom matchmaker server URL")
	noMatchmakerFlag := flag.Bool("no-matchmaker", false, "Do not connect to external matchmaker server (listen locally for direct IP/LAN transfers only)")
	receiverIPFlag := flag.String("receiver-ip", "", "Connect directly to receiver IP address (bypasses external matchmaker)")
	addr := flag.String("addr", ":8080", "Address for matchmaker HTTP server to listen on")
	stunAddr := flag.String("stun-addr", ":3478", "Address for matchmaker STUN server to listen on (UDP)")
	behindProxy := flag.Bool("behind-proxy", false, "Trust proxy headers for rate limiting (X-Forwarded-For, X-Real-IP)")

	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s:\n\n", filepath.Base(os.Args[0]))
		fmt.Fprintf(flag.CommandLine.Output(), "To receive a file:\n")
		fmt.Fprintf(flag.CommandLine.Output(), "  %s [-matchmaker <url>] [-no-matchmaker]\n\n", filepath.Base(os.Args[0]))
		fmt.Fprintf(flag.CommandLine.Output(), "To send files:\n")
		fmt.Fprintf(flag.CommandLine.Output(), "  %s [-matchmaker <url>] [-receiver-ip <ip_addr>] <filename1> [<filename2> ...]\n\n", filepath.Base(os.Args[0]))
		fmt.Fprintf(flag.CommandLine.Output(), "To start a self-hosted match-making server (signalling + STUN):\n")
		fmt.Fprintf(flag.CommandLine.Output(), "  %s -run-matchmaker [-addr <addr>] [-stun-addr <addr>] [-behind-proxy]\n\n", filepath.Base(os.Args[0]))
		fmt.Fprintf(flag.CommandLine.Output(), "Flags:\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	if *runMatchmakerFlag {
		runMatchmaker(*addr, *stunAddr, *behindProxy)
		return
	}

	var serverURL string
	if *noMatchmakerFlag {
		serverURL = ""
	} else {
		serverURL = *matchmakerURL
		if serverURL == "" {
			serverURL = YeetMatchmakerServer
		}
	}

	args := flag.Args()
	if len(args) == 0 {
		runReceive(serverURL)
		return
	}

	runSend(serverURL, *receiverIPFlag, args)
}

func runMatchmaker(addr, stunAddr string, behindProxy bool) {
	mm := NewMatchmaker()
	mm.BehindProxy = behindProxy
	err := mm.Start(addr, stunAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "❌ Error starting matchmaker server: %v\n", err)
		os.Exit(1)
	}
	select {}
}

func runSend(serverURL, receiverIP string, filenames []string) {
	if receiverIP != "" {
		serverURL = FormatReceiverURL(receiverIP)
	}

	fmt.Printf("Enter Session ID: ")
	sessionID := readLine()

	sender, err := NewSender(serverURL, SessionID(sessionID))
	if err != nil {
		errStr := err.Error()
		if strings.Contains(errStr, "no such host") || strings.Contains(errStr, "unreachable") {
			fmt.Fprintf(os.Stderr, "❌ Error: Signalling server unreachable (offline or DNS failed).\n")
			fmt.Fprintf(os.Stderr, "💡 Tip: Check your internet connection, or make sure the receiver is running on the same local network (LAN) and has generated Session ID %s.\n", sessionID)
		} else {
			fmt.Fprintf(os.Stderr, "❌ Error: %v\n", err)
		}
		os.Exit(1)
	}
	defer sender.Close()

	sender.OnProgress = newCLIProgressHandler("📤 Yeeting", "yeeted successfully!")

	if receiverIP != "" {
		fmt.Printf("🔗 Connected directly to receiver at %s! Handshaking...\n", serverURL)
	} else {
		fmt.Println("🔗 Connected to signalling server! Handshaking with receiver...")
	}

	for _, filename := range filenames {
		if err := sender.Send(filename); err != nil {
			fmt.Fprintf(os.Stderr, "\n❌ Error sending file %s: %v\n", filename, err)
			os.Exit(1)
		}
	}
}

func runReceive(serverURL string) {
	receiver, err := NewReceiver(serverURL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "❌ Error starting receiver: %v\n", err)
		os.Exit(1)
	}
	defer receiver.Close()

	receiver.OnProgress = newCLIProgressHandler("📥 Downloading", "received successfully!")

	fmt.Printf("Your Session ID: %s\n", receiver.SessionID)
	if localIP, err := GetLocalIP(); err == nil && localIP != "127.0.0.1" {
		_, portStr, _ := net.SplitHostPort(strings.TrimPrefix(receiver.LocalServerURL, "http://"))
		if portStr != "" && portStr != DefaultReceiverPort {
			fmt.Printf("Local IP: %s:%s\n", localIP, portStr)
		} else {
			fmt.Printf("Local IP: %s\n", localIP)
		}
	}
	fmt.Println("Waiting for a sender to connect...")

	senderName := <-receiver.SenderRequest()
	fmt.Printf("\nConnection request from '%s'. Accept? (y/n): ", senderName)
	answer := readLine()
	if strings.ToLower(answer) != "y" && strings.ToLower(answer) != "yes" {
		if err := receiver.RejectConnection(); err != nil {
			fmt.Printf("❌ Error rejecting connection: %v\n", err)
		}
		fmt.Println("❌ Connection rejected.")
		return
	}

	fmt.Println("Connection accepted! Establishing direct P2P link...")
	if err := receiver.ApproveConnection(); err != nil {
		fmt.Fprintf(os.Stderr, "❌ Error accepting connection: %v\n", err)
		os.Exit(1)
	}

	senderToken := <-receiver.SenderAnswer()
	if err := receiver.Connect(senderToken); err != nil {
		fmt.Fprintf(os.Stderr, "❌ Error establishing connection: %v\n", err)
		os.Exit(1)
	}

	for tr := range receiver.TransferRequest() {
		if err := receiver.Accept(tr); err != nil {
			fmt.Fprintf(os.Stderr, "❌ Error accepting file transfer: %v\n", err)
			os.Exit(1)
		}

		if err := <-receiver.Done(); err != nil {
			fmt.Printf("❌ Error during transfer: %v\n", err)
			os.Exit(1)
		}
	}
}

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

func readLine() string {
	reader := bufio.NewReader(os.Stdin)
	line, err := reader.ReadString('\n')
	if err != nil {
		fmt.Fprintf(os.Stderr, "❌ Error reading input: %v\n", err)
		os.Exit(1)
	}
	return strings.TrimSpace(line)
}

func newCLIProgressHandler(actionLabel, completeSuffix string) ProgressFunc {
	return func(fileName string, current, total int64) {
		if total <= 0 {
			return
		}
		if current >= total {
			fmt.Printf("\r✨ %s %s\033[K\n", fileName, completeSuffix)
			return
		}
		percent := float64(current) / float64(total) * 100
		fmt.Printf("\r%s %s... %.1f%% (%s / %s)\033[K", actionLabel, truncateString(fileName, 40), percent, formatSize(current), formatSize(total))
	}
}