package main

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"time"

	"git.t-juice.club/torjus/apiary"
	"git.t-juice.club/torjus/apiary/config"
	"git.t-juice.club/torjus/apiary/honeypot/ssh"
	"git.t-juice.club/torjus/apiary/honeypot/ssh/store"
	"git.t-juice.club/torjus/apiary/web"
	"github.com/coreos/go-systemd/daemon"
	sshlib "github.com/gliderlabs/ssh"
	"github.com/urfave/cli/v3"
	"golang.org/x/crypto/acme/autocert"
)

func main() {
	app := &cli.Command{
		Name:    "apiary",
		Version: apiary.FullVersion(),
		Commands: []*cli.Command{
			{
				Name:   "serve",
				Action: ActionServe,
				Usage:  "Start Apiary server",
			},
		},
	}

	if err := app.Run(context.Background(), os.Args); err != nil {
		fmt.Printf("Error: %s\n", err)
		os.Exit(1)
	}
}

func ActionServe(ctx context.Context, cmd *cli.Command) error {
	cfg, err := getConfig()
	if err != nil {
		return err
	}

	// Setup logging
	loggers := setupLoggers(cfg)
	loggers.rootLogger.Info("Starting apiary.", "version", apiary.FullVersion())

	// Setup store
	var s store.LoginAttemptStore
	switch cfg.Store.Type {
	case "MEMORY", "memory":
		loggers.rootLogger.Info("Initialized store.", "store_type", "memory")
		s = &store.MemoryStore{}
	case "POSTGRES", "postgres":
		pgStartTime := time.Now()
		loggers.rootLogger.Debug("Initializing store.", "store_type", "postgres")
		pgStore, err := store.NewPostgresStore(cfg.Store.Postgres.DSN)
		if err != nil {
			return err
		}
		if err := pgStore.InitDB(); err != nil {
			return err
		}
		loggers.rootLogger.Info("Initialized store.", "store_type", "postgres", "init_time", time.Since(pgStartTime))
		if cfg.Store.EnableCache {
			loggers.rootLogger.Debug("Initializing store.", "store_type", "cache-postgres")
			startTime := time.Now()
			cachingStore := store.NewCachingStore(pgStore)
			s = cachingStore
			loggers.rootLogger.Info("Initialized store.", "store_type", "cache-postgres", "init_time", time.Since(startTime))
		} else {
			s = pgStore
		}
	default:
		return fmt.Errorf("Invalid store configured")
	}

	// Setup interrupt handling
	interruptChan := make(chan os.Signal, 1)
	signal.Notify(interruptChan, os.Interrupt)

	rootCtx, rootCancel := context.WithCancel(ctx)
	defer rootCancel()
	serversCtx, serversCancel := context.WithCancel(rootCtx)
	defer serversCancel()

	// Setup metrics collection
	s = store.NewMetricsCollectingStore(rootCtx, s)

	// Setup honeypot
	hs, err := ssh.NewHoneypotServer(cfg.Honeypot, s)
	if err != nil {
		return err
	}
	hs.Logger = loggers.honeypotLogger

	// Setup webserver
	web := web.NewServer(cfg.Frontend, hs, s)
	web.AccessLogger = loggers.webAccessLogger
	web.ServerLogger = loggers.webServerLogger
	if cfg.Frontend.Autocert.Enable {
		certManager := autocert.Manager{
			Prompt:     autocert.AcceptTOS,
			HostPolicy: autocert.HostWhitelist(cfg.Frontend.Autocert.Domains...),
			Email:      cfg.Frontend.Autocert.Email,
		}
		if cfg.Frontend.Autocert.CacheDir != "" {
			certManager.Cache = autocert.DirCache(cfg.Frontend.Autocert.CacheDir)
		}

		tlsConfig := certManager.TLSConfig()
		web.TLSConfig = tlsConfig
	}

	// Handle interrupt
	go func() {
		<-interruptChan
		loggers.rootLogger.Info("Interrupt received, shutting down.")
		serversCancel()
	}()

	// Start ssh server
	go func() {
		loggers.rootLogger.Info("Starting SSH server.", "addr", cfg.Honeypot.ListenAddr)
		if err := hs.ListenAndServe(); err != nil && err != sshlib.ErrServerClosed {
			loggers.rootLogger.Warn("SSH server returned error.", "error", err)
		}
		loggers.rootLogger.Info("SSH server stopped.")
	}()

	// Start web server
	go func() {
		loggers.rootLogger.Info("Starting web server.", "addr", cfg.Frontend.ListenAddr)
		if err := web.StartServe(); err != nil && err != http.ErrServerClosed {
			loggers.rootLogger.Warn("Web server returned error.", "error", err)
		}
	}()

	// If run by systemd, enable watchdog and notify ready
	go func() {
		notifyCtx, cancel := context.WithCancel(rootCtx)
		defer cancel()

		_, ok := os.LookupEnv("NOTIFY_SOCKET")
		if !ok {
			return
		}
		loggers.rootLogger.Info("Systemd notify socket detected. Sending ready and enabling watchdog.")
		ok, err := daemon.SdNotify(false, daemon.SdNotifyReady)
		if !ok {
			loggers.rootLogger.Info("Systemd notify not enabled.")
			return
		}
		if err != nil {
			loggers.rootLogger.Warn("Unable to connect to NOTIFY_SOCKET.", "error", err)
			return
		}
		loggers.rootLogger.Debug("Sent READY=1 to NOTIFY_SOCKET.")

		// Setup timer
		timeout, err := daemon.SdWatchdogEnabled(false)
		if err != nil {
			loggers.rootLogger.Warn("Unable to get watchdog timeout.", "error", err)
			return
		}
		if timeout == 0 {
			loggers.rootLogger.Info("Systemd watchdog not enabled.")
			return
		}

		ticker := time.NewTicker(timeout / 2)
		for {
			// TODO: Consider using ctx with deadline to avoid healthcheck taking more time than ticker,
			// which will result in being killed by watchdog.
			healthy := s.IsHealthy()
			select {
			case <-ticker.C:
				if healthy == nil {
					if _, err := daemon.SdNotify(false, daemon.SdNotifyWatchdog); err != nil {
						loggers.rootLogger.Warn("Error notifying watchdog.", "err", err)
					}
					continue
				}
				// TODO: If unhealthy, should we retry healthcheck immediately, otherwise service will most likely get killed by watchdog.
				loggers.rootLogger.Error("Store reported not healthy, might get killed by watchdog.", "err", healthy)
			case <-notifyCtx.Done():
				loggers.rootLogger.Debug("Notify context cancelled.")
				return
			}
		}
	}()

	go func() {
		<-serversCtx.Done()

		// Stop SSH server
		sshShutdownCtx, sshShutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer sshShutdownCancel()
		loggers.rootLogger.Info("SSH server shutdown started.")
		if err := hs.Shutdown(sshShutdownCtx); err != nil {
			loggers.rootLogger.Info("Error shutting down SSH server.", "error", err)
		}
		loggers.rootLogger.Info("SSH server shutdown complete.")

		// Stop Web server
		webShutdownCtx, webShutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer webShutdownCancel()

		loggers.rootLogger.Info("Web server shutdown started.")
		if err := web.Shutdown(webShutdownCtx); err != nil {
			loggers.rootLogger.Info("Error shutting down web server.", "error", err)
		}
		loggers.rootLogger.Info("Web server shutdown complete.")
		rootCancel()
	}()

	<-rootCtx.Done()

	return nil
}

