package mysqlbackend

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"fmt"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	mysql "github.com/go-sql-driver/mysql"
	"go.uber.org/zap"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/metrics"
)

// PoolManager manages per-user connection pools
type PoolManager struct {
	mu       sync.RWMutex
	pools    map[string]*UserPool
	logger   *zap.Logger
	cfg      *PoolConfig
	tlsOnce  sync.Once
	tlsParam string
	tlsErr   error
}

// PoolConfig holds configuration for connection pools
type PoolConfig struct {
	// MySQL host:port
	Host string
	// Connection parameters
	MaxOpenConns    int
	MaxIdleConns    int
	ConnectTimeout  time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration
	ConnMaxLifetime time.Duration
	ConnMaxIdleTime time.Duration
	TLS             TLSConfig
	DSNParams       map[string]string
	// PreparedStatements configures the optional MySQL prepared statement cache.
	PreparedStatements PreparedStatementsConfig
}

// TLSConfig holds MySQL TLS settings for the driver.
type TLSConfig struct {
	// Mode maps to the go-sql-driver/mysql tls parameter (true, skip-verify, preferred, or custom name).
	Mode string
	// CAFile is an optional CA bundle used to validate the server certificate.
	CAFile string
	// CertFile is an optional client certificate for mTLS.
	CertFile string
	// KeyFile is an optional client key for mTLS.
	KeyFile string
	// ServerName overrides the TLS SNI/server name used for verification.
	ServerName string
	// InsecureSkipVerify skips server cert verification when true.
	InsecureSkipVerify bool
}

// PreparedStatementsConfig configures caching for MySQL prepared statements.
type PreparedStatementsConfig struct {
	Enabled          bool
	MaxStatements    int
	IdleTTL          time.Duration
	CleanupInterval  time.Duration
	PrewarmOnPrepare bool
}

// UserPool represents a connection pool for a specific user
type UserPool struct {
	username  string
	password  string
	db        *sql.DB
	logger    *zap.Logger
	sqlURL    string
	cfg       *PoolConfig
	mu        sync.RWMutex
	connCount int
	stmtCache *stmtCache
}

// NewPoolManager creates a new connection pool manager
func NewPoolManager(cfg *PoolConfig, logger *zap.Logger) *PoolManager {
	if logger == nil {
		logger = zap.NewNop()
	}

	pm := &PoolManager{
		pools:  make(map[string]*UserPool),
		logger: logger,
		cfg:    cfg,
	}

	if cfg != nil {
		metrics.SetMetricLabelsFromMySQLHost(cfg.Host)
	}
	metrics.SetPoolStatsProvider(pm.MetricStats)
	metrics.SetPreparedStmtCacheStatsProvider(pm.PreparedMetricStats)

	return pm
}

// GetPool returns the connection pool for a user, creating it if needed
// Uses a composite key of username:password to ensure credentials are isolated
func (pm *PoolManager) GetPool(username string, password string) (*UserPool, error) {
	// Create composite key for the pool
	poolKey := fmt.Sprintf("%s:%s", username, password)

	// Try to get existing pool
	pm.mu.RLock()
	pool, exists := pm.pools[poolKey]
	pm.mu.RUnlock()

	if exists {
		return pool, nil
	}

	// Create new pool
	pm.mu.Lock()
	defer pm.mu.Unlock()

	// Double-check after acquiring write lock
	if pool, exists = pm.pools[poolKey]; exists {
		return pool, nil
	}

	pool, err := pm.createUserPool(username, password)
	if err != nil {
		return nil, fmt.Errorf("failed to create pool: %w", err)
	}

	pm.pools[poolKey] = pool
	pm.logger.Info("created connection pool for user",
		zap.String("user", username))

	return pool, nil
}

