package executor

import (
	"context"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/parser"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/mysqlbackend"
)

var batchHashPartitionClauseRE = regexp.MustCompile("(?is)\\bPARTITION\\s+BY\\s+HASH\\s*\\(\\s*(?:`([^`]+)`|([A-Za-z0-9_]+))\\s*\\)\\s*PARTITIONS\\s+(\\d+)")

type batchHashPartitionHint struct {
	column     string
	partitions int
}

type batchHashPartitionCacheEntry struct {
	expires time.Time
	hint    batchHashPartitionHint
	ok      bool
}

func (e *Executor) batchGroupByTiDBHashPartitionEnabled() bool {
	return e != nil && e.batchGroupByTiDBHashPart
}

func (e *Executor) invalidateBatchPartitionHintCache() {
	if e == nil {
		return
	}
	e.batchPartitionHintMu.Lock()
	defer e.batchPartitionHintMu.Unlock()
	clear(e.batchPartitionHintCache)
}

func (e *Executor) batchHashPartitionSplitKeys(ctx context.Context, user, password string, insert *parser.InsertStatement, params []interface{}) []string {
	hint, ok := e.loadBatchHashPartitionHint(ctx, user, password, insert.Table)
	if !ok {
		return nil
	}
	key, ok := batchHashPartitionBucketKey(insert, params, hint)
	if !ok {
		return nil
	}
	return []string{key}
}

func (e *Executor) loadBatchHashPartitionHint(ctx context.Context, user, password string, table *parser.TableName) (batchHashPartitionHint, bool) {
	if !e.batchGroupByTiDBHashPartitionEnabled() || table == nil {
		return batchHashPartitionHint{}, false
	}
	key := batchTableKey{schema: table.Schema, table: table.Name}
	now := time.Now()

	e.batchPartitionHintMu.RLock()
	entry, ok := e.batchPartitionHintCache[key]
	e.batchPartitionHintMu.RUnlock()
	if ok && (entry.expires.IsZero() || now.Before(entry.expires)) {
		return entry.hint, entry.ok
	}

	hint, ok := e.fetchBatchHashPartitionHint(ctx, user, password, table)
	ttl := e.batchPartitionHintTTL
	if ttl <= 0 {
		ttl = 30 * time.Second
	}
	e.batchPartitionHintMu.Lock()
	e.batchPartitionHintCache[key] = batchHashPartitionCacheEntry{
		expires: now.Add(ttl),
		hint:    hint,
		ok:      ok,
	}
	e.batchPartitionHintMu.Unlock()
	return hint, ok
}

func (e *Executor) fetchBatchHashPartitionHint(ctx context.Context, user, password string, table *parser.TableName) (batchHashPartitionHint, bool) {
	if e == nil || e.schemaManager == nil || table == nil {
		return batchHashPartitionHint{}, false
	}

	metaCols, available, err := e.schemaManager.CQLColumnsForTable(ctx, user, password, table.Schema, table.Name)
	if err != nil || !available {
		return batchHashPartitionHint{}, false
	}
	column, ok := singleHashPartitionKeyColumn(metaCols)
	if !ok {
		return batchHashPartitionHint{}, false
	}

	createSQL, ok := e.loadBatchShowCreateTable(ctx, table.Schema, table.Name, user, password)
	if !ok {
		return batchHashPartitionHint{}, false
	}
	hint, ok := parseBatchHashPartitionHintFromShowCreate(createSQL)
	if !ok || normalizeBatchColumnName(hint.column) != normalizeBatchColumnName(column) || hint.partitions <= 1 {
		return batchHashPartitionHint{}, false
	}
	return hint, true
}

func singleHashPartitionKeyColumn(columns []mysqlbackend.CQLColumnInfo) (string, bool) {
	partitionCols := make([]mysqlbackend.CQLColumnInfo, 0, 1)
	for _, col := range columns {
		if !strings.EqualFold(strings.TrimSpace(col.Kind), "partition_key") {
			continue
		}
		partitionCols = append(partitionCols, col)
	}
	if len(partitionCols) != 1 {
		return "", false
	}
	if !isHashPartitionableCQLType(partitionCols[0].CQLType) {
		return "", false
	}
	return partitionCols[0].Name, true
}