type loggerCollection struct {
	rootLogger      *slog.Logger
	honeypotLogger  *slog.Logger
	webAccessLogger *slog.Logger
	webServerLogger *slog.Logger
}

func setupLoggers(cfg config.Config) *loggerCollection {
	opts := &slog.HandlerOptions{}
	switch strings.ToUpper(cfg.Honeypot.LogLevel) {
	case "INFO":
		opts.Level = slog.LevelInfo
	case "DEBUG":
		opts.Level = slog.LevelDebug
		opts.AddSource = true
	case "WARN", "WARNING":
		opts.Level = slog.LevelWarn
	case "ERR", "ERROR":
		opts.Level = slog.LevelError
	default:
		opts.Level = slog.LevelInfo
	}

	handler := slog.NewTextHandler(os.Stdout, opts)
	rootLogger := slog.New(handler)

	return &loggerCollection{
		rootLogger:      rootLogger.With("module", "application"),
		honeypotLogger:  rootLogger.With("module", "honeypot"),
		webAccessLogger: rootLogger.With("module", "web-access-log"),
		webServerLogger: rootLogger.With("module", "web-server"),
	}
}

func getConfig() (config.Config, error) {
	defaultLocations := []string{
		"apiary.toml",
		"/etc/apiary.toml",
		"/etc/apiary/apiary.toml",
	}

	for _, fname := range defaultLocations {
		if _, err := os.Stat(fname); os.IsNotExist(err) {
			continue
		}
		cfg, err := config.FromFile(fname)
		if err != nil {
			return config.Config{}, err
		}
		return cfg, nil
	}
	return config.Config{}, fmt.Errorf("Could not find config file")
}
