package cql

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/types"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/metrics"
	"go.uber.org/zap"
)

const (
	// connDrainStaggerStep spaces out per-connection drain starts so clients do
	// not all reconnect at the same instant (avoids a reconnect storm).
	connDrainStaggerStep = 25 * time.Millisecond
	// connDrainStaggerBuckets caps the total spread (~1s) so draining still
	// completes well within the server-wide wait_timeout.
	connDrainStaggerBuckets = 40
)

// connDrainStagger returns the delay before connection index begins draining.
func connDrainStagger(index int) time.Duration {
	if index <= 0 {
		return 0
	}
	return time.Duration(index%connDrainStaggerBuckets) * connDrainStaggerStep
}

// Server represents a CQL protocol server
type Server struct {
	addr        string
	listener    net.Listener
	listeners   []net.Listener
	closeOnce   sync.Once
	logger      *zap.Logger
	handler     Handler
	connWG      sync.WaitGroup
	shutdownCh  chan struct{}
	mu          sync.Mutex
	connections map[string]*Conn // Map of connection IDs to connections
	connLimiter chan struct{}
	draining    atomic.Bool
	options     *ServerOptions
}

// ServerOptions configures the CQL server
type ServerOptions struct {
	// Protocol version to support (4 or 5)
	ProtocolVersion byte
	// Authentication provider
	Authenticator Authenticator
	// TLS configuration for the listener.
	TLS *TLSConfig
	// Listeners configures multiple listening addresses with independent TLS.
	Listeners []ListenerConfig
	// MaxConnections enforces a global connection limit when > 0.
	MaxConnections int
	// MaxInFlightFrames sets the per-connection in-flight frame cap (clamped to [1, MaxStreamIDv3p]).
	MaxInFlightFrames int
	// FrameInfo enables per-frame info logging on the write path.
	FrameInfo bool
}

// ListenerConfig configures a single listening address and its TLS settings.
type ListenerConfig struct {
	Name string
	Addr string
	TLS  *TLSConfig
}

// Authenticator interface for handling authentication
type Authenticator interface {
	Authenticate(username, password string) (string, error) // Returns MySQL connection string
	RequiresAuth() bool
}

// Handler interface for handling CQL messages
type Handler interface {
	HandleMessage(ctx context.Context, frame *Frame) (*Frame, error)
}

// NewServer creates a new CQL server
func NewServer(addr string, handler Handler, logger *zap.Logger, options *ServerOptions) *Server {
	if options == nil {
		options = &ServerOptions{
			ProtocolVersion: ProtoVersion4,
		}
	}
	if logger == nil {
		logger = zap.NewNop()
	}

	server := &Server{
		addr:        addr,
		handler:     handler,
		logger:      logger,
		shutdownCh:  make(chan struct{}),
		connections: make(map[string]*Conn),
		options:     options,
	}
	metrics.SetConnectionStatsProvider(server.MetricConnectionStats)
	return server
}

// Start starts the CQL server and begins accepting connections
func (s *Server) Start() error {
	listenerConfigs, err := resolveListeners(s.addr, s.options)
	if err != nil {
		return err
	}
	if s.options != nil && s.options.MaxConnections > 0 && s.connLimiter == nil {
		s.connLimiter = make(chan struct{}, s.options.MaxConnections)
	}

	listeners := make([]net.Listener, 0, len(listenerConfigs))
	for _, cfg := range listenerConfigs {
		listener, err := net.Listen("tcp", cfg.Addr)
		if err != nil {
			for _, existing := range listeners {
				_ = existing.Close()
			}
			return fmt.Errorf("failed to listen on %s: %w", cfg.Addr, err)
		}
		var tlsEnabled bool
		if cfg.TLS != nil {
			tlsConfig, err := cfg.TLS.BuildTLSConfig()
			if err != nil {
				_ = listener.Close()
				for _, existing := range listeners {
					_ = existing.Close()
				}
				return fmt.Errorf("failed to configure tls for %s: %w", cfg.Addr, err)
			}
			if tlsConfig != nil {
				listener = tls.NewListener(listener, tlsConfig)
				tlsEnabled = true
			}
		}
		listeners = append(listeners, listener)
		s.logger.Info("CQL listener started", zap.String("listener", cfg.Name), zap.String("address", cfg.Addr), zap.Bool("tls", tlsEnabled))
		go s.acceptLoop(listener, cfg.Name)
	}

	s.listeners = listeners
	if len(listeners) > 0 {
		s.listener = listeners[0]
	}

	s.logger.Info("CQL server started", zap.Int("listeners", len(listeners)))
	return nil
}

