at-container-registry/pkg/hold/server.go

package hold

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

	"atcr.io/pkg/hold/admin"
	"atcr.io/pkg/hold/billing"
	"atcr.io/pkg/hold/gc"
	"atcr.io/pkg/hold/oci"
	"atcr.io/pkg/hold/pds"
	"atcr.io/pkg/hold/quota"
	"atcr.io/pkg/logging"
	"atcr.io/pkg/s3"

	"github.com/distribution/distribution/v3/registry/storage/driver/factory"
	_ "github.com/distribution/distribution/v3/registry/storage/driver/s3-aws"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
)

// HoldServer is the hold service with an exposed router for extensibility.
// Consumers can add routes to Router before calling Serve().
type HoldServer struct {
	// Router is the chi router. Add routes before calling Serve().
	Router chi.Router

	// PDS is the embedded ATProto PDS. Nil if database path is not configured.
	PDS *pds.HoldPDS

	// QuotaManager manages storage quotas per tier.
	QuotaManager *quota.Manager

	// Config is the hold service configuration.
	Config *Config

	// internal fields for shutdown
	httpServer       *http.Server
	broadcaster      *pds.EventBroadcaster
	scanBroadcaster  *pds.ScanBroadcaster
	garbageCollector *gc.GarbageCollector
	adminUI          *admin.AdminUI
}

