package main

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"net/netip"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"golang.getoutline.org/tunnel-server/ipinfo"
	"github.com/google/uuid"
	"github.com/lmittmann/tint"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/spf13/cobra"
	"golang.org/x/term"
	"gopkg.in/yaml.v3"

	"sourcecraft.dev/bigbes/shroud/internal/api"
	"sourcecraft.dev/bigbes/shroud/internal/awgserver"
	"sourcecraft.dev/bigbes/shroud/internal/config"
	"sourcecraft.dev/bigbes/shroud/internal/metrics"
	"sourcecraft.dev/bigbes/shroud/internal/ssserver"
	"sourcecraft.dev/bigbes/shroud/internal/store"
)

var version = "dev"

func main() {
	if err := newRootCmd().Execute(); err != nil {
		os.Exit(1)
	}
}

func newRootCmd() *cobra.Command {
	var (
		configFile string
		verbose    bool
	)

	root := &cobra.Command{
		Use:     "shroud",
		Short:   "Shadowsocks + AmneziaWG VPN server",
		Version: version,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runServe(configFile, verbose)
		},
		SilenceUsage:  true,
		SilenceErrors: true,
	}

	pf := root.PersistentFlags()
	pf.StringVarP(&configFile, "config", "c", "config.yaml", "configuration file path")
	pf.BoolVarP(&verbose, "verbose", "v", false, "enable debug logging")

	root.AddCommand(
		newCompletionCmd(),
		newKeyCmd(&configFile),
		newServerCmd(&configFile),
	)

	return root
}

// --- completion ---

func newCompletionCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "completion [bash|zsh|fish|powershell]",
		Short: "Generate shell completion scripts",
		Long: `Generate shell completion scripts for shroud.

To load completions:

Bash:
  $ source <(shroud completion bash)
  # To install permanently:
  $ shroud completion bash > /etc/bash_completion.d/shroud

Zsh:
  $ source <(shroud completion zsh)
  # To install permanently:
  $ shroud completion zsh > "${fpath[1]}/_shroud"

Fish:
  $ shroud completion fish | source
  # To install permanently:
  $ shroud completion fish > ~/.config/fish/completions/shroud.fish

PowerShell:
  PS> shroud completion powershell | Out-String | Invoke-Expression
`,
		DisableFlagsInUseLine: true,
		ValidArgs:             []string{"bash", "zsh", "fish", "powershell"},
		Args:                  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			switch args[0] {
			case "bash":
				return cmd.Root().GenBashCompletion(os.Stdout)
			case "zsh":
				return cmd.Root().GenZshCompletion(os.Stdout)
			case "fish":
				return cmd.Root().GenFishCompletion(os.Stdout, true)
			case "powershell":
				return cmd.Root().GenPowerShellCompletionWithDesc(os.Stdout)
			default:
				return fmt.Errorf("unsupported shell: %s", args[0])
			}
		},
	}
	return cmd
}

// --- key commands ---

func newKeyCmd(configFile *string) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "key",
		Short: "Manage access keys",
	}
	cmd.AddCommand(
		newKeyListCmd(configFile),
		newKeyAddCmd(configFile),
		newKeyRemoveCmd(configFile),
		newKeyRenameCmd(configFile),
	)
	return cmd
}

func newKeyListCmd(configFile *string) *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "List all access keys",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			keys := st.ListKeys()
			if len(keys) == 0 {
				fmt.Println("No access keys configured.")
				return nil
			}
			for _, k := range keys {
				limit := "none"
				if k.DataLimit != nil {
					limit = fmt.Sprintf("%d bytes", k.DataLimit.Bytes)
				}
				awgInfo := ""
				if k.AWG != nil {
					awgInfo = fmt.Sprintf("  awg=%s", k.AWG.AllowedIP)
				}
				fmt.Printf("%-6s  %-20s  port=%-6d  cipher=%s  limit=%s%s\n",
					k.ID, k.Name, k.Port, k.Method, limit, awgInfo)
			}
			return nil
		},
	}
}