// createUserPool creates a new user-specific connection pool
func (pm *PoolManager) createUserPool(username string, password string) (*UserPool, error) {
	dsn, err := pm.buildDSN(username, password)
	if err != nil {
		return nil, err
	}

	// Create database connection
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	// Configure connection pool
	pm.configureDB(db)

	// Test connection
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := db.PingContext(ctx); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("failed to connect to database: %w", err)
	}

	pool := &UserPool{
		username: username,
		password: password,
		db:       db,
		logger:   pm.logger.With(zap.String("user", username)),
		sqlURL:   fmt.Sprintf("%s@tcp(%s)", username, pm.cfg.Host),
		cfg:      pm.cfg,
	}

	if pm.cfg != nil {
		if cache := newStmtCache(pm.cfg.PreparedStatements, pool.logger); cache != nil {
			cache.startCleanup()
			pool.stmtCache = cache
		}
	}

	return pool, nil
}

func (pm *PoolManager) buildDSN(username string, password string) (string, error) {
	params := url.Values{}
	params.Set("parseTime", "true")
	params.Set("loc", "UTC")
	params.Set("charset", "utf8mb4")
	params.Set("collation", "utf8mb4_bin")
	if pm.cfg.ConnectTimeout > 0 {
		params.Set("timeout", pm.cfg.ConnectTimeout.String())
	}
	if pm.cfg.ReadTimeout > 0 {
		params.Set("readTimeout", pm.cfg.ReadTimeout.String())
	}
	if pm.cfg.WriteTimeout > 0 {
		params.Set("writeTimeout", pm.cfg.WriteTimeout.String())
	}

	tlsParam, err := pm.tlsParamValue()
	if err != nil {
		return "", fmt.Errorf("failed to configure tls: %w", err)
	}
	if tlsParam != "" {
		params.Set("tls", tlsParam)
	}

	if pm.cfg != nil && len(pm.cfg.DSNParams) > 0 {
		for key, value := range pm.cfg.DSNParams {
			trimmedKey := strings.TrimSpace(key)
			if trimmedKey == "" || value == "" {
				continue
			}
			params.Set(trimmedKey, value)
		}
	}

	dsn := fmt.Sprintf("%s:%s@tcp(%s)/?%s",
		url.QueryEscape(username),
		url.QueryEscape(password),
		pm.cfg.Host,
		params.Encode())
	return dsn, nil
}

func (pm *PoolManager) configureDB(db *sql.DB) {
	db.SetMaxOpenConns(pm.cfg.MaxOpenConns)
	db.SetMaxIdleConns(pm.cfg.MaxIdleConns)

	if pm.cfg.ConnMaxLifetime > 0 {
		db.SetConnMaxLifetime(pm.cfg.ConnMaxLifetime)
	}

	if pm.cfg.ConnMaxIdleTime > 0 {
		db.SetConnMaxIdleTime(pm.cfg.ConnMaxIdleTime)
	}
}

// ValidateCredentials verifies credentials using a transient connection, bypassing cached pools.
func (pm *PoolManager) ValidateCredentials(ctx context.Context, username string, password string) error {
	if pm == nil {
		return fmt.Errorf("pool manager is not configured")
	}
	dsn, err := pm.buildDSN(username, password)
	if err != nil {
		return err
	}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return fmt.Errorf("failed to open database: %w", err)
	}
	defer db.Close()
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(0)
	return db.PingContext(ctx)
}

func (pm *PoolManager) tlsParamValue() (string, error) {
	if pm == nil || pm.cfg == nil {
		return "", nil
	}
	pm.tlsOnce.Do(func() {
		pm.tlsParam, pm.tlsErr = buildTLSParam(pm.cfg.TLS)
	})
	return pm.tlsParam, pm.tlsErr
}

func buildTLSParam(cfg TLSConfig) (string, error) {
	if tlsNeedsCustomConfig(cfg) {
		name := tlsConfigKey(cfg)
		tlsConfig, err := buildTLSConfig(cfg)
		if err != nil {
			return "", err
		}
		if err := mysql.RegisterTLSConfig(name, tlsConfig); err != nil {
			return "", err
		}
		return name, nil
	}

	mode := strings.TrimSpace(cfg.Mode)
	if mode != "" {
		lower := strings.ToLower(mode)
		switch lower {
		case "false", "0":
			return "", nil
		case "true", "1", "skip-verify", "preferred":
			return lower, nil
		default:
			return mode, nil
		}
	}

	if cfg.InsecureSkipVerify {
		return "skip-verify", nil
	}

	return "", nil
}