// NewHoldServer initializes PDS, storage, quota, XRPC handlers, and returns
// before starting. Consumer can add routes to Router before calling Serve().
func NewHoldServer(cfg *Config) (*HoldServer, error) {
	// Initialize structured logging with optional remote shipping
	logging.InitLoggerWithShipper(cfg.LogLevel, logging.ShipperConfig{
		Backend:       cfg.LogShipper.Backend,
		URL:           cfg.LogShipper.URL,
		BatchSize:     cfg.LogShipper.BatchSize,
		FlushInterval: cfg.LogShipper.FlushInterval,
		Service:       "hold",
		Username:      cfg.LogShipper.Username,
		Password:      cfg.LogShipper.Password,
	})

	s := &HoldServer{
		Config: cfg,
	}

	// Initialize embedded PDS if database path is configured
	var xrpcHandler *pds.XRPCHandler
	if cfg.Database.Path != "" {
		holdDID := pds.GenerateDIDFromURL(cfg.Server.PublicURL)
		slog.Info("Initializing embedded PDS", "did", holdDID)

		ctx := context.Background()
		var err error
		s.PDS, err = pds.NewHoldPDS(ctx, holdDID, cfg.Server.PublicURL, cfg.Database.Path, cfg.Database.KeyPath, cfg.Registration.EnableBlueskyPosts)
		if err != nil {
			return nil, fmt.Errorf("failed to initialize embedded PDS: %w", err)
		}

		// Create storage driver from config (needed for bootstrap profile avatar)
		driver, err := factory.Create(ctx, cfg.Storage.Type(), cfg.Storage.Parameters())
		if err != nil {
			return nil, fmt.Errorf("failed to create storage driver: %w", err)
		}

		// Bootstrap PDS with captain record, hold owner as first crew member, and profile
		if err := s.PDS.Bootstrap(ctx, driver, cfg.Registration.OwnerDID, cfg.Server.Public, cfg.Registration.AllowAllCrew, cfg.Registration.ProfileAvatarURL, cfg.Registration.Region); err != nil {
			return nil, fmt.Errorf("failed to bootstrap PDS: %w", err)
		}

		// Create event broadcaster for subscribeRepos firehose
		var dbPath string
		if cfg.Database.Path != ":memory:" {
			dbPath = cfg.Database.Path + "/db.sqlite3"
		} else {
			dbPath = ":memory:"
		}
		s.broadcaster = pds.NewEventBroadcaster(holdDID, 100, dbPath)

		// Bootstrap events from existing repo records (one-time migration)
		if err := s.broadcaster.BootstrapFromRepo(s.PDS); err != nil {
			slog.Warn("Failed to bootstrap events from repo", "error", err)
		}

		// Backfill records index from existing MST data (one-time on startup)
		if err := s.PDS.BackfillRecordsIndex(ctx); err != nil {
			slog.Warn("Failed to backfill records index", "error", err)
		}

		// Wire up repo event handler with records indexing + broadcaster
		indexingHandler := s.PDS.CreateRecordsIndexEventHandler(s.broadcaster.SetRepoEventHandler())
		s.PDS.RepomgrRef().SetEventHandler(indexingHandler, true)

		slog.Info("Embedded PDS initialized successfully with firehose and records index enabled")
	} else {
		return nil, fmt.Errorf("database path is required for embedded PDS authorization")
	}

	// Initialize quota manager from config
	var err error
	s.QuotaManager, err = quota.NewManagerFromConfig(&cfg.Quota)
	if err != nil {
		return nil, fmt.Errorf("failed to load quota config: %w", err)
	}
	if s.QuotaManager.IsEnabled() {
		slog.Info("Quota enforcement enabled", "tiers", s.QuotaManager.TierCount(), "defaultTier", s.QuotaManager.GetDefaultTier())
	} else {
		slog.Info("Quota enforcement disabled (no quota tiers configured)")
	}

	// Create blob store adapter and XRPC handlers
	var ociHandler *oci.XRPCHandler
	if s.PDS != nil {
		ctx := context.Background()
		driver, err := factory.Create(ctx, cfg.Storage.Type(), cfg.Storage.Parameters())
		if err != nil {
			return nil, fmt.Errorf("failed to create storage driver: %w", err)
		}

		s3Service, err := s3.NewS3Service(cfg.Storage.Parameters())
		if err != nil {
			return nil, fmt.Errorf("failed to create S3 service: %w", err)
		}

		xrpcHandler = pds.NewXRPCHandler(s.PDS, *s3Service, driver, s.broadcaster, nil, s.QuotaManager)
		ociHandler = oci.NewXRPCHandler(s.PDS, *s3Service, driver, cfg.Registration.EnableBlueskyPosts, nil, s.QuotaManager)

		// Initialize scan broadcaster if scanner secret is configured
		if cfg.Scanner.Secret != "" {
			holdDID := pds.GenerateDIDFromURL(cfg.Server.PublicURL)
			scanDBPath := cfg.Database.Path + "/db.sqlite3"
			sb, err := pds.NewScanBroadcaster(holdDID, cfg.Server.PublicURL, cfg.Scanner.Secret, scanDBPath, driver, s.PDS)
			if err != nil {
				return nil, fmt.Errorf("failed to initialize scan broadcaster: %w", err)
			}
			s.scanBroadcaster = sb
			xrpcHandler.SetScanBroadcaster(sb)
			ociHandler.SetScanBroadcaster(sb)
			slog.Info("Scan broadcaster initialized (scanner WebSocket enabled)")
		}

		// Initialize garbage collector
		gcConfig := gc.LoadConfigFromEnv()
		s.garbageCollector = gc.NewGarbageCollector(s.PDS, driver, gcConfig)
		slog.Info("Garbage collector initialized",
			"enabled", gcConfig.Enabled,
			"dryRun", gcConfig.DryRun)
	}

	// Setup HTTP routes with chi router
	r := chi.NewRouter()
	r.Use(middleware.RealIP)
	r.Use(middleware.Logger)

	if xrpcHandler != nil {
		r.Use(xrpcHandler.CORSMiddleware())
	}

	// Root page
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		fmt.Fprintf(w, "This is a hold server. More info at https://atcr.io")
	})

	// Register XRPC/ATProto PDS endpoints
	if xrpcHandler != nil {
		slog.Info("Registering ATProto PDS endpoints")
		xrpcHandler.RegisterHandlers(r)
	}

	// Register OCI multipart upload endpoints
	if ociHandler != nil {
		slog.Info("Registering OCI multipart upload endpoints")
		ociHandler.RegisterHandlers(r)
	}

	// Initialize and register admin panel if enabled
	if cfg.Admin.Enabled && s.PDS != nil {
		adminCfg := admin.AdminConfig{
			Enabled:   true,
			PublicURL: cfg.Server.PublicURL,
		}

		s.adminUI, err = admin.NewAdminUI(context.Background(), s.PDS, s.QuotaManager, adminCfg)
		if err != nil {
			return nil, fmt.Errorf("failed to initialize admin panel: %w", err)
		}

		if s.adminUI != nil {
			slog.Info("Registering admin panel routes")
			s.adminUI.RegisterRoutes(r)
		}
	}

	// Initialize billing manager (compile-time optional via -tags billing)
	billingMgr := billing.New(s.QuotaManager, cfg.Server.PublicURL, cfg.ConfigPath())
	if billingMgr.Enabled() {
		slog.Info("Billing enabled (Stripe integration active)")
	} else {
		slog.Info("Billing disabled (not compiled or not configured)")
	}

	// Register billing endpoints (if configured and PDS available)
	if s.PDS != nil && billingMgr.Enabled() {
		billingHandler := billing.NewXRPCHandler(billingMgr, s.PDS, http.DefaultClient)
		billingHandler.RegisterHandlers(r)
	}

	s.Router = r

	return s, nil
}