func isHashPartitionableCQLType(cqlType string) bool {
	switch strings.ToUpper(strings.TrimSpace(cqlType)) {
	case "TINYINT", "SMALLINT", "INT", "BIGINT", "COUNTER":
		return true
	default:
		return false
	}
}

func (e *Executor) loadBatchShowCreateTable(ctx context.Context, keyspaceName, tableName, user, password string) (string, bool) {
	if e == nil {
		return "", false
	}
	conn, release, err := e.getConn(ctx, user, password)
	if err != nil {
		return "", false
	}
	if release != nil {
		defer release()
	}

	showSQL := fmt.Sprintf("SHOW CREATE TABLE %s.%s", quoteMySQLIdentifier(keyspaceName), quoteMySQLIdentifier(tableName))
	rows, err := conn.QueryContext(ctx, showSQL)
	if err != nil {
		return "", false
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil || len(columns) == 0 || !rows.Next() {
		return "", false
	}

	values := make([]interface{}, len(columns))
	dest := make([]interface{}, len(columns))
	for i := range values {
		dest[i] = &values[i]
	}
	if err := rows.Scan(dest...); err != nil {
		return "", false
	}
	if err := rows.Err(); err != nil {
		return "", false
	}

	createSQL := ""
	for i, col := range columns {
		if strings.EqualFold(col, "Create Table") {
			createSQL = describeCellString(values[i])
			break
		}
	}
	if createSQL == "" {
		for i := len(values) - 1; i >= 0; i-- {
			if s := strings.TrimSpace(describeCellString(values[i])); s != "" {
				createSQL = s
				break
			}
		}
	}
	if strings.TrimSpace(createSQL) == "" {
		return "", false
	}
	return createSQL, true
}

func parseBatchHashPartitionHintFromShowCreate(createSQL string) (batchHashPartitionHint, bool) {
	matches := batchHashPartitionClauseRE.FindStringSubmatch(createSQL)
	if len(matches) != 4 {
		return batchHashPartitionHint{}, false
	}
	column := matches[1]
	if column == "" {
		column = matches[2]
	}
	partitions, err := strconv.Atoi(strings.TrimSpace(matches[3]))
	if err != nil || partitions <= 0 {
		return batchHashPartitionHint{}, false
	}
	return batchHashPartitionHint{
		column:     strings.TrimSpace(column),
		partitions: partitions,
	}, true
}

func batchHashPartitionBucketKey(insert *parser.InsertStatement, params []interface{}, hint batchHashPartitionHint) (string, bool) {
	if insert == nil || insert.Table == nil || len(insert.Values) != 1 || hint.partitions <= 1 {
		return "", false
	}
	idx := -1
	for i, col := range insert.Columns {
		if normalizeBatchColumnName(col) == normalizeBatchColumnName(hint.column) {
			idx = i
			break
		}
	}
	if idx < 0 || idx >= len(insert.Values[0]) {
		return "", false
	}
	value, err := resolveScalarExpr(insert.Values[0][idx], params, buildBindMarkerIndex(insert))
	if err != nil {
		return "", false
	}
	bucket, ok := batchHashPartitionBucket(value, hint.partitions)
	if !ok {
		return "", false
	}
	return fmt.Sprintf("%s.%s#hash:%d", insert.Table.Schema, insert.Table.Name, bucket), true
}

func batchHashPartitionBucket(value interface{}, partitions int) (int, bool) {
	if partitions <= 1 {
		return 0, false
	}
	numeric, err := batchNumericAsInt64(value)
	if err != nil {
		return 0, false
	}
	mod := numeric % int64(partitions)
	if mod < 0 {
		mod += int64(partitions)
	}
	return int(mod), true
}