// Stop gracefully shuts down the server
func (s *Server) Stop() error {
	close(s.shutdownCh)
	s.closeListeners()
	s.ForceCloseConnections()

	// Wait for all connection handlers to finish
	s.connWG.Wait()
	s.logger.Info("CQL server stopped")
	return nil
}

// BroadcastSchemaChange sends a schema change EVENT frame to subscribed connections.
func (s *Server) BroadcastSchemaChange(change *types.SchemaChange) {
	s.BroadcastSchemaChangeForTenant(change, "")
}

// BroadcastSchemaChangeForTenant sends a schema change EVENT frame to subscribed tenant connections.
func (s *Server) BroadcastSchemaChangeForTenant(change *types.SchemaChange, tenantID string) {
	if change == nil {
		return
	}
	body, err := buildSchemaChangeEventBody(change)
	if err != nil {
		s.logger.Error("failed to build schema change event", zap.Error(err))
		return
	}

	s.mu.Lock()
	conns := make([]*Conn, 0, len(s.connections))
	for _, conn := range s.connections {
		conns = append(conns, conn)
	}
	s.mu.Unlock()

	for _, conn := range conns {
		if tenantID != "" && conn.GetTenantID() != tenantID {
			continue
		}
		if !conn.HasEventSubscription(schemaChangeEventType) {
			continue
		}
		event := Frame{
			Version: ResponseVersion(conn.GetProtocolVersion()),
			Flags:   0,
			Stream:  -1,
			Opcode:  OpcodeEvent,
			Length:  int32(len(body)),
			Body:    body,
		}
		conn.enqueueEvent(&event)
	}
}

// acceptLoop accepts incoming connections
func (s *Server) acceptLoop(listener net.Listener, listenerName string) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			select {
			case <-s.shutdownCh:
				return
			default:
				if s.IsDraining() || isListenerClosedError(err) {
					return
				}
				s.logger.Error("failed to accept connection", zap.String("listener", listenerName), zap.Error(err))
				continue
			}
		}

		if !s.acquireConn() {
			s.logger.Warn("connection limit reached", zap.String("listener", listenerName), zap.String("remote", conn.RemoteAddr().String()))
			_ = conn.Close()
			continue
		}

		s.connWG.Add(1)
		go s.handleConnection(conn, listenerName)
	}
}

func (s *Server) acquireConn() bool {
	if s == nil || s.connLimiter == nil {
		return true
	}
	select {
	case s.connLimiter <- struct{}{}:
		return true
	default:
		return false
	}
}

func (s *Server) releaseConn() {
	if s == nil || s.connLimiter == nil {
		return
	}
	select {
	case <-s.connLimiter:
	default:
	}
}

// IsDraining reports whether the server is currently draining.
func (s *Server) IsDraining() bool {
	if s == nil {
		return false
	}
	return s.draining.Load()
}

// BeginDrain stops accepting new connections and asks existing connections to
// close after they finish the current response.
func (s *Server) BeginDrain() int {
	if s == nil {
		return 0
	}

	s.draining.Store(true)
	s.closeListeners()

	conns := s.snapshotConnections()
	for i, conn := range conns {
		if conn == nil {
			continue
		}
		delay := connDrainStagger(i)
		if delay <= 0 {
			conn.beginDrain()
			continue
		}
		go func(conn *Conn, delay time.Duration) {
			timer := time.NewTimer(delay)
			defer timer.Stop()
			<-timer.C
			conn.beginDrain()
		}(conn, delay)
	}

	active := len(conns)
	s.logger.Info("CQL server entered drain mode", zap.Int("active_connections", active))
	return active
}

// ActiveConnectionCount returns the number of active client connections.
func (s *Server) ActiveConnectionCount() int {
	if s == nil {
		return 0
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.connections)
}