func newKeyAddCmd(configFile *string) *cobra.Command {
	var (
		name   string
		port   int
		cipher string
	)
	cmd := &cobra.Command{
		Use:   "add",
		Short: "Add a new access key",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load(*configFile)
			if err != nil {
				return fmt.Errorf("loading config: %w", err)
			}
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			srv := st.GetServer()

			password, err := store.GeneratePassword()
			if err != nil {
				return err
			}

			id := fmt.Sprintf("%d", srv.NextID)
			if err := st.UpdateServer(func(s *store.ServerState) {
				s.NextID++
			}); err != nil {
				return err
			}

			if cipher == "" {
				cipher = srv.DefaultCipher
			}
			if port == 0 {
				port = srv.PortForNewAccessKeys
			}

			ak := store.AccessKey{
				ID:       id,
				Name:     name,
				Password: password,
				Port:     port,
				Method:   cipher,
			}

			// Generate AWG credentials if AWG is enabled.
			if cfg.AmneziaWG.Enabled && cfg.AmneziaWG.Address != "" {
				privKey, err := awgserver.GeneratePrivateKey()
				if err != nil {
					return fmt.Errorf("generating AWG key: %w", err)
				}
				pubKey, err := awgserver.PublicKeyFromPrivate(privKey)
				if err != nil {
					return fmt.Errorf("deriving AWG public key: %w", err)
				}
				keys := st.ListKeys()
				used := make([]string, 0, len(keys))
				for _, k := range keys {
					if k.AWG != nil {
						used = append(used, k.AWG.AllowedIP)
					}
				}
				alloc, err := awgserver.NewIPAllocator(cfg.AmneziaWG.Address)
				if err != nil {
					return fmt.Errorf("creating IP allocator: %w", err)
				}
				ip, err := alloc.Allocate(used)
				if err != nil {
					return fmt.Errorf("allocating AWG IP: %w", err)
				}
				ak.AWG = &store.AWGKeyData{
					PrivateKey: privKey,
					PublicKey:  pubKey,
					AllowedIP:  ip,
				}
			}

			if err := st.CreateKey(ak); err != nil {
				return err
			}
			fmt.Printf("Created key %s (port=%d, cipher=%s)\n", id, port, cipher)
			fmt.Printf("  ss://%s\n", password)
			if ak.AWG != nil {
				fmt.Printf("  awg: %s (pubkey=%s)\n", ak.AWG.AllowedIP, ak.AWG.PublicKey)
			}
			return nil
		},
	}
	cmd.Flags().StringVarP(&name, "name", "n", "", "key name")
	cmd.Flags().IntVarP(&port, "port", "p", 0, "port (default: server's default port)")
	cmd.Flags().StringVar(&cipher, "cipher", "", "cipher (default: server's default cipher)")
	return cmd
}

func newKeyRemoveCmd(configFile *string) *cobra.Command {
	return &cobra.Command{
		Use:   "remove <id>",
		Short: "Remove an access key",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			if err := st.DeleteKey(args[0]); err != nil {
				return err
			}
			fmt.Printf("Removed key %s\n", args[0])
			return nil
		},
	}
}

func newKeyRenameCmd(configFile *string) *cobra.Command {
	return &cobra.Command{
		Use:   "rename <id> <name>",
		Short: "Rename an access key",
		Args:  cobra.ExactArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			return st.UpdateKey(args[0], func(k *store.AccessKey) {
				k.Name = args[1]
			})
		},
	}
}

// --- server commands ---

func newServerCmd(configFile *string) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "server",
		Short: "Server configuration commands",
	}
	cmd.AddCommand(
		newServerInfoCmd(configFile),
		newServerSetPortCmd(configFile),
		newServerSetHostnameCmd(configFile),
		newServerSetNameCmd(configFile),
	)
	return cmd
}

func newServerInfoCmd(configFile *string) *cobra.Command {
	return &cobra.Command{
		Use:   "info",
		Short: "Show server configuration",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			srv := st.GetServer()
			fmt.Printf("Server ID:        %s\n", srv.ID)
			fmt.Printf("Name:             %s\n", srv.Name)
			fmt.Printf("Hostname:         %s\n", srv.Hostname)
			fmt.Printf("Default Port:     %d\n", srv.PortForNewAccessKeys)
			fmt.Printf("Default Cipher:   %s\n", srv.DefaultCipher)
			fmt.Printf("Metrics Enabled:  %v\n", srv.MetricsEnabled)
			fmt.Printf("Access Keys:      %d\n", len(st.ListKeys()))
			if srv.AccessKeyDataLimit != nil {
				fmt.Printf("Data Limit:       %d bytes\n", srv.AccessKeyDataLimit.Bytes)
			}
			return nil
		},
	}
}

