package main

import (
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"time"
	"unicode"

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

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

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

// cleanHost strips BOM, NULs, and all Unicode whitespace (including NBSP) from ends.
func cleanHost(s string) string {
	s = strings.ReplaceAll(s, "\x00", "")
	s = strings.TrimPrefix(s, "\ufeff")
	s = strings.TrimFunc(s, unicode.IsSpace)
	return s
}

// resolveBackendHost picks backend hostname. Priority:
//  1) TUNNEL_HOST_FILE (path to a file containing hostname, e.g. Docker secret mount)
//  2) TUNNEL_HOST
//  3) TUNNEL_BACKEND_HOST (alias for stacks that reserve TUNNEL_HOST)
//  4) 127.0.0.1
func resolveBackendHost() (host string, source string) {
	path := cleanHost(os.Getenv("TUNNEL_HOST_FILE"))
	if path != "" {
		b, err := os.ReadFile(path)
		if err != nil {
			log.Fatalf("TUNNEL_HOST_FILE %q: %v", path, err)
		}
		h := cleanHost(string(b))
		if h != "" {
			return h, fmt.Sprintf("TUNNEL_HOST_FILE(%s)", path)
		}
		log.Fatalf("TUNNEL_HOST_FILE %q is empty after trim", path)
	}
	for _, key := range []string{"TUNNEL_HOST", "TUNNEL_BACKEND_HOST"} {
		raw := os.Getenv(key)
		h := cleanHost(raw)
		if h != "" {
			return h, key
		}
		if raw != "" && h == "" {
			log.Printf("tunnel-client: %s is set but only whitespace/control after clean (len=%d); trying next key / default", key, len(raw))
		}
	}
	return "127.0.0.1", "default(127.0.0.1)"
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Printf("tunnel-client starting commit=%s", buildinfo.Commit)

	serverAddr := envRequired("TUNNEL_SERVER")
	domain := envRequired("TUNNEL_DOMAIN")
	keyPath := envOr("TUNNEL_KEY", "/keys/id_ed25519")

	// Optional HTTP Basic Auth credentials for Traefik middleware.
	authUser := envOr("TUNNEL_AUTH_USER", "")
	authPass := envOr("TUNNEL_AUTH_PASS", "")

	localHost, hostSrc := resolveBackendHost()
	localPortStr := cleanHost(os.Getenv("TUNNEL_PORT"))
	if localPortStr == "" {
		localPortStr = "8080"
	}
	localPort, err := strconv.Atoi(localPortStr)
	if err != nil {
		log.Fatalf("Invalid TUNNEL_PORT=%q: %v", localPortStr, err)
	}
	log.Printf("tunnel-client config: backend host %q (source=%s); port=%d (TUNNEL_PORT raw=%q)",
		localHost, hostSrc, localPort, os.Getenv("TUNNEL_PORT"))
	log.Printf("tunnel-client env (for debugging): TUNNEL_HOST=%q TUNNEL_BACKEND_HOST=%q TUNNEL_HOST_FILE=%q",
		os.Getenv("TUNNEL_HOST"), os.Getenv("TUNNEL_BACKEND_HOST"), os.Getenv("TUNNEL_HOST_FILE"))

	// Load the private key.
	signer, err := client.LoadPrivateKey(keyPath)
	if err != nil {
		log.Fatalf("Failed to load private key: %v", err)
	}
	log.Printf("Loaded key from %s", keyPath)

	// Reconnect loop.
	backoff := time.Second
	maxBackoff := 30 * time.Second

	for {
		if authUser != "" {
			log.Printf("Connecting to %s (domain=%s, local=%s:%d, basicauth=enabled)", serverAddr, domain, localHost, localPort)
		} else {
			log.Printf("Connecting to %s (domain=%s, local=%s:%d)", serverAddr, domain, localHost, localPort)
		}

		sshClient, err := client.Connect(serverAddr, signer)
		if err != nil {
			log.Printf("Connection failed: %v (retry in %s)", err, backoff)
			time.Sleep(backoff)
			backoff = min(backoff*2, maxBackoff)
			continue
		}

		// Reset backoff on successful connection.
		backoff = time.Second
		log.Printf("Connected to %s", serverAddr)

		// Set up the reverse tunnel (blocks until disconnected).
		if err := client.SetupTunnel(sshClient, domain, localHost, localPort, authUser, authPass); err != nil {
			log.Printf("Tunnel error: %v (reconnecting in %s)", err, backoff)
		}

		sshClient.Close()
		time.Sleep(backoff)
		backoff = min(backoff*2, maxBackoff)
	}
}