main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/schematichq/rulesengine"
	schematicdatastreamws "github.com/schematichq/schematic-datastream-ws"
	"github.com/schematichq/schematic-go/client"
	"github.com/schematichq/schematic-go/option"
)

// Populated at build time via -ldflags="-X main.version=... -X main.commit=... -X main.buildTime=..."
var (
	version   string
	commit    string
	buildTime string
)

func valueOrUnknown(s string) string {
	if s == "" {
		return "unknown"
	}
	return s
}

// redisConfigFromEnv reads the Redis connection settings from the environment.
// Keeping all env access in main means the rest of the package takes plain
// config values.
func redisConfigFromEnv() RedisConfig {
	cfg := RedisConfig{
		ClusterMode:              strings.EqualFold(os.Getenv("REDIS_CLUSTER_MODE"), "true"),
		Addr:                     os.Getenv("REDIS_ADDR"),
		Password:                 os.Getenv("REDIS_PASSWORD"),
		TLS:                      os.Getenv("REDIS_TLS") == "true",
		MaintenanceNotifications: os.Getenv("REDIS_ENABLE_MAINTENANCE_NOTIFICATIONS") == "true",
	}
	if dbStr := os.Getenv("REDIS_DB"); dbStr != "" {
		if parsed, err := strconv.Atoi(dbStr); err == nil {
			cfg.DB = parsed
		}
	}
	if addrsStr := os.Getenv("REDIS_CLUSTER_ADDRS"); addrsStr != "" {
		addrs := strings.Split(addrsStr, ",")
		for i := range addrs {
			addrs[i] = strings.TrimSpace(addrs[i])
		}
		cfg.ClusterAddrs = addrs
	}
	return cfg
}

const (
	defaultAPIURL               = "https://api.schematichq.com"
	apiKeyEnvVar                = "SCHEMATIC_API_KEY"
	defaultCacheTTL             = 0 * time.Second // Unlimited cache by default
	defaultCacheCleanupInterval = 1 * time.Hour   // Clean up stale cache entries every hour
	defaultHealthPort           = 8090
	healthCheckTimeout          = 5 * time.Second // Timeout for the self-health-check probe
	cacheKeyPrefix              = "schematic"
	cacheKeyPrefixCompany       = "company"
	cacheKeyPrefixUser          = "user"
	cacheKeyPrefixFlags         = "flags"
)

// HealthServer provides health and readiness endpoints for container orchestration
type HealthServer struct {
	datastreamClient *schematicdatastreamws.Client
	redisClient      interface{}
	logger           *SchematicLogger
	server           *http.Server
	mu               sync.RWMutex

	// Replay introspection (optional; populated once replay components exist).
	replayCursor *ReplayCursor
	replayStats  *ReplayStats
	msgHandler   *AsyncReplicatorMessageHandler
}

// HealthStatusType represents the overall health status
type HealthStatusType string

const (
	HealthStatusHealthy   HealthStatusType = "healthy"
	HealthStatusUnhealthy HealthStatusType = "unhealthy"
)

// ReadinessStatusType represents the readiness status
type ReadinessStatusType string

const (
	ReadinessStatusReady    ReadinessStatusType = "ready"
	ReadinessStatusNotReady ReadinessStatusType = "not_ready"
)

// ComponentStatusType represents individual component status
type ComponentStatusType string

const (
	ComponentStatusConnected        ComponentStatusType = "connected"
	ComponentStatusDisconnected     ComponentStatusType = "disconnected"
	ComponentStatusReady            ComponentStatusType = "ready"
	ComponentStatusConnectedLoading ComponentStatusType = "connected_loading"
	ComponentStatusNotReady         ComponentStatusType = "not_ready"
	ComponentStatusUnknown          ComponentStatusType = "unknown"
)

// HealthStatus represents the health status response
type HealthStatus struct {
	Status       HealthStatusType               `json:"status"`
	Ready        bool                           `json:"ready"`
	Connected    bool                           `json:"connected"`
	Components   map[string]ComponentStatusType `json:"components"`
	CacheVersion string                         `json:"cache_version"`
	Timestamp    time.Time                      `json:"timestamp"`
}