func newServerSetPortCmd(configFile *string) *cobra.Command {
	return &cobra.Command{
		Use:   "set-port <port>",
		Short: "Set the default port for new access keys",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			var port int
			if _, err := fmt.Sscanf(args[0], "%d", &port); err != nil || port < 1 || port > 65535 {
				return fmt.Errorf("invalid port: %s", args[0])
			}
			if err := st.UpdateServer(func(s *store.ServerState) {
				s.PortForNewAccessKeys = port
			}); err != nil {
				return err
			}
			fmt.Printf("Default port set to %d\n", port)
			return nil
		},
	}
}

func newServerSetHostnameCmd(configFile *string) *cobra.Command {
	return &cobra.Command{
		Use:   "set-hostname <hostname>",
		Short: "Set the hostname for access key URLs",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			if err := st.UpdateServer(func(s *store.ServerState) {
				s.Hostname = args[0]
			}); err != nil {
				return err
			}
			fmt.Printf("Hostname set to %s\n", args[0])
			return nil
		},
	}
}

func newServerSetNameCmd(configFile *string) *cobra.Command {
	return &cobra.Command{
		Use:   "set-name <name>",
		Short: "Set the server display name",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			st, err := openStore(*configFile)
			if err != nil {
				return err
			}
			if err := st.UpdateServer(func(s *store.ServerState) {
				s.Name = args[0]
			}); err != nil {
				return err
			}
			fmt.Printf("Server name set to %s\n", args[0])
			return nil
		},
	}
}

// --- helpers ---

func openStore(configFile string) (*store.YAMLFileStore, error) {
	cfg, err := config.Load(configFile)
	if err != nil {
		return nil, fmt.Errorf("loading config: %w", err)
	}
	return store.NewYAMLFileStore(
		cfg.StateFile(),
		"", cfg.Server.Name, cfg.Server.Hostname,
		cfg.Shadowsocks.DefaultCipher, cfg.Shadowsocks.DefaultPort,
	)
}

// --- access info YAML ---

type accessInfo struct {
	APIURL     string `yaml:"api_url"`
	CertSHA256 string `yaml:"cert_sha256"`
}

func writeAccessInfo(apiAddr, secret, certFile, infoFile string) error {
	certSHA256 := ""
	if certFile != "" {
		cert, err := tls.LoadX509KeyPair(certFile, certFile)
		if err == nil && len(cert.Certificate) > 0 {
			hash := sha256.Sum256(cert.Certificate[0])
			certSHA256 = hex.EncodeToString(hash[:])
		}
	}

	info := accessInfo{
		APIURL:     fmt.Sprintf("https://%s/%s", apiAddr, secret),
		CertSHA256: certSHA256,
	}

	data, err := yaml.Marshal(info)
	if err != nil {
		return err
	}
	return os.WriteFile(infoFile, data, 0600)
}

// --- serve ---