// ForceCloseConnections closes all active client connections immediately.
func (s *Server) ForceCloseConnections() int {
	if s == nil {
		return 0
	}
	conns := s.snapshotConnections()
	for _, conn := range conns {
		if conn == nil {
			continue
		}
		_ = conn.Close()
	}
	return len(conns)
}

func (s *Server) snapshotConnections() []*Conn {
	s.mu.Lock()
	defer s.mu.Unlock()
	conns := make([]*Conn, 0, len(s.connections))
	for _, conn := range s.connections {
		conns = append(conns, conn)
	}
	return conns
}

func (s *Server) closeListeners() {
	if s == nil {
		return
	}
	s.closeOnce.Do(func() {
		if s.listener != nil {
			if err := s.listener.Close(); err != nil && !isListenerClosedError(err) {
				s.logger.Error("failed to close listener", zap.Error(err))
			}
		}
		for _, listener := range s.listeners {
			if listener == nil || listener == s.listener {
				continue
			}
			if err := listener.Close(); err != nil && !isListenerClosedError(err) {
				s.logger.Error("failed to close listener", zap.Error(err))
			}
		}
	})
}

func isListenerClosedError(err error) bool {
	if err == nil {
		return false
	}
	return errors.Is(err, net.ErrClosed) || strings.Contains(err.Error(), "use of closed network connection")
}

// handleConnection handles a single client connection
func (s *Server) handleConnection(conn net.Conn, listenerName string) {
	defer s.connWG.Done()
	defer s.releaseConn()

	var tlsState *tls.ConnectionState
	if tlsConn, ok := conn.(*tls.Conn); ok {
		if err := tlsConn.Handshake(); err != nil {
			s.logger.Error("TLS handshake failed", zap.String("listener", listenerName), zap.String("remote", conn.RemoteAddr().String()), zap.Error(err))
			_ = conn.Close()
			return
		}
		state := tlsConn.ConnectionState()
		tlsState = &state
	}

	c := NewConn(conn, s)
	if listenerName != "" {
		c.logger = c.logger.With(zap.String("listener", listenerName))
	}
	if tlsState != nil {
		c.SetTLSState(*tlsState)
	}
	connID := conn.RemoteAddr().String()

	// Add connection to map
	s.mu.Lock()
	s.connections[connID] = c
	s.mu.Unlock()

	defer func() {
		// Remove connection from map
		s.mu.Lock()
		delete(s.connections, connID)
		s.mu.Unlock()
		metrics.RecordCQLDisconnection(c.GetUsername())
		c.Close()
	}()

	s.logger.Info("new connection", zap.String("listener", listenerName), zap.String("remote", conn.RemoteAddr().String()))

	// Handle connection
	if err := c.Serve(); err != nil {
		s.logger.Error("connection error", zap.String("listener", listenerName), zap.String("remote", conn.RemoteAddr().String()), zap.Error(err))
	}

	s.logger.Info("connection closed", zap.String("listener", listenerName), zap.String("remote", conn.RemoteAddr().String()))
}

// MetricConnectionStats returns a metrics-friendly snapshot for active CQL client connections.
func (s *Server) MetricConnectionStats() []metrics.ConnectionStat {
	if s == nil {
		return nil
	}
	type statKey struct {
		scope        string
		keyspaceName string
		clusterID    string
	}

	counts := make(map[statKey]int)
	for _, conn := range s.snapshotConnections() {
		if conn == nil {
			continue
		}
		labels := metrics.ResolveMetricLabels(context.Background(), conn.GetUsername())
		allKey := statKey{scope: "all", keyspaceName: labels.KeyspaceName, clusterID: labels.ClusterID}
		counts[allKey]++
		if conn.IsStartupComplete() {
			activeKey := statKey{scope: "active", keyspaceName: labels.KeyspaceName, clusterID: labels.ClusterID}
			counts[activeKey]++
		}
	}

	stats := make([]metrics.ConnectionStat, 0, len(counts))
	for key, count := range counts {
		stats = append(stats, metrics.ConnectionStat{
			Scope:        key.scope,
			KeyspaceName: key.keyspaceName,
			ClusterID:    key.clusterID,
			Count:        count,
		})
	}
	return stats
}