func tlsNeedsCustomConfig(cfg TLSConfig) bool {
	return cfg.CAFile != "" || cfg.CertFile != "" || cfg.KeyFile != "" || cfg.ServerName != ""
}

func tlsConfigKey(cfg TLSConfig) string {
	mode := strings.TrimSpace(cfg.Mode)
	parts := []string{
		mode,
		cfg.CAFile,
		cfg.CertFile,
		cfg.KeyFile,
		cfg.ServerName,
		fmt.Sprintf("%t", cfg.InsecureSkipVerify),
	}
	sum := sha256.Sum256([]byte(strings.Join(parts, "|")))
	return fmt.Sprintf("proxy-tls-%x", sum[:8])
}

func buildTLSConfig(cfg TLSConfig) (*tls.Config, error) {
	mode := strings.TrimSpace(cfg.Mode)
	config := &tls.Config{
		InsecureSkipVerify: cfg.InsecureSkipVerify || strings.EqualFold(mode, "skip-verify"),
	}
	if cfg.ServerName != "" {
		config.ServerName = cfg.ServerName
	}
	if cfg.CAFile != "" {
		pem, err := os.ReadFile(cfg.CAFile)
		if err != nil {
			return nil, fmt.Errorf("read tls ca file: %w", err)
		}
		pool := x509.NewCertPool()
		if ok := pool.AppendCertsFromPEM(pem); !ok {
			return nil, fmt.Errorf("failed to parse tls ca file %q", cfg.CAFile)
		}
		config.RootCAs = pool
	}
	if cfg.CertFile != "" || cfg.KeyFile != "" {
		if cfg.CertFile == "" || cfg.KeyFile == "" {
			return nil, fmt.Errorf("tls cert and key files must both be set")
		}
		cert, err := tls.LoadX509KeyPair(cfg.CertFile, cfg.KeyFile)
		if err != nil {
			return nil, fmt.Errorf("load tls client cert/key: %w", err)
		}
		config.Certificates = []tls.Certificate{cert}
	}
	return config, nil
}

// Close closes all connection pools
func (pm *PoolManager) Close() error {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	for username, pool := range pm.pools {
		if err := pool.Close(); err != nil {
			pm.logger.Error("failed to close pool",
				zap.String("user", username),
				zap.Error(err))
		}
	}

	pm.pools = make(map[string]*UserPool)
	return nil
}

// InvalidateUser closes and removes all pools for a username, regardless of password.
func (pm *PoolManager) InvalidateUser(username string) {
	if pm == nil || username == "" {
		return
	}
	pm.mu.Lock()
	var pools []*UserPool
	for key, pool := range pm.pools {
		if pool != nil && pool.username == username {
			pools = append(pools, pool)
			delete(pm.pools, key)
		}
	}
	pm.mu.Unlock()

	for _, pool := range pools {
		if err := pool.Close(); err != nil {
			pm.logger.Error("failed to close invalidated user pool",
				zap.String("user", username),
				zap.Error(err))
		}
	}
}

// GetStats returns statistics for all pools
func (pm *PoolManager) GetStats() map[string]*PoolStats {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	stats := make(map[string]*PoolStats)
	for username, pool := range pm.pools {
		stats[username] = pool.GetStats()
	}

	return stats
}

// PoolStats contains statistics for a connection pool
type PoolStats struct {
	Username          string
	KeyspaceName      string
	ClusterID         string
	ConnURL           string
	MaxOpenConns      int
	MaxIdleConns      int
	OpenConns         int
	InUseConns        int
	IdleConns         int
	WaitCount         int64
	WaitDuration      int64
	MaxIdleClosed     int64
	MaxLifetimeClosed int64
}