func runServe(configFile string, verbose bool) error {
	logLevel := new(slog.LevelVar)
	logHandler := tint.NewHandler(
		os.Stderr,
		&tint.Options{NoColor: !term.IsTerminal(int(os.Stderr.Fd())), Level: logLevel},
	)
	logger := slog.New(logHandler)
	slog.SetDefault(logger)

	if verbose {
		logLevel.Set(slog.LevelDebug)
	}

	cfg, err := config.Load(configFile)
	if err != nil {
		return fmt.Errorf("loading config: %w", err)
	}

	// Generate server ID if not set.
	serverID := cfg.Server.ID
	if serverID == "" {
		serverID = uuid.New().String()
	}

	// Generate API secret if not set.
	apiSecret := cfg.API.Secret
	if apiSecret == "" {
		buf := make([]byte, 16)
		if _, err := rand.Read(buf); err != nil {
			return fmt.Errorf("generating API secret: %w", err)
		}
		apiSecret = hex.EncodeToString(buf)
		logger.Info("Generated API secret.", "secret", apiSecret)
	}

	// Set up IP info for geo-metrics.
	var ip2info ipinfo.IPInfoMap
	if cfg.Shadowsocks.IPCountryDB != "" || cfg.Shadowsocks.IPASNDB != "" {
		mmdb, err := ipinfo.NewMMDBIPInfoMap(cfg.Shadowsocks.IPCountryDB, cfg.Shadowsocks.IPASNDB)
		if err != nil {
			return fmt.Errorf("loading IP info databases: %w", err)
		}
		defer mmdb.Close()
		ip2info = mmdb
	}

	// Set up metrics.
	serverMetrics := metrics.NewServerMetrics()
	registry, transferTracker, _, err := metrics.SetupRegistry(
		ip2info,
		serverMetrics,
		version,
		cfg.Metrics.NodeExporterCollectors,
		logger,
	)
	if err != nil {
		return fmt.Errorf("setting up metrics: %w", err)
	}

	// Set up persistent store.
	st, err := store.NewYAMLFileStore(
		cfg.StateFile(),
		serverID,
		cfg.Server.Name,
		cfg.Server.Hostname,
		cfg.Shadowsocks.DefaultCipher,
		cfg.Shadowsocks.DefaultPort,
	)
	if err != nil {
		return fmt.Errorf("initializing store: %w", err)
	}

	// Set up Shadowsocks server (optional).
	var ss *ssserver.Server
	if *cfg.Shadowsocks.Enabled {
		natTimeout, err := time.ParseDuration(cfg.Shadowsocks.NATTimeout)
		if err != nil {
			return fmt.Errorf("invalid NAT timeout: %w", err)
		}

		// Pick a random unused port for Shadowsocks if not configured.
		ssDefaultPort := cfg.Shadowsocks.DefaultPort
		if ssDefaultPort == 0 {
			p, err := pickRandomPort()
			if err != nil {
				return fmt.Errorf("picking random port for shadowsocks: %w", err)
			}
			ssDefaultPort = p
			logger.Info("Selected random port for new access keys.", "port", ssDefaultPort)
		}

		ss = ssserver.New(
			natTimeout,
			cfg.Shadowsocks.ReplayHistory,
			serverMetrics,
			transferTracker,
			logger,
		)

		// Start with existing keys.
		if keys := st.ListKeys(); len(keys) > 0 {
			if err := ss.SyncKeys(keys); err != nil {
				return fmt.Errorf("starting Shadowsocks with existing keys: %w", err)
			}
			logger.Info("Loaded existing access keys.", "count", len(keys))
		}
		logger.Info("Shadowsocks server enabled.")
	}

	// Set up AmneziaWG server (optional).
	var awg *awgserver.Server
	if cfg.AmneziaWG.Enabled {
		// Ensure server has AWG keypair.
		srv := st.GetServer()
		if srv.AWGPrivateKey == "" {
			privKey, err := awgserver.GeneratePrivateKey()
			if err != nil {
				return fmt.Errorf("generating AWG server key: %w", err)
			}
			pubKey, err := awgserver.PublicKeyFromPrivate(privKey)
			if err != nil {
				return fmt.Errorf("deriving AWG server public key: %w", err)
			}
			if err := st.UpdateServer(func(s *store.ServerState) {
				s.AWGPrivateKey = privKey
				s.AWGPublicKey = pubKey
			}); err != nil {
				return fmt.Errorf("persisting AWG server key: %w", err)
			}
			logger.Info("Generated AWG server keypair.")
			srv = st.GetServer()
		}

		// Parse H1-H4 ranges.
		h1, err := awgserver.ParseHeaderRange(cfg.AmneziaWG.H1)
		if err != nil {
			return fmt.Errorf("parsing AWG H1: %w", err)
		}
		h2, err := awgserver.ParseHeaderRange(cfg.AmneziaWG.H2)
		if err != nil {
			return fmt.Errorf("parsing AWG H2: %w", err)
		}
		h3, err := awgserver.ParseHeaderRange(cfg.AmneziaWG.H3)
		if err != nil {
			return fmt.Errorf("parsing AWG H3: %w", err)
		}
		h4, err := awgserver.ParseHeaderRange(cfg.AmneziaWG.H4)
		if err != nil {
			return fmt.Errorf("parsing AWG H4: %w", err)
		}

		// Build server address from subnet (take .1).
		alloc, err := awgserver.NewIPAllocator(cfg.AmneziaWG.Address)
		if err != nil {
			return fmt.Errorf("parsing AWG address: %w", err)
		}
		serverAddr := fmt.Sprintf("%s/%d", alloc.ServerIP().String(), netip.MustParsePrefix(cfg.AmneziaWG.Address).Bits())

		awgCfg := awgserver.Config{
			ListenPort:   cfg.AmneziaWG.ListenPort,
			TUNName:      cfg.AmneziaWG.TUNName,
			Address:      serverAddr,
			MTU:          cfg.AmneziaWG.MTU,
			PrivateKey:   srv.AWGPrivateKey,
			MuxEnabled:   cfg.AmneziaWG.AWGMuxEnabled(),
			Domain:       cfg.AmneziaWG.Domain,
			CertCache:    cfg.ACME.CertCache,
			ACMEHTTPPort: cfg.ACME.HTTPPort,
			Jc:           cfg.AmneziaWG.Jc,
			Jmin:         cfg.AmneziaWG.Jmin,
			Jmax:         cfg.AmneziaWG.Jmax,
			S1:           cfg.AmneziaWG.S1,
			S2:           cfg.AmneziaWG.S2,
			S3:           cfg.AmneziaWG.S3,
			S4:           cfg.AmneziaWG.S4,
			H1:           h1,
			H2:           h2,
			H3:           h3,
			H4:           h4,
		}

		awg = awgserver.New(awgCfg, logger)
		if err := awg.Start(); err != nil {
			return fmt.Errorf("starting AWG server: %w", err)
		}

		if keys := st.ListKeys(); len(keys) > 0 {
			if err := awg.SyncKeys(keys); err != nil {
				return fmt.Errorf("syncing AWG peers: %w", err)
			}
		}
		logger.Info("AmneziaWG server started.", "port", cfg.AmneziaWG.ListenPort)
	}

	// Start Prometheus metrics + /healthz endpoint.
	metricsHandler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
	metricsMux := http.NewServeMux()
	metricsMux.Handle("/metrics", metricsHandler)
	metricsMux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(map[string]any{
			"healthy":    true,
			"version":    version,
			"accessKeys": len(st.ListKeys()),
		})
	})
	metricsServer := &http.Server{Addr: cfg.Metrics.ListenAddr, Handler: metricsMux}
	go func() {
		logger.Info("Prometheus metrics server started.", "address", cfg.Metrics.ListenAddr)
		if err := metricsServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logger.Error("Metrics server failed.", "err", err)
		}
	}()

	// Start REST API.
	var awgCfgPtr *config.AmneziaWGConfig
	if cfg.AmneziaWG.Enabled {
		awgCfgPtr = &cfg.AmneziaWG
	}
	handler := api.NewHandler(st, ss, awg, awgCfgPtr, transferTracker, version, logger)
	router := api.NewRouter(apiSecret, handler)
	apiServer := &http.Server{Addr: cfg.API.ListenAddr, Handler: router}
	go func() {
		logger.Info("Management API started.", "address", cfg.API.ListenAddr, "prefix", "/"+apiSecret)
		if err := apiServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logger.Error("API server failed.", "err", err)
		}
	}()

	// Write access info.
	infoFile := filepath.Join(filepath.Dir(cfg.StateFile()), "access.yaml")
	if err := writeAccessInfo(cfg.API.ListenAddr, apiSecret, cfg.API.CertFile, infoFile); err != nil {
		logger.Warn("Failed to write access info.", "err", err)
	} else {
		logger.Info("Access info written.", "file", infoFile)
	}

	logger.Info("shroud started.",
		"version", version,
		"api", cfg.API.ListenAddr,
		"metrics", cfg.Metrics.ListenAddr,
	)

	// Wait for shutdown signal.
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	<-sigCh

	logger.Info("Shutting down...")
	if awg != nil {
		awg.Stop()
	}
	if ss != nil {
		ss.Stop()
	}
	apiServer.Close()
	metricsServer.Close()

	return nil
}

func pickRandomPort() (int, error) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	port := ln.Addr().(*net.TCPAddr).Port
	ln.Close()
	return port, nil
}