// NewHealthServer creates a new health server
func NewHealthServer(port int, datastreamClient *schematicdatastreamws.Client, redisClient interface{}, logger *SchematicLogger) *HealthServer {
	hs := &HealthServer{
		datastreamClient: datastreamClient,
		redisClient:      redisClient,
		logger:           logger,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/health", hs.healthHandler)
	mux.HandleFunc("/ready", hs.readinessHandler)
	mux.HandleFunc("/debug", hs.debugHandler)

	hs.server = &http.Server{
		Addr:    fmt.Sprintf(":%d", port),
		Handler: mux,
	}

	return hs
}

// SetDatastreamClient attaches the datastream client once it exists. The server
// can be started before this (e.g. while waiting to acquire the writer lock):
// until the client is set, /health reports alive and /ready reports not-ready.
func (hs *HealthServer) SetDatastreamClient(client *schematicdatastreamws.Client) {
	hs.mu.Lock()
	hs.datastreamClient = client
	hs.mu.Unlock()
}

// SetReplayDebug wires the replay introspection sources exposed at /debug. Any
// may be nil (e.g. the cursor when replay is disabled).
func (hs *HealthServer) SetReplayDebug(cursor *ReplayCursor, stats *ReplayStats, handler *AsyncReplicatorMessageHandler) {
	hs.mu.Lock()
	hs.replayCursor = cursor
	hs.replayStats = stats
	hs.msgHandler = handler
	hs.mu.Unlock()
}

// debugReplay / debugMessages / debugResponse are the JSON shape of the /debug
// endpoint.
type debugReplay struct {
	Enabled bool   `json:"enabled"`
	Cursor  string `json:"cursor"`
	ReplayStatsSnapshot
}

type debugMessages struct {
	Processed int64 `json:"processed"`
	Applied   int64 `json:"applied"`
	Dropped   int64 `json:"dropped"`
}

type debugResponse struct {
	Replay    debugReplay   `json:"replay"`
	Messages  debugMessages `json:"messages"`
	Timestamp time.Time     `json:"timestamp"`
}

// debugHandler exposes replay/processing introspection for operators and
// integration tests: the committed cursor, reconnect/replay/reload counters, and
// message-processing totals.
func (hs *HealthServer) debugHandler(w http.ResponseWriter, r *http.Request) {
	hs.mu.RLock()
	cursor, stats, handler := hs.replayCursor, hs.replayStats, hs.msgHandler
	hs.mu.RUnlock()

	resp := debugResponse{
		Replay:    debugReplay{Enabled: cursor != nil, ReplayStatsSnapshot: stats.Snapshot()},
		Timestamp: time.Now(),
	}
	if cursor != nil {
		resp.Replay.Cursor = cursor.Get()
	}
	if handler != nil {
		p, a, d := handler.GetMetrics()
		resp.Messages = debugMessages{Processed: p, Applied: a, Dropped: d}
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		hs.logger.Error(context.Background(), fmt.Sprintf("Failed to encode debug status: %v", err))
	}
}

// Start starts the health server
func (hs *HealthServer) Start() {
	go func() {
		defer func() {
			if r := recover(); r != nil {
				hs.logger.Error(context.Background(), fmt.Sprintf("Panic in health server: %v", r))
			}
		}()
		hs.logger.Info(context.Background(), fmt.Sprintf("Health server starting on port %s", hs.server.Addr))
		if err := hs.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			hs.logger.Error(context.Background(), fmt.Sprintf("Health server error: %v", err))
		}
	}()
}

// Stop stops the health server
func (hs *HealthServer) Stop() {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := hs.server.Shutdown(ctx); err != nil {
		hs.logger.Error(context.Background(), fmt.Sprintf("Health server shutdown error: %v", err))
	}
}

