package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"net/http"
	"net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/everything-on-mysql/cassandra-on-mysql/internal/config"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/tenancy"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/tenantmanager"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/logger"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/metrics"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/mysqlbackend"
	"go.uber.org/zap"
)

const (
	drainPollInterval    = 250 * time.Millisecond
	adminShutdownTimeout = 5 * time.Second
	warmupProbeInterval  = 1 * time.Second
	warmupProbeTimeout   = 5 * time.Second
)

func main() {
	configPath := flag.String("config", "config.yaml", "Path to config file")
	advertisedAddress := flag.String("advertised-address", "", "Override system_tables.advertised_address (or env SYSTEM_TABLES_ADVERTISED_ADDRESS)")
	flag.Parse()

	// Load configuration first
	cfg, err := config.LoadConfig(*configPath)
	if err != nil {
		// If config fails, use defaults and create a basic logger
		cfg = config.DefaultConfig()
		fmt.Printf("Warning: Failed to load config %q, using defaults: %v\n", *configPath, err)
	}
	if *advertisedAddress != "" {
		cfg.SystemTables.AdvertisedAddress = *advertisedAddress
	} else if value, ok := os.LookupEnv("SYSTEM_TABLES_ADVERTISED_ADDRESS"); ok && value != "" {
		cfg.SystemTables.AdvertisedAddress = value
	}
	config.SetConfig(cfg)

	// Initialize logger with configuration
	log, err := logger.NewLogger(&cfg.Logging)
	if err != nil {
		panic(fmt.Sprintf("Failed to initialize logger: %v", err))
	}

	// Get the underlying zap.Logger for components that need it
	zapLog := log.GetZapLogger()

	poolCfg := &mysqlbackend.PoolConfig{
		Host:            cfg.MySQL.Host,
		MaxOpenConns:    cfg.MySQL.MaxOpenConns,
		MaxIdleConns:    cfg.MySQL.MaxIdleConns,
		ConnectTimeout:  cfg.MySQL.ConnectTimeout,
		ReadTimeout:     cfg.MySQL.ReadTimeout,
		WriteTimeout:    cfg.MySQL.WriteTimeout,
		ConnMaxLifetime: cfg.MySQL.ConnMaxLifetime,
		ConnMaxIdleTime: cfg.MySQL.ConnMaxIdleTime,
		DSNParams:       cfg.MySQL.DSNParams,
		PreparedStatements: mysqlbackend.PreparedStatementsConfig{
			Enabled:          cfg.MySQL.Prepared.Enabled,
			MaxStatements:    cfg.MySQL.Prepared.MaxStatements,
			IdleTTL:          cfg.MySQL.Prepared.IdleTTL,
			CleanupInterval:  cfg.MySQL.Prepared.CleanupInterval,
			PrewarmOnPrepare: cfg.MySQL.Prepared.PrewarmOnPrepare,
		},
		TLS: mysqlbackend.TLSConfig{
			Mode:               cfg.MySQL.TLS.Mode,
			CAFile:             cfg.MySQL.TLS.CAFile,
			CertFile:           cfg.MySQL.TLS.CertFile,
			KeyFile:            cfg.MySQL.TLS.KeyFile,
			ServerName:         cfg.MySQL.TLS.ServerName,
			InsecureSkipVerify: cfg.MySQL.TLS.InsecureSkipVerify,
		},
	}

	resolver, err := tenancy.NewResolver(cfg.Tenancy.TenantRegex, cfg.Tenancy.TenantRegexFallback)
	if err != nil {
		log.Fatal("Invalid tenant configuration", zap.Error(err))
	}
	tenantManager, err := tenantmanager.NewTenantManager(cfg.Tenancy, poolCfg, zapLog)
	if err != nil {
		log.Fatal("Failed to initialize tenant manager", zap.Error(err))
	}
	log.Info("Created tenant manager", zap.String("mysql", cfg.MySQL.Host))

	metrics.Init(metrics.Settings{
		Enabled:              cfg.Server.Admin.Metrics.Enabled,
		SlowRequestThreshold: cfg.Server.Admin.Metrics.SlowRequestThreshold,
	}, func() []metrics.PoolStat {
		stats := tenantManager.PoolStats()
		if len(stats) == 0 {
			return nil
		}
		out := make([]metrics.PoolStat, 0)
		for _, tenantStats := range stats {
			for _, pool := range tenantStats.Pools {
				if pool == nil {
					continue
				}
				out = append(out, metrics.PoolStat{
					TenantID:     tenantStats.TenantID,
					User:         pool.Username,
					Open:         pool.OpenConns,
					InUse:        pool.InUseConns,
					WaitCount:    pool.WaitCount,
					WaitDuration: time.Duration(pool.WaitDuration) * time.Millisecond,
				})
			}
		}
		return out
	})

	var pruneStop chan struct{}
	if cfg.Tenancy.PruneInterval > 0 {
		pruneStop = make(chan struct{})
		ticker := time.NewTicker(cfg.Tenancy.PruneInterval)
		go func() {
			defer ticker.Stop()
			for {
				select {
				case <-ticker.C:
					pruned := tenantManager.PruneIdleTenants()
					stats := tenantManager.Stats()
					fields := []zap.Field{
						zap.Int("tenants_pruned", pruned),
						zap.Int("tenants_active", stats.Tenants),
						zap.Int("proxy_in_use", stats.ProxyInUse),
						zap.Int("mysql_in_use", stats.MySQLInUse),
					}
					if pruned > 0 {
						log.Info("Pruned idle tenants", fields...)
					} else {
						log.Debug("Tenant activity snapshot", fields...)
					}
				case <-pruneStop:
					return
				}
			}
		}()
	}

	// Create password authenticator - validates credentials are non-empty
	// Actual credential validation happens when connecting to MySQL backend
	authenticator := cql.NewPasswordAuthenticator()

	// Create message handler with query executor and pool manager
	handler := cql.NewMessageHandlerWithTenancy(zapLog, authenticator, tenantManager, resolver)

	// Create CQL server
	listeners := make([]cql.ListenerConfig, 0, len(cfg.Server.Listeners))
	for i, listener := range cfg.Server.Listeners {
		name := listener.Name
		if name == "" {
			name = fmt.Sprintf("listener-%d", i)
		}
		listeners = append(listeners, cql.ListenerConfig{
			Name: name,
			Addr: listener.Address,
			TLS: &cql.TLSConfig{
				Mode:         listener.TLS.Mode,
				CertFile:     listener.TLS.CertFile,
				KeyFile:      listener.TLS.KeyFile,
				ClientCAFile: listener.TLS.ClientCAFile,
				ClientAuth:   listener.TLS.ClientAuth,
				MinVersion:   listener.TLS.MinVersion,
				MaxVersion:   listener.TLS.MaxVersion,
			},
		})
	}
	cqlServer := cql.NewServer(
		"",
		handler,
		zapLog,
		&cql.ServerOptions{
			ProtocolVersion:   cql.ProtoVersion4,
			Listeners:         listeners,
			MaxConnections:    cfg.Server.MaxConnections,
			MaxInFlightFrames: cfg.Server.MaxInFlightFrames,
			FrameInfo:         cfg.Logging.FrameInfo,
		},
	)

	// Start server
	addresses := make([]string, 0, len(cfg.Server.Listeners))
	for i, listener := range cfg.Server.Listeners {
		name := listener.Name
		if name == "" {
			name = fmt.Sprintf("listener-%d", i)
		}
		addresses = append(addresses, fmt.Sprintf("%s=%s", name, listener.Address))
	}
	log.Info("Starting Cassandra Proxy Server",
		zap.Strings("addresses", addresses),
		zap.String("mysql", cfg.MySQL.Host))

	if err := cqlServer.Start(); err != nil {
		log.Fatal("Failed to start server", zap.Error(err))
	}

	// readiness starts not-ready: /ready returns 503 until warmup verifies the
	// MySQL backend is reachable, so the load balancer does not route traffic to
	// a pod whose backend it cannot reach yet.
	rdy := &readiness{}

	drain := newDrainController(rdy, cfg.Server.Drain.ReadinessDelay, func() int {
		return cqlServer.BeginDrain()
	}, zapLog)

	// warmupCtx bounds the warmup probe loop. Cancelling it (on drain or signal)
	// stops warmup so a draining pod is never marked ready afterwards.
	warmupCtx, warmupCancel := context.WithCancel(context.Background())
	defer warmupCancel()

	var adminServer *http.Server

	// Start admin HTTP server with pprof endpoints if configured
	if cfg.Server.Admin.Address != "" {
		mux := http.NewServeMux()
		mux.HandleFunc("/ready", func(w http.ResponseWriter, _ *http.Request) {
			if !rdy.Ready() {
				http.Error(w, "not ready", http.StatusServiceUnavailable)
				return
			}
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ok\n"))
		})
		mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ok\n"))
		})
		mux.HandleFunc("/drain", func(w http.ResponseWriter, r *http.Request) {
			if r.Method != http.MethodGet && r.Method != http.MethodPost {
				w.Header().Set("Allow", "GET, POST")
				w.WriteHeader(http.StatusMethodNotAllowed)
				return
			}
			// Mark draining (synchronous) before cancelling warmup so a racing
			// warmup completion observes the drain and cannot mark this pod ready.
			drain.Enter("admin endpoint")
			warmupCancel()
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("draining\n"))
		})

		if cfg.Server.Admin.PprofEnabled() {
			pprofBase := cfg.Server.Admin.PprofPath()
			mux.HandleFunc(pprofBase+"/", pprof.Index)
			mux.HandleFunc(pprofBase+"/cmdline", pprof.Cmdline)
			mux.HandleFunc(pprofBase+"/profile", pprof.Profile)
			mux.HandleFunc(pprofBase+"/symbol", pprof.Symbol)
			mux.HandleFunc(pprofBase+"/trace", pprof.Trace)
		}
		if cfg.Server.Admin.Metrics.Enabled {
			path := cfg.Server.Admin.Metrics.Path
			if path == "" {
				path = "/metrics"
			}
			if metricsHandler := metrics.Handler(); metricsHandler != nil {
				// Return 503 only while draining so probe/NLB can remove this pod
				// quickly. Metrics stay available during warmup so a pod stuck
				// warming up remains observable.
				mux.Handle(path, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					if rdy.Draining() {
						http.Error(w, "draining", http.StatusServiceUnavailable)
						return
					}
					metricsHandler.ServeHTTP(w, r)
				}))
			}
		}

		adminServer = &http.Server{Addr: cfg.Server.Admin.Address, Handler: mux}
		go func() {
			log.Info("Starting admin HTTP server",
				zap.String("address", cfg.Server.Admin.Address))
			if err := adminServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
				log.Error("Admin HTTP server failed", zap.Error(err))
			}
		}()
	}

	// Warm up the backend before advertising readiness: /ready stays 503 until a
	// TCP probe confirms the MySQL backend (TiProxy) is reachable, so the load
	// balancer does not route traffic to a pod whose backend it cannot reach.
	go func() {
		probe := tcpReachableProbe(cfg.MySQL.Host, warmupProbeTimeout)
		if err := warmUp(warmupCtx, probe, warmupProbeInterval, zapLog); err != nil {
			log.Warn("Warmup aborted before completion", zap.Error(err))
			return
		}
		if !rdy.MarkWarm() {
			log.Info("Warmup completed after drain started; staying not ready")
			return
		}
		log.Info("Warmup complete, proxy marked ready", zap.String("mysql", cfg.MySQL.Host))
	}()

	// Wait for interrupt signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	sig := <-sigChan

	log.Info("Shutting down server...", zap.String("signal", sig.String()))
	// Mark draining (synchronous) before cancelling warmup so a racing warmup
	// completion observes the drain and cannot mark this pod ready.
	drain.Enter("signal")
	warmupCancel()
	drain.WaitForBeginDrain()

	if pruneStop != nil {
		close(pruneStop)
	}

	remaining := waitForConnectionsToDrain(cqlServer, cfg.Server.Drain.WaitTimeout)
	if remaining > 0 {
		log.Warn("Drain timeout reached, force closing remaining connections",
			zap.Int("remaining_connections", remaining))
		closed := cqlServer.ForceCloseConnections()
		log.Info("Forced connection close completed", zap.Int("connections", closed))
	}

	// Stop server
	if err := cqlServer.Stop(); err != nil {
		log.Error("Error stopping server", zap.Error(err))
	}

	// Stop admin HTTP server
	if adminServer != nil {
		ctx, cancel := context.WithTimeout(context.Background(), adminShutdownTimeout)
		if err := adminServer.Shutdown(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Error("Error shutting down admin server", zap.Error(err))
		}
		cancel()
	}

	// Close connection pool manager
	if err := tenantManager.Close(); err != nil {
		log.Error("Error closing tenant manager", zap.Error(err))
	}

	log.Info("Server stopped")
}

func waitForConnectionsToDrain(server *cql.Server, timeout time.Duration) int {
	if server == nil {
		return 0
	}
	if timeout <= 0 {
		return server.ActiveConnectionCount()
	}
	deadline := time.Now().Add(timeout)
	for {
		remaining := server.ActiveConnectionCount()
		if remaining == 0 {
			return 0
		}
		if time.Now().After(deadline) {
			return remaining
		}
		time.Sleep(drainPollInterval)
	}
}