// GetStats returns statistics for this pool
func (up *UserPool) GetStats() *PoolStats {
	up.mu.RLock()
	defer up.mu.RUnlock()

	stats := up.db.Stats()
	metricLabels := metrics.ResolveMetricLabels(context.Background(), up.username)

	return &PoolStats{
		Username:          up.username,
		KeyspaceName:      metricLabels.KeyspaceName,
		ClusterID:         metricLabels.ClusterID,
		ConnURL:           up.sqlURL,
		MaxOpenConns:      stats.MaxOpenConnections,
		MaxIdleConns:      up.cfg.MaxIdleConns,
		OpenConns:         stats.OpenConnections,
		InUseConns:        stats.InUse,
		IdleConns:         stats.Idle,
		WaitCount:         stats.WaitCount,
		WaitDuration:      stats.WaitDuration.Milliseconds(),
		MaxIdleClosed:     stats.MaxIdleClosed,
		MaxLifetimeClosed: stats.MaxLifetimeClosed,
	}
}

// MetricStats returns a metrics-friendly snapshot for all pools.
func (pm *PoolManager) MetricStats() []metrics.PoolStat {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	stats := make([]metrics.PoolStat, 0, len(pm.pools))
	for _, pool := range pm.pools {
		if pool == nil {
			continue
		}
		poolStats := pool.GetStats()
		if poolStats == nil {
			continue
		}
		stats = append(stats, metrics.PoolStat{
			User:         poolStats.Username,
			KeyspaceName: poolStats.KeyspaceName,
			ClusterID:    poolStats.ClusterID,
			Open:         poolStats.OpenConns,
			InUse:        poolStats.InUseConns,
			WaitCount:    poolStats.WaitCount,
			WaitDuration: time.Duration(poolStats.WaitDuration) * time.Millisecond,
		})
	}
	return stats
}

// Execute executes a query on this pool
func (up *UserPool) Execute(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
	return up.db.ExecContext(ctx, query, args...)
}

// Query executes a query that returns rows
func (up *UserPool) Query(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	return up.db.QueryContext(ctx, query, args...)
}

// QueryRow executes a query that returns at most one row
func (up *UserPool) QueryRow(ctx context.Context, query string, args ...interface{}) *sql.Row {
	return up.db.QueryRowContext(ctx, query, args...)
}

// Begin starts a transaction
func (up *UserPool) Begin(ctx context.Context) (*sql.Tx, error) {
	ctx = metrics.WithUsername(ctx, up.username)
	start := time.Now()
	tx, err := up.db.BeginTx(ctx, nil)
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpBeginTx, err == nil, time.Since(start))
	return tx, err
}

// BeginTx starts a transaction with options
func (up *UserPool) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) {
	ctx = metrics.WithUsername(ctx, up.username)
	start := time.Now()
	tx, err := up.db.BeginTx(ctx, opts)
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpBeginTx, err == nil, time.Since(start))
	return tx, err
}

// Close closes this specific pool
func (up *UserPool) Close() error {
	up.mu.Lock()
	defer up.mu.Unlock()

	if up.stmtCache != nil {
		up.stmtCache.stopCleanup()
		up.stmtCache.clear()
	}
	return up.db.Close()
}

// Ping verifies the pool is healthy
func (up *UserPool) Ping(ctx context.Context) error {
	return up.db.PingContext(ctx)
}

// IsHealthy returns true if the pool can connect to MySQL
func (up *UserPool) IsHealthy() bool {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	return up.Ping(ctx) == nil
}

// Conn returns a connection from the pool
func (up *UserPool) Conn(ctx context.Context) (*sql.Conn, error) {
	ctx = metrics.WithUsername(ctx, up.username)
	start := time.Now()
	conn, err := up.db.Conn(ctx)
	d := time.Since(start)
	metrics.ObserveStageFromContext(ctx, metrics.StageMySQLWait, d)
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpPoolWait, err == nil, d)
	return conn, err
}