// Serve starts the HTTP server and blocks until shutdown signal.
func (s *HoldServer) Serve() error {
	s.httpServer = &http.Server{
		Addr:         s.Config.Server.Addr,
		Handler:      s.Router,
		ReadTimeout:  s.Config.Server.ReadTimeout,
		WriteTimeout: s.Config.Server.WriteTimeout,
	}

	// Set up signal handling for graceful shutdown
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	// Start server in goroutine
	serverErr := make(chan error, 1)
	go func() {
		slog.Info("Starting hold service", "addr", s.Config.Server.Addr)
		if err := s.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			serverErr <- err
		}
	}()

	// Update status post to "online" after server starts
	if s.PDS != nil {
		ctx := context.Background()
		if err := s.PDS.SetStatus(ctx, "online"); err != nil {
			slog.Warn("Failed to set status post to online", "error", err)
		} else {
			slog.Info("Status post set to online")
		}
	}

	// Request crawl from relay to make PDS discoverable
	if s.Config.Server.RelayEndpoint != "" {
		slog.Info("Requesting crawl from relay", "relay", s.Config.Server.RelayEndpoint)
		if err := RequestCrawl(s.Config.Server.RelayEndpoint, s.Config.Server.PublicURL); err != nil {
			slog.Warn("Failed to request crawl from relay", "error", err)
		} else {
			slog.Info("Crawl requested successfully")
		}
	}

	// Start garbage collector (runs on startup + nightly)
	if s.garbageCollector != nil {
		s.garbageCollector.Start(context.Background())
	}

	// Wait for signal or server error
	select {
	case err := <-serverErr:
		slog.Error("Server failed", "error", err)
		logging.Shutdown()
		return err
	case sig := <-sigChan:
		slog.Info("Received signal, shutting down gracefully", "signal", sig)
		s.shutdown()
	}

	return nil
}

func (s *HoldServer) shutdown() {
	// Update status post to "offline" before shutdown
	if s.PDS != nil {
		ctx := context.Background()
		if err := s.PDS.SetStatus(ctx, "offline"); err != nil {
			slog.Warn("Failed to set status post to offline", "error", err)
		} else {
			slog.Info("Status post set to offline")
		}
	}

	// Stop garbage collector
	if s.garbageCollector != nil {
		s.garbageCollector.Stop()
		slog.Info("Garbage collector stopped")
	}

	// Close scan broadcaster database connection
	if s.scanBroadcaster != nil {
		if err := s.scanBroadcaster.Close(); err != nil {
			slog.Warn("Failed to close scan broadcaster database", "error", err)
		} else {
			slog.Info("Scan broadcaster database closed")
		}
	}

	// Close broadcaster database connection
	if s.broadcaster != nil {
		if err := s.broadcaster.Close(); err != nil {
			slog.Warn("Failed to close broadcaster database", "error", err)
		} else {
			slog.Info("Broadcaster database closed")
		}
	}

	// Close admin panel
	if s.adminUI != nil {
		if err := s.adminUI.Close(); err != nil {
			slog.Warn("Failed to close admin panel", "error", err)
		} else {
			slog.Info("Admin panel closed")
		}
	}

	// Graceful shutdown with 10 second timeout
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := s.httpServer.Shutdown(shutdownCtx); err != nil {
		slog.Error("Server shutdown error", "error", err)
	} else {
		slog.Info("Server shutdown complete")
	}

	logging.Shutdown()
}