package mysqlbackend

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

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

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

const defaultSchemaCacheTTL = 5 * time.Second

type userKey struct {
	user     string
	password string
}

type tableKey struct {
	schema string
	table  string
}

type indexKey struct {
	schema string
	name   string
}

type schemaCacheEntry struct {
	expires  time.Time
	list     []string
	set      map[string]struct{}
	lowerSet map[string]struct{}
}

type tableCacheEntry struct {
	expires             time.Time
	list                []TableInfo
	tableSet            map[tableKey]struct{}
	bySchema            map[string][]TableInfo
	tableSetLower       map[tableKey]struct{}
	tableByLower        map[tableKey]TableInfo
	tableLowerAmbiguous map[tableKey]struct{}
}

type columnCacheEntry struct {
	expires      time.Time
	list         []ColumnInfo
	byTable      map[tableKey][]ColumnInfo
	byTableLower map[tableKey][]ColumnInfo
}

type indexCacheEntry struct {
	expires      time.Time
	list         []IndexInfo
	byTable      map[tableKey][]IndexInfo
	byName       map[indexKey][]IndexInfo
	byTableLower map[tableKey][]IndexInfo
	byNameLower  map[indexKey][]IndexInfo
}

// SchemaCache keeps a small TTL cache for information_schema queries.
type SchemaCache struct {
	ttl         time.Duration
	mu          sync.RWMutex
	schemas     map[userKey]schemaCacheEntry
	tables      map[userKey]tableCacheEntry
	columns     map[userKey]columnCacheEntry
	indexes     map[userKey]indexCacheEntry
	refreshMu   sync.Mutex
	schemaLocks map[userKey]*sync.Mutex
	tableLocks  map[userKey]*sync.Mutex
	columnLocks map[userKey]*sync.Mutex
	indexLocks  map[userKey]*sync.Mutex
}

// NewSchemaCache creates a SchemaCache with the provided TTL.
func NewSchemaCache(ttl time.Duration) *SchemaCache {
	if ttl <= 0 {
		ttl = defaultSchemaCacheTTL
	}
	return &SchemaCache{
		ttl:     ttl,
		schemas: make(map[userKey]schemaCacheEntry),
		tables:  make(map[userKey]tableCacheEntry),
		columns: make(map[userKey]columnCacheEntry),
		indexes: make(map[userKey]indexCacheEntry),
	}
}

func (c *SchemaCache) lockFor(m *map[userKey]*sync.Mutex, key userKey) *sync.Mutex {
	if c == nil {
		return nil
	}
	c.refreshMu.Lock()
	defer c.refreshMu.Unlock()
	if *m == nil {
		*m = make(map[userKey]*sync.Mutex)
	}
	lock := (*m)[key]
	if lock == nil {
		lock = &sync.Mutex{}
		(*m)[key] = lock
	}
	return lock
}

func (c *SchemaCache) schemaLock(key userKey) *sync.Mutex {
	return c.lockFor(&c.schemaLocks, key)
}

func (c *SchemaCache) tableLock(key userKey) *sync.Mutex {
	return c.lockFor(&c.tableLocks, key)
}

func (c *SchemaCache) columnLock(key userKey) *sync.Mutex {
	return c.lockFor(&c.columnLocks, key)
}

func (c *SchemaCache) indexLock(key userKey) *sync.Mutex {
	return c.lockFor(&c.indexLocks, key)
}

// InvalidateUser clears cached metadata for a specific user.
func (c *SchemaCache) InvalidateUser(user, password string) {
	if c == nil {
		return
	}
	key := userKey{user: user, password: password}
	c.mu.Lock()
	delete(c.schemas, key)
	delete(c.tables, key)
	delete(c.columns, key)
	delete(c.indexes, key)
	c.mu.Unlock()

	c.refreshMu.Lock()
	delete(c.schemaLocks, key)
	delete(c.tableLocks, key)
	delete(c.columnLocks, key)
	delete(c.indexLocks, key)
	c.refreshMu.Unlock()
}