// healthHandler handles the /health endpoint (liveness probe)
// Returns healthy if the process is alive and Redis is connected
func (hs *HealthServer) healthHandler(w http.ResponseWriter, r *http.Request) {
	hs.mu.RLock()
	defer hs.mu.RUnlock()

	status := HealthStatus{
		Status:    HealthStatusHealthy,
		Ready:     false,
		Connected: false,
		Components: map[string]ComponentStatusType{
			"redis":      ComponentStatusConnected, // Redis is assumed healthy if we got this far (connection tested at startup)
			"datastream": ComponentStatusUnknown,
		},
		CacheVersion: rulesengine.GetVersionKey(),
		Timestamp:    time.Now(),
	}

	// Check datastream connection (for informational purposes)
	if hs.datastreamClient != nil {
		status.Connected = hs.datastreamClient.IsConnected()
		status.Ready = hs.datastreamClient.IsReady()

		if status.Connected {
			status.Components["datastream"] = ComponentStatusConnected
		} else {
			status.Components["datastream"] = ComponentStatusDisconnected
		}
	}

	// Liveness check: Process is healthy if Redis is working
	// Datastream connectivity issues don't make the process unhealthy (it can retry)
	w.WriteHeader(http.StatusOK)
	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(status); err != nil {
		hs.logger.Error(context.Background(), fmt.Sprintf("Failed to encode health status: %v", err))
	}
}

