package main

import (
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/nixc/reverse-ssh-traefik/internal/server"
)

func envOr(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func envRequired(key string) string {
	v := os.Getenv(key)
	if v == "" {
		log.Fatalf("Required environment variable %s is not set", key)
	}
	return v
}

func envInt(key string, fallback int) int {
	v := os.Getenv(key)
	if v == "" {
		return fallback
	}
	n, err := strconv.Atoi(v)
	if err != nil {
		log.Printf("WARN: invalid %s=%q, using default %d", key, v, fallback)
		return fallback
	}
	return n
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Println("tunnel-server starting")

	// SSH server config (for accepting tunnel clients).
	sshPort := envOr("SSH_PORT", "2222")
	hostKeyPath := envOr("SSH_HOST_KEY", "/keys/host_key")
	authKeysPath := envOr("AUTHORIZED_KEYS", "/keys/authorized_keys")
	portStart := envInt("PORT_RANGE_START", 10000)
	portEnd := envInt("PORT_RANGE_END", 10100)

	// Swarm manager SSH config (for updating service labels).
	traefikHost := envRequired("TRAEFIK_SSH_HOST")
	traefikUser := envOr("TRAEFIK_SSH_USER", "root")
	traefikKey := envRequired("TRAEFIK_SSH_KEY")
	serviceName := envOr("SWARM_SERVICE_NAME", "better-argo-tunnels_tunnel-server")
	entrypoint := envOr("TRAEFIK_ENTRYPOINT", "websecure")
	certResolver := envOr("TRAEFIK_CERT_RESOLVER", "letsencryptresolver")

	// Load the SSH key for connecting to the Swarm manager.
	traefikSigner, err := server.LoadSigner(traefikKey)
	if err != nil {
		log.Fatalf("Failed to load Traefik SSH key: %v", err)
	}
	log.Printf("Loaded Swarm manager SSH key")

	// Initialize port pool.
	pool := server.NewPortPool(portStart, portEnd)
	log.Printf("Port pool: %d-%d (%d ports)", portStart, portEnd, portEnd-portStart+1)

	// Initialize label manager (Swarm service update via SSH).
	labels, err := server.NewLabelManager(
		traefikHost, traefikUser, traefikSigner,
		serviceName, entrypoint, certResolver,
	)
	if err != nil {
		log.Fatalf("Failed to init label manager: %v", err)
	}
	defer labels.Close()

	// Initialize SSH server for tunnel clients.
	sshSrv, err := server.NewSSHServer(hostKeyPath, authKeysPath, pool, labels)
	if err != nil {
		log.Fatalf("Failed to init SSH server: %v", err)
	}

	// Handle graceful shutdown.
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigCh
		log.Printf("Received signal %s, shutting down", sig)
		os.Exit(0)
	}()

	// Start SSH server.
	addr := "0.0.0.0:" + sshPort
	if err := sshSrv.ListenAndServe(addr); err != nil {
		log.Fatalf("SSH server error: %v", err)
	}
}