// PrepareCached returns a cached prepared statement for the SQL string.
// When caching is disabled, it falls back to db.PrepareContext.
func (up *UserPool) PrepareCached(ctx context.Context, query string) (*sql.Stmt, error) {
	if up == nil || up.db == nil {
		return nil, fmt.Errorf("pool not initialized")
	}
	if query == "" {
		return nil, fmt.Errorf("query is required")
	}
	ctx = metrics.WithUsername(ctx, up.username)
	if up.stmtCache == nil {
		start := time.Now()
		stmt, err := up.db.PrepareContext(ctx, query)
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpPrepare, err == nil, time.Since(start))
		return stmt, err
	}
	start := time.Now()
	handle, err := up.stmtCache.prepare(ctx, dbStmtPreparer{db: up.db}, query)
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpPrepare, err == nil, time.Since(start))
	if err != nil {
		return nil, err
	}
	stmt, ok := handle.(*sql.Stmt)
	if !ok || stmt == nil {
		up.stmtCache.recordError()
		return nil, fmt.Errorf("unexpected prepared statement type %T", handle)
	}
	return stmt, nil
}

// EvictPrepared removes a single prepared statement from the cache.
func (up *UserPool) EvictPrepared(query string) {
	if up == nil || up.stmtCache == nil || query == "" {
		return
	}
	up.stmtCache.evict(query)
}

// ClearPrepared closes and removes all cached prepared statements.
func (up *UserPool) ClearPrepared() {
	if up == nil || up.stmtCache == nil {
		return
	}
	up.stmtCache.clear()
}

// PreparedStats returns prepared statement cache statistics.
func (up *UserPool) PreparedStats() PreparedStmtCacheStats {
	if up == nil || up.stmtCache == nil {
		return PreparedStmtCacheStats{}
	}
	return up.stmtCache.stats()
}

// RecordPreparedError increments the prepared statement cache error counter.
func (up *UserPool) RecordPreparedError() {
	if up == nil || up.stmtCache == nil {
		return
	}
	up.stmtCache.recordError()
}

// Metric formatters
func (pm *PoolManager) GetMetricSummary() string {
	stats := pm.GetStats()
	var totalConns int
	for _, stat := range stats {
		totalConns += stat.OpenConns
	}

	return fmt.Sprintf("Total pools: %d, Total connections: %d", len(stats), totalConns)
}

// GetConfig returns the pool configuration
func (pm *PoolManager) GetConfig() *PoolConfig {
	return pm.cfg
}

// ClearPrepared closes cached prepared statements for all user pools.
func (pm *PoolManager) ClearPrepared() {
	if pm == nil {
		return
	}
	pm.mu.RLock()
	pools := make([]*UserPool, 0, len(pm.pools))
	for _, pool := range pm.pools {
		if pool != nil {
			pools = append(pools, pool)
		}
	}
	pm.mu.RUnlock()

	for _, pool := range pools {
		pool.ClearPrepared()
	}
}

// PreparedStats returns prepared statement cache statistics for all pools.
func (pm *PoolManager) PreparedStats() map[string]PreparedStmtCacheStats {
	if pm == nil {
		return nil
	}
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	stats := make(map[string]PreparedStmtCacheStats, len(pm.pools))
	for username, pool := range pm.pools {
		if pool == nil {
			continue
		}
		stats[username] = pool.PreparedStats()
	}
	return stats
}

// PreparedMetricStats returns a metrics-friendly snapshot for all prepared statement caches.
func (pm *PoolManager) PreparedMetricStats() []metrics.PreparedStmtCacheStat {
	if pm == nil {
		return nil
	}
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	stats := make([]metrics.PreparedStmtCacheStat, 0, len(pm.pools))
	for _, pool := range pm.pools {
		if pool == nil {
			continue
		}
		cacheStats := pool.PreparedStats()
		metricLabels := metrics.ResolveMetricLabels(context.Background(), pool.username)
		stats = append(stats, metrics.PreparedStmtCacheStat{
			User:         pool.username,
			KeyspaceName: metricLabels.KeyspaceName,
			ClusterID:    metricLabels.ClusterID,
			Size:         cacheStats.Size,
			Hits:         cacheStats.Hits,
			Misses:       cacheStats.Misses,
			Evictions:    cacheStats.Evictions,
			Errors:       cacheStats.Errors,
		})
	}
	return stats
}