// readinessHandler handles the /ready endpoint (readiness probe)
// Returns ready only when connected to datastream and initial data is loaded
func (hs *HealthServer) readinessHandler(w http.ResponseWriter, r *http.Request) {
	hs.mu.RLock()
	defer hs.mu.RUnlock()

	ready := false
	connected := false

	if hs.datastreamClient != nil {
		connected = hs.datastreamClient.IsConnected()
		ready = hs.datastreamClient.IsReady()
	}

	status := HealthStatus{
		Status:    HealthStatusHealthy, // Will be updated based on readiness
		Ready:     ready,
		Connected: connected,
		Components: map[string]ComponentStatusType{
			"redis":      ComponentStatusConnected,
			"datastream": ComponentStatusNotReady,
		},
		CacheVersion: rulesengine.GetVersionKey(),
		Timestamp:    time.Now(),
	}

	if ready {
		// Fully ready: connected and initial data loaded
		status.Status = HealthStatusHealthy
		status.Components["datastream"] = ComponentStatusReady
		w.WriteHeader(http.StatusOK)
	} else if connected {
		// Connected but still loading initial data
		status.Status = HealthStatusHealthy // Still healthy, just not ready yet
		status.Components["datastream"] = ComponentStatusConnectedLoading
		w.WriteHeader(http.StatusServiceUnavailable)
	} else {
		// Not connected at all
		status.Status = HealthStatusHealthy // Health endpoint should still return healthy
		status.Components["datastream"] = ComponentStatusDisconnected
		w.WriteHeader(http.StatusServiceUnavailable)
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(status); err != nil {
		hs.logger.Error(context.Background(), fmt.Sprintf("Failed to encode readiness status: %v", err))
	}
}

// healthPortFromEnv reads the health server port, falling back to the default
// when HEALTH_PORT is unset or unparseable.
func healthPortFromEnv() int {
	if portStr := os.Getenv("HEALTH_PORT"); portStr != "" {
		if parsedPort, err := strconv.Atoi(portStr); err == nil && parsedPort > 0 {
			return parsedPort
		}
	}
	return defaultHealthPort
}

// runHealthCheck probes an endpoint on the local health server and reports the
// exit code the process should use. It lets a container health-check itself, so
// the runtime image doesn't need to ship curl and its dependency chain.
func runHealthCheck(path string) int {
	url := fmt.Sprintf("http://127.0.0.1:%d%s", healthPortFromEnv(), path)

	client := &http.Client{Timeout: healthCheckTimeout}
	resp, err := client.Get(url)
	if err != nil {
		fmt.Fprintf(os.Stderr, "health check failed: %v\n", err)
		return 1
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		fmt.Fprintf(os.Stderr, "health check failed: %s returned %d\n", url, resp.StatusCode)
		return 1
	}

	return 0
}

func main() {

	// Self-check mode: `schematic-datastream-replicator healthcheck [path]`.
	// Handled before any other setup so it stays a cheap, dependency-free probe.
	if len(os.Args) > 1 && os.Args[1] == "healthcheck" {
		path := "/health"
		if len(os.Args) > 2 {
			path = os.Args[2]
		}
		os.Exit(runHealthCheck(path))
	}

	// Get API key from environment
	apiKey := os.Getenv(apiKeyEnvVar)
	if apiKey == "" {
		log.Fatalf("Please set %s environment variable", apiKeyEnvVar)
	}

	// Get API base URL from environment or use default
	apiBaseURL := os.Getenv("SCHEMATIC_API_URL")
	if apiBaseURL == "" {
		apiBaseURL = defaultAPIURL
	}

	// Get datastream WebSocket URL from environment
	// If not set, the schematic-go datastream client will automatically
	// convert the API URL to a WebSocket URL (http->ws, https->wss, +/datastream path)
	datastreamURL := os.Getenv("SCHEMATIC_DATASTREAM_URL")
	if datastreamURL == "" {
		datastreamURL = apiBaseURL // Let the datastream client handle URL conversion
	}

	// Get cache TTL from environment
	// Default is unlimited cache (0). Override with CACHE_TTL env var (e.g., "5m", "1h", "30s")
	// Set CACHE_TTL="0" or CACHE_TTL="0s" for unlimited cache explicitly
	cacheTTL := defaultCacheTTL
	if ttlStr := os.Getenv("CACHE_TTL"); ttlStr != "" {
		if parsedTTL, err := time.ParseDuration(ttlStr); err == nil {
			cacheTTL = parsedTTL
		}
	}

	// Get cache cleanup interval from environment
	// Default is 1 hour. Override with CACHE_CLEANUP_INTERVAL env var (e.g., "30m", "2h")
	// Set to "0" or "0s" to disable cache cleanup (not recommended with unlimited cache)
	cacheCleanupInterval := defaultCacheCleanupInterval
	if cleanupStr := os.Getenv("CACHE_CLEANUP_INTERVAL"); cleanupStr != "" {
		if parsedInterval, err := time.ParseDuration(cleanupStr); err == nil {
			cacheCleanupInterval = parsedInterval
		}
	}

	// Create logger
	logger := NewSchematicLogger()

	// Set log level from environment
	if logLevelStr := os.Getenv("LOG_LEVEL"); logLevelStr != "" {
		switch strings.ToLower(logLevelStr) {
		case "debug":
			logger.SetLevel(LogLevelDebug)
		case "info":
			logger.SetLevel(LogLevelInfo)
		case "warn":
			logger.SetLevel(LogLevelWarn)
		case "error":
			logger.SetLevel(LogLevelError)
		}
	}

	// Get health server port from environment
	healthPort := healthPortFromEnv()

	// Configure WebSocket options early so we can populate them as we parse environment variables
	wsOptions := schematicdatastreamws.ClientOptions{
		URL:                  datastreamURL,
		ApiKey:               apiKey,
		MaxReconnectAttempts: 10,
		MinReconnectDelay:    1 * time.Second,
		MaxReconnectDelay:    30 * time.Second,
		ExtraHeaders:         datastreamHandshakeHeaders(),
	}

	// Get WebSocket ping/pong intervals from environment (optional)
	// If not set, will use defaults from schematic-datastream-ws library (30s ping, 40s pong)
	if pingStr := os.Getenv("WS_PING_INTERVAL"); pingStr != "" {
		logger.Debug(context.Background(), "Setting ping interval from environment variable WS_PING_INTERVAL="+pingStr)
		if parsed, err := time.ParseDuration(pingStr); err == nil {
			wsOptions.PingPeriod = parsed
		} else {
			logger.Warn(context.Background(), fmt.Sprintf("Invalid WS_PING_INTERVAL value '%s', ignoring", pingStr))
		}
	}

	if pongStr := os.Getenv("WS_PONG_WAIT"); pongStr != "" {
		logger.Debug(context.Background(), "Setting pong wait from environment variable WS_PONG_WAIT="+pongStr)
		if parsed, err := time.ParseDuration(pongStr); err == nil {
			wsOptions.PongWait = parsed
		} else {
			logger.Warn(context.Background(), fmt.Sprintf("Invalid WS_PONG_WAIT value '%s', ignoring", pongStr))
		}
	}

	logger.Info(context.Background(), "Starting Schematic Datastream Replicator...")
	logger.Info(context.Background(), fmt.Sprintf("Version: %s, Commit: %s, Build time: %s", ClientVersion(), valueOrUnknown(commit), valueOrUnknown(buildTime)))
	logger.Info(context.Background(), fmt.Sprintf("API URL: %s", apiBaseURL))
	if os.Getenv("SCHEMATIC_DATASTREAM_URL") != "" {
		logger.Info(context.Background(), fmt.Sprintf("Datastream URL: %s (explicit)", datastreamURL))
	} else {
		logger.Info(context.Background(), fmt.Sprintf("Datastream URL: %s (derived from API URL)", datastreamURL))
	}
	logger.Info(context.Background(), fmt.Sprintf("Health server port: %d", healthPort))

	// Create cache providers - only Redis, no local cache fallback
	var companiesCache CacheProvider[*rulesengine.Company]
	var usersCache CacheProvider[*rulesengine.User]
	var featuresCache CacheProvider[*rulesengine.Flag]

	// Create Redis client - if this fails, the application will exit
	redisClient := setupRedisClient(redisConfigFromEnv())

	// Test Redis connection
	if err := testRedisConnection(redisClient, logger); err != nil {
		log.Fatalf("Redis connection failed: %v", err)
	}

	// Start the health server before acquiring the writer lock so probes get a
	// response during a contended startup wait (rolling deploy / lease expiry):
	// /health reports alive and /ready reports not-ready until setup completes.
	// The datastream client is attached later via SetDatastreamClient.
	healthServer := NewHealthServer(healthPort, nil, redisClient, logger)
	healthServer.Start()

	// Enforce the single-writer contract: only one instance may consume the
	// datastream and write this Redis. Wait briefly for a previous writer to
	// release (rolling deploy / lease expiry) before giving up. Disable with
	// WRITER_LOCK_DISABLED=true (e.g. for future read-only instances).
	var writerLock *WriterLock
	var lockLost chan struct{}
	var lockCancel context.CancelFunc
	if os.Getenv("WRITER_LOCK_DISABLED") != "true" {
		writerLockTTL := defaultWriterLockTTL
		if v := os.Getenv("WRITER_LOCK_TTL"); v != "" {
			if d, err := time.ParseDuration(v); err == nil && d > 0 {
				writerLockTTL = d
			}
		}
		writerLock = NewWriterLock(redisClient, logger, os.Getenv("WRITER_LOCK_KEY"), writerLockTTL)
		if err := writerLock.Acquire(context.Background(), 2*writerLock.ttl); err != nil {
			log.Fatalf("Could not acquire writer lock (another replicator instance may be running against this Redis): %v", err)
		}
		lockLost = make(chan struct{})
		var lockLostOnce sync.Once
		var lockCtx context.Context
		lockCtx, lockCancel = context.WithCancel(context.Background())
		defer lockCancel()
		go writerLock.KeepAlive(lockCtx, func() { lockLostOnce.Do(func() { close(lockLost) }) })
	}

	logger.Info(context.Background(), "Using Redis cache with batch support")
	companiesCache = NewRedisBatchCache[*rulesengine.Company](redisClient, cacheTTL)
	usersCache = NewRedisBatchCache[*rulesengine.User](redisClient, cacheTTL)
	featuresCache = NewRedisBatchCache[*rulesengine.Flag](redisClient, cacheTTL)

	// Company lookup cache stores company ID strings at versioned lookup keys
	companyLookupCache := NewRedisBatchCache[string](redisClient, cacheTTL)

	// User lookup cache stores user ID strings at versioned lookup keys
	userLookupCache := NewRedisBatchCache[string](redisClient, cacheTTL)

	// Initialize cache cleanup manager (only if cleanup is enabled)
	var cacheCleanupManager *CacheCleanupManager
	if cacheCleanupInterval > 0 {
		cacheCleanupManager = NewCacheCleanupManager(
			featuresCache,
			companiesCache,
			usersCache,
			logger,
			cacheCleanupInterval,
		)
		logger.Info(context.Background(), fmt.Sprintf("Cache cleanup enabled with interval: %v", cacheCleanupInterval))
	} else {
		logger.Info(context.Background(), "Cache cleanup disabled")
	}

	// Initialize Schematic Go client with proper configuration
	schematicClient := client.NewClient(
		option.WithAPIKey(apiKey),
		option.WithBaseURL(apiBaseURL),
	)

	// Initialize datastream WebSocket client (will be configured later)
	var datastreamClient *schematicdatastreamws.Client

	// Get async configuration from environment or use defaults
	asyncConfig := DefaultAsyncConfig()

	// Auto-detect or configure number of workers
	if numWorkersStr := os.Getenv("NUM_WORKERS"); numWorkersStr != "" {
		if numWorkers, err := strconv.Atoi(numWorkersStr); err == nil && numWorkers > 0 {
			asyncConfig.NumWorkers = numWorkers
		}
	}

	// If still 0 (default), auto-detect based on CPU cores with reasonable bounds
	if asyncConfig.NumWorkers == 0 {
		numCPU := runtime.NumCPU()
		workers := numCPU
		if workers < 2 {
			workers = 2 // Minimum for small systems
		}
		if workers > 16 {
			workers = 16 // Cap for large systems to avoid resource exhaustion
		}
		asyncConfig.NumWorkers = workers
	}

	if batchSizeStr := os.Getenv("BATCH_SIZE"); batchSizeStr != "" {
		if batchSize, err := strconv.Atoi(batchSizeStr); err == nil && batchSize > 0 {
			asyncConfig.BatchSize = batchSize
		}
	}

	if batchTimeoutStr := os.Getenv("BATCH_TIMEOUT"); batchTimeoutStr != "" {
		if batchTimeout, err := time.ParseDuration(batchTimeoutStr); err == nil {
			asyncConfig.BatchTimeout = batchTimeout
		}
	}

	// Channel size configuration for memory management
	if companyChanSizeStr := os.Getenv("COMPANY_CHANNEL_SIZE"); companyChanSizeStr != "" {
		if size, err := strconv.Atoi(companyChanSizeStr); err == nil && size > 0 {
			asyncConfig.CompanyChannelSize = size
		}
	}

	if userChanSizeStr := os.Getenv("USER_CHANNEL_SIZE"); userChanSizeStr != "" {
		if size, err := strconv.Atoi(userChanSizeStr); err == nil && size > 0 {
			asyncConfig.UserChannelSize = size
		}
	}

	if flagsChanSizeStr := os.Getenv("FLAGS_CHANNEL_SIZE"); flagsChanSizeStr != "" {
		if size, err := strconv.Atoi(flagsChanSizeStr); err == nil && size > 0 {
			asyncConfig.FlagsChannelSize = size
		}
	}

	// Circuit breaker configuration for customer environment resilience
	if cbThresholdStr := os.Getenv("CIRCUIT_BREAKER_THRESHOLD"); cbThresholdStr != "" {
		if threshold, err := strconv.Atoi(cbThresholdStr); err == nil && threshold > 0 {
			asyncConfig.CircuitBreakerThreshold = threshold
		}
	}

	if cbTimeoutStr := os.Getenv("CIRCUIT_BREAKER_TIMEOUT"); cbTimeoutStr != "" {
		if timeout, err := time.ParseDuration(cbTimeoutStr); err == nil {
			asyncConfig.CircuitBreakerTimeout = timeout
		}
	}

	logger.Info(context.Background(), fmt.Sprintf("Async processing config: workers=%d, batch_size=%d, batch_timeout=%v, channels=[company:%d, user:%d, flags:%d], circuit_breaker=[threshold:%d, timeout:%v]",
		asyncConfig.NumWorkers, asyncConfig.BatchSize, asyncConfig.BatchTimeout,
		asyncConfig.CompanyChannelSize, asyncConfig.UserChannelSize, asyncConfig.FlagsChannelSize,
		asyncConfig.CircuitBreakerThreshold, asyncConfig.CircuitBreakerTimeout))

	// Parse async loading configuration. Async is the default: it loads once and
	// resumes via replay on reconnect, avoiding the full-reload-on-every-reconnect
	// behavior of the sync path (and replay is only wired on the async path). Opt
	// into the legacy synchronous path with USE_ASYNC_LOADING=false (intended for
	// small datasets / local development).
	useAsyncLoading := true
	if asyncLoadingStr := os.Getenv("USE_ASYNC_LOADING"); asyncLoadingStr != "" {
		if asyncLoadingStr == "false" || asyncLoadingStr == "0" {
			useAsyncLoading = false
		}
	}

	// Configure async loader (only if async loading is enabled)
	var asyncLoaderConfig AsyncLoaderConfig
	if useAsyncLoading {
		asyncLoaderConfig = DefaultAsyncLoaderConfig()

		// Parse async loader specific environment variables
		if pageSize := os.Getenv("ASYNC_LOADER_PAGE_SIZE"); pageSize != "" {
			if size, err := strconv.Atoi(pageSize); err == nil && size > 0 {
				asyncLoaderConfig.PageSize = size
			}
		}
		if cbThresholdStr := os.Getenv("ASYNC_LOADER_CIRCUIT_BREAKER_THRESHOLD"); cbThresholdStr != "" {
			if threshold, err := strconv.Atoi(cbThresholdStr); err == nil && threshold > 0 {
				asyncLoaderConfig.CircuitBreakerThreshold = threshold
			}
		}
		if cbTimeoutStr := os.Getenv("ASYNC_LOADER_CIRCUIT_BREAKER_TIMEOUT"); cbTimeoutStr != "" {
			if timeout, err := time.ParseDuration(cbTimeoutStr); err == nil {
				asyncLoaderConfig.CircuitBreakerTimeout = timeout
			}
		}
		// Concurrent loading settings (always enabled in async mode)
		if maxConcurrentStr := os.Getenv("ASYNC_LOADER_MAX_CONCURRENT_REQUESTS"); maxConcurrentStr != "" {
			if maxConcurrent, err := strconv.Atoi(maxConcurrentStr); err == nil && maxConcurrent > 0 {
				asyncLoaderConfig.MaxConcurrentRequests = maxConcurrent
			}
		}
		if rateLimitStr := os.Getenv("ASYNC_LOADER_RATE_LIMIT_RPS"); rateLimitStr != "" {
			if rateLimit, err := strconv.Atoi(rateLimitStr); err == nil && rateLimit > 0 {
				asyncLoaderConfig.RateLimitRPS = rateLimit
			}
		}
	}

	// Create async message handler with caching and batching
	messageHandler := NewAsyncReplicatorMessageHandler(companiesCache, usersCache, featuresCache, companyLookupCache, userLookupCache, logger, cacheTTL, asyncConfig)

	// Counters for replay/reconnect behavior, exposed at /debug.
	replayStats := &ReplayStats{}

	// Replay cursor: lets a reconnect resume from the last processed message
	// instead of a full reload. Disable with REPLAY_DISABLED=true to revert to
	// full-reload-on-reconnect behavior — a rollout safety switch, and the
	// fallback when the datastream server doesn't support replay. When the
	// cursor is nil, the handler skips Track/Complete and never sets ReplayFrom,
	// so the server behaves exactly as it did before replay existed.
	var replayCursor *ReplayCursor
	if os.Getenv("REPLAY_DISABLED") == "true" {
		logger.Info(context.Background(), "Replay disabled (REPLAY_DISABLED=true); reconnects will do a full reload")
	} else {
		replayCursor = NewReplayCursor(redisClient, logger, os.Getenv("REPLAY_CURSOR_KEY"))
		replayCursor.Load(context.Background())
		messageHandler.SetReplayCursor(replayCursor)
	}

	// Create connection ready handler (wsClient will be set later)
	var connectionReadyHandlerFunc schematicdatastreamws.ConnectionReadyHandlerFunc
	var syncHandler *ConnectionReadyHandler
	var asyncHandler *AsyncConnectionReadyHandler

	if useAsyncLoading {
		logger.Info(context.Background(), fmt.Sprintf("Using async initial loading: page_size=%d, circuit_breaker=[threshold:%d, timeout:%v], concurrency=[max_requests:%d, rate_limit:%d_rps]",
			asyncLoaderConfig.PageSize, asyncLoaderConfig.CircuitBreakerThreshold, asyncLoaderConfig.CircuitBreakerTimeout,
			asyncLoaderConfig.MaxConcurrentRequests, asyncLoaderConfig.RateLimitRPS))
		asyncHandler = NewAsyncConnectionReadyHandler(schematicClient, nil, companiesCache, usersCache, featuresCache, companyLookupCache, userLookupCache, logger, cacheTTL, asyncLoaderConfig)
		asyncHandler.SetStats(replayStats)
		if replayCursor != nil {
			asyncHandler.SetReplayCursor(replayCursor)
			messageHandler.SetReloadFunc(asyncHandler.TriggerReload)
			// If the cursor's in-flight set overflows (a gap that narrow replay can
			// no longer close), fall back to a full reload.
			replayCursor.SetOverflowHandler(func() {
				replayStats.IncOverflowEscalations()
				asyncHandler.TriggerReload(context.Background())
			})
		}
		connectionReadyHandlerFunc = asyncHandler.OnConnectionReady
	} else {
		logger.Info(context.Background(), "Using synchronous initial loading (USE_ASYNC_LOADING=false); replay is disabled on this path")
		syncHandler = NewConnectionReadyHandler(schematicClient, nil, companiesCache, usersCache, featuresCache, companyLookupCache, userLookupCache, logger, cacheTTL)
		connectionReadyHandlerFunc = syncHandler.OnConnectionReady
	}

	// Set message and connection handlers on wsOptions
	wsOptions.MessageHandler = messageHandler.HandleMessage
	wsOptions.ConnectionReadyHandler = connectionReadyHandlerFunc
	wsOptions.Logger = logger

	datastreamClient, err := schematicdatastreamws.NewClient(wsOptions)
	if err != nil {
		log.Fatalf("Failed to create WebSocket client: %v", err)
	}

	// Set the WebSocket client in the connection ready handler
	if asyncHandler != nil {
		asyncHandler.SetWebSocketClient(datastreamClient)
	} else if syncHandler != nil {
		syncHandler.SetWebSocketClient(datastreamClient)
	}

	// Attach the datastream client to the already-running health server so
	// readiness now reflects datastream connectivity and initial load, and wire
	// the replay introspection exposed at /debug.
	healthServer.SetDatastreamClient(datastreamClient)
	healthServer.SetReplayDebug(replayCursor, replayStats, messageHandler)

	// Start the WebSocket connection
	datastreamClient.Start()

	// Periodically persist the replay cursor so a process restart can resume.
	cursorCtx, cursorCancel := context.WithCancel(context.Background())
	defer cursorCancel()
	if replayCursor != nil {
		replayCursor.StartFlusher(cursorCtx, 5*time.Second)
	}

	// Start cache cleanup manager if enabled
	if cacheCleanupManager != nil {
		cacheCleanupManager.Start(context.Background())
	}

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

	// Monitor for errors
	go func() {
		defer func() {
			if r := recover(); r != nil {
				logger.Error(context.Background(), fmt.Sprintf("Panic in error monitor: %v", r))
			}
		}()
		errorChan := datastreamClient.GetErrorChannel()
		for err := range errorChan {
			logger.Error(context.Background(), fmt.Sprintf("WebSocket error: %v", err))
		}
	}()

	// Monitor async handler metrics
	go func() {
		defer func() {
			if r := recover(); r != nil {
				logger.Error(context.Background(), fmt.Sprintf("Panic in metrics monitor: %v", r))
			}
		}()
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				processed, applied, dropped := messageHandler.GetMetrics()
				if processed > 0 || dropped > 0 {
					logger.Info(context.Background(), fmt.Sprintf("Message metrics: processed=%d, applied=%d, dropped=%d", processed, applied, dropped))
				}
			case <-sigChan:
				return
			}
		}
	}()

	// Wait for a shutdown signal, or for the writer lock to be lost (another
	// instance took over) — in which case we must stop to preserve the
	// single-writer guarantee.
	select {
	case <-sigChan:
		logger.Info(context.Background(), "Received shutdown signal, closing connection...")
	case <-lockLost:
		logger.Error(context.Background(), "Lost writer lock; shutting down to preserve the single-writer guarantee")
	}

	// Stop the cursor flusher (flushes once more) so the latest position is
	// durable across the restart.
	cursorCancel()

	// Stop cache cleanup manager
	if cacheCleanupManager != nil {
		cacheCleanupManager.Stop()
		logger.Info(context.Background(), "Cache cleanup manager stopped")
	}

	// Stop async message handler
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer shutdownCancel()

	if err := messageHandler.Shutdown(shutdownCtx); err != nil {
		logger.Error(context.Background(), fmt.Sprintf("Error shutting down async handler: %v", err))
	} else {
		logger.Info(context.Background(), "Async message handler stopped")
	}

	// Stop health server
	healthServer.Stop()

	// Close the WebSocket connection
	datastreamClient.Close()

	// Release the writer lock last, after we've stopped consuming and writing,
	// so a replacement instance doesn't start writing while we're still flushing.
	// Stop the renewer first so it can't re-extend the lease after we release.
	if writerLock != nil {
		lockCancel()
		writerLock.Release(context.Background())
	}

	logger.Info(context.Background(), "Datastream replicator stopped")
}