// Addr returns the server's listening address
func (s *Server) Addr() net.Addr {
	if len(s.listeners) > 0 && s.listeners[0] != nil {
		return s.listeners[0].Addr()
	}
	if s.listener != nil {
		return s.listener.Addr()
	}
	return nil
}

// Addrs returns all active listener addresses.
func (s *Server) Addrs() []net.Addr {
	if len(s.listeners) == 0 {
		if s.listener != nil {
			return []net.Addr{s.listener.Addr()}
		}
		return nil
	}
	addrs := make([]net.Addr, 0, len(s.listeners))
	for _, listener := range s.listeners {
		if listener != nil {
			addrs = append(addrs, listener.Addr())
		}
	}
	return addrs
}

func resolveListeners(addr string, opts *ServerOptions) ([]ListenerConfig, error) {
	if opts != nil && len(opts.Listeners) > 0 {
		listeners := make([]ListenerConfig, 0, len(opts.Listeners))
		for i, listener := range opts.Listeners {
			listenAddr := strings.TrimSpace(listener.Addr)
			if listenAddr == "" {
				return nil, fmt.Errorf("listener[%d] address is required", i)
			}
			name := strings.TrimSpace(listener.Name)
			if name == "" {
				name = fmt.Sprintf("listener-%d", i)
			}
			listeners = append(listeners, ListenerConfig{
				Name: name,
				Addr: listenAddr,
				TLS:  listener.TLS,
			})
		}
		return listeners, nil
	}

	listenAddr := strings.TrimSpace(addr)
	if listenAddr == "" {
		return nil, fmt.Errorf("listen address is required")
	}
	var tlsConfig *TLSConfig
	if opts != nil {
		tlsConfig = opts.TLS
	}
	return []ListenerConfig{{Name: "default", Addr: listenAddr, TLS: tlsConfig}}, nil
}

// ConnectionInfo represents information about an active connection
type ConnectionInfo struct {
	Address         string
	Port            int
	ProtocolVersion int
	RequestCount    int64
	Username        string
	TenantID        string
	SSL             bool
	SSLProtocol     string
	SSLCipherSuite  string
	DriverName      string
	DriverVersion   string
	Hostname        string
	Keyspace        string
	ClientOptions   map[string]string
}

// GetActiveConnections returns information about all active connections
func (s *Server) GetActiveConnections() []ConnectionInfo {
	s.mu.Lock()
	defer s.mu.Unlock()

	connections := make([]ConnectionInfo, 0, len(s.connections))
	for connID, conn := range s.connections {
		// Parse connection ID to get address and port
		// Format: "ip:port"
		addr := connID
		port := 0
		if idx := strings.LastIndex(connID, ":"); idx >= 0 {
			addr = connID[:idx]
			// Try to parse port
			if portStr := connID[idx+1:]; portStr != "" {
				if p, err := strconv.Atoi(portStr); err == nil {
					port = p
				}
			}
		}

		// Get connection details
		username := conn.GetUsername()
		protocolVersion := int(conn.GetProtocolVersion())
		sslEnabled, sslProtocol, sslCipherSuite := conn.TLSInfo()
		options := conn.GetStartupOptions()
		driverName := "Unknown"
		driverVersion := "Unknown"
		if len(options) > 0 {
			if value := options["DRIVER_NAME"]; value != "" {
				driverName = value
			}
			if value := options["DRIVER_VERSION"]; value != "" {
				driverVersion = value
			}
		}

		// For now, use mock data for fields we can't easily get
		// In a real implementation, we'd track these in the Conn struct
		connections = append(connections, ConnectionInfo{
			Address:         addr,
			Port:            port,
			ProtocolVersion: protocolVersion,
			RequestCount:    conn.GetRequestCount(),
			Username:        username,
			TenantID:        conn.GetTenantID(),
			SSL:             sslEnabled,
			SSLProtocol:     sslProtocol,
			SSLCipherSuite:  sslCipherSuite,
			DriverName:      driverName,
			DriverVersion:   driverVersion,
			Hostname:        addr,
			Keyspace:        conn.GetKeyspace(),
			ClientOptions:   options,
		})
	}
	return connections
}