// InvalidateAll clears cached metadata for all users.
func (c *SchemaCache) InvalidateAll() {
	if c == nil {
		return
	}
	c.mu.Lock()
	c.schemas = make(map[userKey]schemaCacheEntry)
	c.tables = make(map[userKey]tableCacheEntry)
	c.columns = make(map[userKey]columnCacheEntry)
	c.indexes = make(map[userKey]indexCacheEntry)
	c.mu.Unlock()

	c.refreshMu.Lock()
	c.schemaLocks = nil
	c.tableLocks = nil
	c.columnLocks = nil
	c.indexLocks = nil
	c.refreshMu.Unlock()
}

// Schemas returns cached schema entries for the user, refreshing as needed.
func (c *SchemaCache) Schemas(ctx context.Context, pm *PoolManager, user, password string) (schemaCacheEntry, error) {
	if c == nil {
		return schemaCacheEntry{}, fmt.Errorf("schema cache not initialized")
	}
	key := userKey{user: user, password: password}
	now := time.Now()

	c.mu.RLock()
	entry, ok := c.schemas[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	lock := c.schemaLock(key)
	if lock != nil {
		lock.Lock()
		defer lock.Unlock()
	}

	now = time.Now()
	c.mu.RLock()
	entry, ok = c.schemas[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	start := time.Now()
	list, err := fetchSchemas(ctx, pm, user, password)
	if err != nil {
		return schemaCacheEntry{}, err
	}
	set := make(map[string]struct{}, len(list))
	lowerSet := make(map[string]struct{}, len(list))
	for _, name := range list {
		set[name] = struct{}{}
		lowerSet[normalizeLookup(name)] = struct{}{}
	}

	entry = schemaCacheEntry{
		expires:  now.Add(c.ttl),
		list:     list,
		set:      set,
		lowerSet: lowerSet,
	}

	metrics.ObserveStageFromContext(ctx, metrics.StageSchemaCache, time.Since(start))

	c.mu.Lock()
	c.schemas[key] = entry
	c.mu.Unlock()

	return entry, nil
}

// Tables returns cached table entries for the user, refreshing as needed.
func (c *SchemaCache) Tables(ctx context.Context, pm *PoolManager, user, password string) (tableCacheEntry, error) {
	if c == nil {
		return tableCacheEntry{}, fmt.Errorf("schema cache not initialized")
	}
	key := userKey{user: user, password: password}
	now := time.Now()

	c.mu.RLock()
	entry, ok := c.tables[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	lock := c.tableLock(key)
	if lock != nil {
		lock.Lock()
		defer lock.Unlock()
	}

	now = time.Now()
	c.mu.RLock()
	entry, ok = c.tables[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	start := time.Now()
	list, err := fetchTables(ctx, pm, user, password)
	if err != nil {
		return tableCacheEntry{}, err
	}

	tableSet := make(map[tableKey]struct{}, len(list))
	tableSetLower := make(map[tableKey]struct{}, len(list))
	tableByLower := make(map[tableKey]TableInfo, len(list))
	tableLowerAmbiguous := make(map[tableKey]struct{})
	bySchema := make(map[string][]TableInfo)
	for _, tbl := range list {
		tableSet[tableKey{schema: tbl.Schema, table: tbl.Name}] = struct{}{}
		bySchema[tbl.Schema] = append(bySchema[tbl.Schema], tbl)
		keyLower := normalizeTableKey(tbl.Schema, tbl.Name)
		tableSetLower[keyLower] = struct{}{}
		if _, ok := tableByLower[keyLower]; ok {
			tableLowerAmbiguous[keyLower] = struct{}{}
		} else {
			tableByLower[keyLower] = tbl
		}
	}

	entry = tableCacheEntry{
		expires:             now.Add(c.ttl),
		list:                list,
		tableSet:            tableSet,
		bySchema:            bySchema,
		tableSetLower:       tableSetLower,
		tableByLower:        tableByLower,
		tableLowerAmbiguous: tableLowerAmbiguous,
	}

	metrics.ObserveStageFromContext(ctx, metrics.StageSchemaCache, time.Since(start))

	c.mu.Lock()
	c.tables[key] = entry
	c.mu.Unlock()

	return entry, nil
}

// Columns returns cached column entries for the user, refreshing as needed.
func (c *SchemaCache) Columns(ctx context.Context, pm *PoolManager, user, password string) (columnCacheEntry, error) {
	if c == nil {
		return columnCacheEntry{}, fmt.Errorf("schema cache not initialized")
	}
	key := userKey{user: user, password: password}
	now := time.Now()

	c.mu.RLock()
	entry, ok := c.columns[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	lock := c.columnLock(key)
	if lock != nil {
		lock.Lock()
		defer lock.Unlock()
	}

	now = time.Now()
	c.mu.RLock()
	entry, ok = c.columns[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	start := time.Now()
	list, err := fetchColumns(ctx, pm, user, password)
	if err != nil {
		return columnCacheEntry{}, err
	}

	byTable := make(map[tableKey][]ColumnInfo)
	byTableLower := make(map[tableKey][]ColumnInfo)
	for _, col := range list {
		k := tableKey{schema: col.Schema, table: col.Table}
		byTable[k] = append(byTable[k], col)
		kLower := normalizeTableKey(col.Schema, col.Table)
		byTableLower[kLower] = append(byTableLower[kLower], col)
	}

	entry = columnCacheEntry{
		expires:      now.Add(c.ttl),
		list:         list,
		byTable:      byTable,
		byTableLower: byTableLower,
	}

	metrics.ObserveStageFromContext(ctx, metrics.StageSchemaCache, time.Since(start))

	c.mu.Lock()
	c.columns[key] = entry
	c.mu.Unlock()

	return entry, nil
}

// Indexes returns cached index entries for the user, refreshing as needed.
func (c *SchemaCache) Indexes(ctx context.Context, pm *PoolManager, user, password string) (indexCacheEntry, error) {
	if c == nil {
		return indexCacheEntry{}, fmt.Errorf("schema cache not initialized")
	}
	key := userKey{user: user, password: password}
	now := time.Now()

	c.mu.RLock()
	entry, ok := c.indexes[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	lock := c.indexLock(key)
	if lock != nil {
		lock.Lock()
		defer lock.Unlock()
	}

	now = time.Now()
	c.mu.RLock()
	entry, ok = c.indexes[key]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry, nil
	}

	start := time.Now()
	list, err := fetchIndexes(ctx, pm, user, password)
	if err != nil {
		return indexCacheEntry{}, err
	}

	byTable := make(map[tableKey][]IndexInfo)
	byName := make(map[indexKey][]IndexInfo)
	byTableLower := make(map[tableKey][]IndexInfo)
	byNameLower := make(map[indexKey][]IndexInfo)
	for _, idx := range list {
		tKey := tableKey{schema: idx.Schema, table: idx.Table}
		byTable[tKey] = append(byTable[tKey], idx)
		iKey := indexKey{schema: idx.Schema, name: idx.Name}
		byName[iKey] = append(byName[iKey], idx)
		tKeyLower := normalizeTableKey(idx.Schema, idx.Table)
		byTableLower[tKeyLower] = append(byTableLower[tKeyLower], idx)
		iKeyLower := normalizeIndexKey(idx.Schema, idx.Name)
		byNameLower[iKeyLower] = append(byNameLower[iKeyLower], idx)
	}

	entry = indexCacheEntry{
		expires:      now.Add(c.ttl),
		list:         list,
		byTable:      byTable,
		byName:       byName,
		byTableLower: byTableLower,
		byNameLower:  byNameLower,
	}

	metrics.ObserveStageFromContext(ctx, metrics.StageSchemaCache, time.Since(start))

	c.mu.Lock()
	c.indexes[key] = entry
	c.mu.Unlock()

	return entry, nil
}

func fetchSchemas(ctx context.Context, pm *PoolManager, user, password string) ([]string, error) {
	conn, release, err := getConn(ctx, pm, user, password)
	if err != nil {
		return nil, err
	}
	if release != nil {
		defer release()
	}

	rows, err := queryContextWithMetrics(ctx, conn, metrics.MySQLOpQueryRead, `
		SELECT schema_name
		FROM information_schema.schemata
		ORDER BY schema_name
	`)
	if err != nil {
		return nil, fmt.Errorf("failed to query schemata: %w", err)
	}
	defer rows.Close()

	list := make([]string, 0)
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, fmt.Errorf("failed to scan schema name: %w", err)
		}
		list = append(list, name)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("failed to read schema names: %w", err)
	}
	return list, nil
}

func fetchTables(ctx context.Context, pm *PoolManager, user, password string) ([]TableInfo, error) {
	conn, release, err := getConn(ctx, pm, user, password)
	if err != nil {
		return nil, err
	}
	if release != nil {
		defer release()
	}

	rows, err := queryContextWithMetrics(ctx, conn, metrics.MySQLOpQueryRead, `
		SELECT table_schema, table_name, table_comment
		FROM information_schema.tables
		WHERE table_type = 'BASE TABLE' OR LOWER(table_schema) = 'information_schema'
		ORDER BY table_schema, table_name
	`)
	if err != nil {
		return nil, fmt.Errorf("failed to query tables: %w", err)
	}
	defer rows.Close()

	list := make([]TableInfo, 0)
	for rows.Next() {
		var info TableInfo
		if err := rows.Scan(&info.Schema, &info.Name, &info.Comment); err != nil {
			return nil, fmt.Errorf("failed to scan table info: %w", err)
		}
		list = append(list, info)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("failed to read table info: %w", err)
	}
	return list, nil
}

func fetchColumns(ctx context.Context, pm *PoolManager, user, password string) ([]ColumnInfo, error) {
	_, rows, release, err := queryMetadataContextWithRetry(ctx, pm, user, password, `
		SELECT table_schema, table_name, column_name, data_type, numeric_precision, numeric_scale, ordinal_position
		FROM information_schema.columns
		ORDER BY table_schema, table_name, ordinal_position
	`)
	defer releaseRows(rows, release)
	if err != nil {
		return nil, fmt.Errorf("failed to query columns: %w", err)
	}

	list := make([]ColumnInfo, 0)
	for rows.Next() {
		var info ColumnInfo
		var precision sql.NullInt64
		var scale sql.NullInt64
		if err := rows.Scan(&info.Schema, &info.Table, &info.Name, &info.DataType, &precision, &scale, &info.Ordinal); err != nil {
			return nil, fmt.Errorf("failed to scan column info: %w", err)
		}
		if precision.Valid {
			info.NumericPrecision = int(precision.Int64)
		}
		if scale.Valid {
			info.NumericScale = int(scale.Int64)
		}
		list = append(list, info)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("failed to read column info: %w", err)
	}
	return list, nil
}

func fetchIndexes(ctx context.Context, pm *PoolManager, user, password string) ([]IndexInfo, error) {
	conn, release, err := getConn(ctx, pm, user, password)
	if err != nil {
		return nil, err
	}
	if release != nil {
		defer release()
	}

	rows, err := queryContextWithMetrics(ctx, conn, metrics.MySQLOpQueryRead, `
		SELECT table_schema, table_name, index_name, column_name, seq_in_index, non_unique, collation
		FROM information_schema.statistics
		ORDER BY table_schema, table_name, index_name, seq_in_index
	`)
	if err != nil {
		return nil, fmt.Errorf("failed to query indexes: %w", err)
	}
	defer rows.Close()

	list := make([]IndexInfo, 0)
	for rows.Next() {
		var info IndexInfo
		var collation sql.NullString
		var nonUnique int
		if err := rows.Scan(&info.Schema, &info.Table, &info.Name, &info.Column, &info.Seq, &nonUnique, &collation); err != nil {
			return nil, fmt.Errorf("failed to scan index info: %w", err)
		}
		if collation.Valid {
			info.Collation = collation.String
		}
		info.NonUnique = nonUnique != 0
		info.IsPrimary = info.Name == "PRIMARY"
		list = append(list, info)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("failed to read index info: %w", err)
	}
	return list, nil
}

func normalizeLookup(name string) string {
	for i := 0; i < len(name); i++ {
		ch := name[i]
		if ch >= 'A' && ch <= 'Z' {
			return strings.ToLower(name)
		}
	}
	return name
}

func normalizeTableKey(schema, table string) tableKey {
	return tableKey{
		schema: normalizeLookup(schema),
		table:  normalizeLookup(table),
	}
}

func normalizeIndexKey(schema, name string) indexKey {
	return indexKey{
		schema: normalizeLookup(schema),
		name:   normalizeLookup(name),
	}
}

func getConn(ctx context.Context, pm *PoolManager, user, password string) (*sql.Conn, func(), error) {
	if conn, ok := ConnFromContext(ctx); ok {
		return conn, nil, nil
	}
	return getPoolConn(ctx, pm, user, password)
}

func getPoolConn(ctx context.Context, pm *PoolManager, user, password string) (*sql.Conn, func(), error) {
	leaseRelease, err := AcquireConnLeaseFromContext(ctx, ConnLeaseKindStandard)
	if err != nil {
		return nil, nil, err
	}
	if pm == nil {
		if leaseRelease != nil {
			leaseRelease()
		}
		return nil, nil, fmt.Errorf("pool manager is nil")
	}
	pool, err := pm.GetPool(user, password)
	if err != nil {
		if leaseRelease != nil {
			leaseRelease()
		}
		return nil, nil, fmt.Errorf("failed to get connection pool: %w", err)
	}
	conn, err := pool.Conn(ctx)
	if err != nil {
		if leaseRelease != nil {
			leaseRelease()
		}
		return nil, nil, fmt.Errorf("failed to get connection: %w", err)
	}
	var once sync.Once
	release := func() {
		once.Do(func() {
			_ = conn.Close()
			if leaseRelease != nil {
				leaseRelease()
			}
		})
	}
	return conn, release, nil
}

// IsMySQLConnectionError reports whether err means the underlying MySQL
// connection state is no longer safe to reuse.
func IsMySQLConnectionError(err error) bool {
	if err == nil {
		return false
	}
	return errors.Is(err, mysql.ErrInvalidConn) || isMySQLBusyBufferError(err)
}

func isMySQLBusyBufferError(err error) bool {
	if err == nil {
		return false
	}
	return errors.Is(err, mysql.ErrBusyBuffer) ||
		strings.Contains(strings.ToLower(err.Error()), "busy buffer")
}

func queryMetadataContextWithRetry(ctx context.Context, pm *PoolManager, user, password, query string, args ...interface{}) (*sql.Conn, *sql.Rows, func(), error) {
	logger := zap.NewNop()
	if pm != nil && pm.logger != nil {
		logger = pm.logger
	}

	conn, release, err := getConn(ctx, pm, user, password)
	if err != nil {
		return nil, nil, nil, err
	}

	rows, err := queryContextWithMetrics(ctx, conn, metrics.MySQLOpQueryRead, query, args...)
	if err == nil {
		logger.Debug("metadata query succeeded")
		return conn, rows, release, nil
	}

	if !IsMySQLConnectionError(err) {
		logger.Error("metadata query failed with non connection error", zap.Error(err))
		return conn, nil, release, err
	}

	pinnedConnection := release == nil
	logFields := []zap.Field{
		zap.Error(err),
		zap.Bool("pinned_connection", pinnedConnection),
	}
	if isMySQLBusyBufferError(err) {
		logger.Warn("metadata query hit busy buffer, retrying with fresh connection", logFields...)
	} else {
		logger.Warn("metadata query failed, retrying with fresh connection", logFields...)
	}

	if release != nil {
		release()
	}
	// A nil release means the connection is pinned in ctx and owned by the
	// caller. Do not mark it bad here: sql.Conn.Raw may block while the caller
	// still owns the connection, and metadata refreshes hold per-user locks.

	retryConn, retryRelease, retryConnErr := getPoolConn(ctx, pm, user, password)
	if retryConnErr != nil {
		return nil, nil, nil, errors.Join(err, retryConnErr)
	}
	retryRows, retryErr := queryContextWithMetrics(ctx, retryConn, metrics.MySQLOpQueryRead, query, args...)
	if retryErr != nil {
		if isMySQLBusyBufferError(retryErr) {
			logger.Warn("metadata retry query hit busy buffer",
				zap.Error(retryErr),
				zap.Bool("pinned_connection", false))
		}
		if retryRelease != nil {
			retryRelease()
		}
		return nil, nil, nil, errors.Join(err, retryErr)
	}
	return retryConn, retryRows, retryRelease, nil
}

func releaseRows(rows *sql.Rows, release func()) {
	if rows != nil {
		_ = rows.Close()
	}
	if release != nil {
		release()
	}
}