package executor

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

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

type batchPrimaryKeyDebugGroup struct {
	ExecuteItemIndex       int      `json:"execute_item_index"`
	OriginalStatementCount int      `json:"original_statement_count"`
	GroupingStrategy       string   `json:"grouping_strategy,omitempty"`
	BatchTemplateKey       string   `json:"batch_template_key,omitempty"`
	GroupBreakReason       string   `json:"group_break_reason,omitempty"`
	PartitionHits          []string `json:"partition_hits,omitempty"`
	PrimaryKeys            []string `json:"primary_keys"`
}

type batchPrimaryKeyDebugKey struct {
	Prefix string
	Value  string
	Full   string
}

type batchPrimaryKeyDebugGroupKeys struct {
	ExecuteItemIndex       int
	OriginalStatementCount int
	GroupingStrategy       batchGroupingStrategy
	BatchTemplateKey       string
	GroupBreakReason       string
	PartitionHits          []string
	PrimaryKeys            []batchPrimaryKeyDebugKey
}

func (e *Executor) logBatchPrimaryKeys(ctx context.Context, user, password, batchType string, groups []coalescedBatchGroup, planningDebug *batchBestEffortPlanningDebug) {
	if e == nil || e.logger == nil || len(groups) == 0 || !e.logger.Core().Enabled(zap.DebugLevel) {
		return
	}

	pkCache := make(map[batchTableKey][]string)
	resolvePK := func(table *parser.TableName) ([]string, error) {
		if table == nil {
			return nil, fmt.Errorf("table is required")
		}
		key := batchTableKey{schema: table.Schema, table: table.Name}
		if cols, ok := pkCache[key]; ok {
			return cols, nil
		}
		cols, err := e.primaryKeyColumns(ctx, user, password, table)
		if err != nil {
			return nil, err
		}
		pkCache[key] = cols
		return cols, nil
	}

	prefix, inputKeys, mergeAppliedKeys, mergedKeys := collectBatchPrimaryKeyDebugInfo(groups, resolvePK)
	fields := []zap.Field{
		zap.String("batch_type", batchType),
		zap.Bool("batch_partition_group_strategy_enabled", e.batchGroupByTiDBHashPartitionEnabled()),
		zap.Bool("batch_partition_group_strategy_applied", batchPartitionGroupStrategyApplied(groups)),
		zap.Strings("batch_input_primary_keys", inputKeys),
		zap.Strings("batch_merge_applied_primary_keys", mergeAppliedKeys),
		zap.Any("batch_merged_primary_keys", mergedKeys),
	}
	if planningDebug != nil {
		fields = append(fields,
			zap.Bool("batch_global_fold_planner_enabled", planningDebug.globalFoldPlannerEnabled),
			zap.Bool("batch_global_fold_planner_applied", planningDebug.globalFoldPlannerApplied),
			zap.String("batch_global_fold_fallback_reason", planningDebug.globalFoldFallbackReason),
		)
	}
	if prefix != "" {
		fields = append(fields, zap.String("batch_primary_key_prefix", prefix))
	}
	if tracker := metrics.FromContext(ctx); tracker != nil {
		if payload := tracker.ClientPayload(); len(payload) > 0 {
			fields = append(fields, zap.Any("client_payload", payload))
		}
	}
	e.logger.Debug("batch merge primary keys", fields...)
}

func collectBatchPrimaryKeyDebugInfo(groups []coalescedBatchGroup, resolvePK batchUpdatePKResolver) (string, []string, []string, []batchPrimaryKeyDebugGroup) {
	if len(groups) == 0 {
		return "", nil, nil, nil
	}

	inputKeys := make([]batchPrimaryKeyDebugKey, 0, len(groups))
	coalescedKeys := make([]batchPrimaryKeyDebugKey, 0, len(groups))
	mergedKeys := make([]batchPrimaryKeyDebugGroupKeys, 0, len(groups))
	for i, group := range groups {
		for _, entry := range group.originals {
			inputKeys = append(inputKeys, batchEntryPrimaryKeyDebugKeys(entry, resolvePK)...)
		}
		groupMergedKeys := batchEntryPrimaryKeyDebugKeys(group.merged, resolvePK)
		coalescedKeys = append(coalescedKeys, groupMergedKeys...)
		mergedKeys = append(mergedKeys, batchPrimaryKeyDebugGroupKeys{
			ExecuteItemIndex:       i,
			OriginalStatementCount: len(group.originals),
			GroupingStrategy:       group.groupingStrategy,
			BatchTemplateKey:       group.debug.templateKey,
			GroupBreakReason:       group.debug.breakReason,
			PartitionHits:          append([]string(nil), group.partitionHits...),
			PrimaryKeys:            groupMergedKeys,
		})
	}
	if prefix, ok := batchPrimaryKeySharedPrefix(inputKeys, mergedKeys); ok {
		return prefix,
			batchPrimaryKeyValues(inputKeys, true),
			batchPrimaryKeyValues(batchMergeAppliedKeys(inputKeys, coalescedKeys), true),
			batchPrimaryKeyGroupValues(mergedKeys, true)
	}
	return "",
		batchPrimaryKeyValues(inputKeys, false),
		batchPrimaryKeyValues(batchMergeAppliedKeys(inputKeys, coalescedKeys), false),
		batchPrimaryKeyGroupValues(mergedKeys, false)
}

func batchEntryPrimaryKeyStrings(entry BatchEntry, resolvePK batchUpdatePKResolver) []string {
	keys := batchEntryPrimaryKeyDebugKeys(entry, resolvePK)
	out := make([]string, 0, len(keys))
	for _, key := range keys {
		out = append(out, key.Full)
	}
	return out
}

func batchEntryPrimaryKeyDebugKeys(entry BatchEntry, resolvePK batchUpdatePKResolver) []batchPrimaryKeyDebugKey {
	switch stmt := entry.Statement.(type) {
	case *parser.InsertStatement:
		return batchInsertPrimaryKeyDebugKeys(stmt, entry.Params, resolvePK)
	case *parser.UpdateStatement:
		return batchWherePrimaryKeyDebugKeys(stmt, stmt.Table, stmt.Where, entry.Params, resolvePK)
	case *parser.DeleteStatement:
		return batchWherePrimaryKeyDebugKeys(stmt, stmt.From, stmt.Where, entry.Params, resolvePK)
	default:
		return []batchPrimaryKeyDebugKey{{Full: fmt.Sprintf("<unsupported:%T>", entry.Statement)}}
	}
}

func batchInsertPrimaryKeyDebugKeys(stmt *parser.InsertStatement, params []interface{}, resolvePK batchUpdatePKResolver) []batchPrimaryKeyDebugKey {
	if stmt == nil {
		return []batchPrimaryKeyDebugKey{{Full: "<nil_insert_statement>"}}
	}
	if stmt.Table == nil {
		return []batchPrimaryKeyDebugKey{{Full: "<missing_table>"}}
	}
	if resolvePK == nil {
		return nil
	}

	pkCols, err := resolvePK(stmt.Table)
	if err != nil {
		return []batchPrimaryKeyDebugKey{{Full: fmt.Sprintf("%s(<pk_lookup_error:%v>)", batchTableLogName(stmt.Table), err)}}
	}
	if len(pkCols) == 0 {
		return []batchPrimaryKeyDebugKey{{Full: fmt.Sprintf("%s(<missing_primary_key>)", batchTableLogName(stmt.Table))}}
	}

	columnIndex := make(map[string]int, len(stmt.Columns))
	for i, col := range stmt.Columns {
		columnIndex[normalizeBatchColumnName(col)] = i
	}
	markerIndex := buildBindMarkerIndex(stmt)
	keys := make([]batchPrimaryKeyDebugKey, 0, len(stmt.Values))
	prefix := batchPrimaryKeyPrefix(stmt.Table, pkCols)
	for _, row := range stmt.Values {
		labeledParts := make([]string, 0, len(pkCols))
		rawValues := make([]string, 0, len(pkCols))
		for _, pk := range pkCols {
			idx, ok := columnIndex[normalizeBatchColumnName(pk)]
			if !ok {
				labeledParts = append(labeledParts, fmt.Sprintf("%s=<missing>", pk))
				rawValues = append(rawValues, "<missing>")
				continue
			}
			if idx >= len(row) {
				labeledParts = append(labeledParts, fmt.Sprintf("%s=<row_index_out_of_range>", pk))
				rawValues = append(rawValues, "<row_index_out_of_range>")
				continue
			}
			value := batchPrimaryKeyExprValueString(row[idx], params, markerIndex)
			labeledParts = append(labeledParts, fmt.Sprintf("%s=%s", pk, value))
			rawValues = append(rawValues, value)
		}
		keys = append(keys, batchPrimaryKeyDebugKey{
			Prefix: prefix,
			Value:  strings.Join(rawValues, ","),
			Full:   fmt.Sprintf("%s(%s)", batchTableLogName(stmt.Table), strings.Join(labeledParts, ", ")),
		})
	}
	if len(keys) == 0 {
		return []batchPrimaryKeyDebugKey{{Full: fmt.Sprintf("%s(<empty_values>)", batchTableLogName(stmt.Table))}}
	}
	return keys
}

func batchWherePrimaryKeyDebugKeys(stmt parser.Statement, table *parser.TableName, where parser.Expression, params []interface{}, resolvePK batchUpdatePKResolver) []batchPrimaryKeyDebugKey {
	if table == nil {
		return []batchPrimaryKeyDebugKey{{Full: "<missing_table>"}}
	}
	if resolvePK == nil {
		return nil
	}

	pkCols, err := resolvePK(table)
	if err != nil {
		return []batchPrimaryKeyDebugKey{{Full: fmt.Sprintf("%s(<pk_lookup_error:%v>)", batchTableLogName(table), err)}}
	}
	if len(pkCols) == 0 {
		return []batchPrimaryKeyDebugKey{{Full: fmt.Sprintf("%s(<missing_primary_key>)", batchTableLogName(table))}}
	}

	conditions, err := extractEqualityConditions(where)
	if err != nil {
		return []batchPrimaryKeyDebugKey{{Full: fmt.Sprintf("%s(<where_parse_error:%v>)", batchTableLogName(table), err)}}
	}
	condByColumn := make(map[string]equalityCondition, len(conditions))
	for _, cond := range conditions {
		condByColumn[normalizeBatchColumnName(cond.column)] = cond
	}

	markerIndex := buildBindMarkerIndex(stmt)
	labeledParts := make([]string, 0, len(pkCols))
	rawValues := make([]string, 0, len(pkCols))
	for _, pk := range pkCols {
		cond, ok := condByColumn[normalizeBatchColumnName(pk)]
		if !ok {
			labeledParts = append(labeledParts, fmt.Sprintf("%s=<missing>", pk))
			rawValues = append(rawValues, "<missing>")
			continue
		}
		if cond.op != "" && !strings.EqualFold(cond.op, "=") {
			value := fmt.Sprintf("<op:%s:%s>", cond.op, cond.value.String())
			labeledParts = append(labeledParts, fmt.Sprintf("%s=%s", pk, value))
			rawValues = append(rawValues, value)
			continue
		}
		value := batchPrimaryKeyExprValueString(cond.value, params, markerIndex)
		labeledParts = append(labeledParts, fmt.Sprintf("%s=%s", pk, value))
		rawValues = append(rawValues, value)
	}
	return []batchPrimaryKeyDebugKey{{
		Prefix: batchPrimaryKeyPrefix(table, pkCols),
		Value:  strings.Join(rawValues, ","),
		Full:   fmt.Sprintf("%s(%s)", batchTableLogName(table), strings.Join(labeledParts, ", ")),
	}}
}

func batchPrimaryKeyPrefix(table *parser.TableName, pkCols []string) string {
	return fmt.Sprintf("%s(%s)", batchTableLogName(table), strings.Join(pkCols, ","))
}

func batchPrimaryKeySharedPrefix(inputKeys []batchPrimaryKeyDebugKey, mergedKeys []batchPrimaryKeyDebugGroupKeys) (string, bool) {
	prefix := ""
	visit := func(key batchPrimaryKeyDebugKey) bool {
		if key.Prefix == "" || key.Value == "" {
			return false
		}
		if prefix == "" {
			prefix = key.Prefix
			return true
		}
		return key.Prefix == prefix
	}

	for _, key := range inputKeys {
		if !visit(key) {
			return "", false
		}
	}
	for _, group := range mergedKeys {
		for _, key := range group.PrimaryKeys {
			if !visit(key) {
				return "", false
			}
		}
	}
	return prefix, prefix != ""
}

func batchPrimaryKeyValues(keys []batchPrimaryKeyDebugKey, compact bool) []string {
	out := make([]string, 0, len(keys))
	for _, key := range keys {
		if compact && key.Value != "" {
			out = append(out, key.Value)
			continue
		}
		out = append(out, key.Full)
	}
	return out
}

func batchPrimaryKeyGroupValues(groups []batchPrimaryKeyDebugGroupKeys, compact bool) []batchPrimaryKeyDebugGroup {
	out := make([]batchPrimaryKeyDebugGroup, 0, len(groups))
	for _, group := range groups {
		item := batchPrimaryKeyDebugGroup{
			ExecuteItemIndex:       group.ExecuteItemIndex,
			OriginalStatementCount: group.OriginalStatementCount,
			PrimaryKeys:            batchPrimaryKeyValues(group.PrimaryKeys, compact),
		}
		if group.GroupingStrategy != "" {
			item.GroupingStrategy = string(group.GroupingStrategy)
		}
		if group.BatchTemplateKey != "" {
			item.BatchTemplateKey = group.BatchTemplateKey
		}
		if group.GroupBreakReason != "" {
			item.GroupBreakReason = group.GroupBreakReason
		}
		if len(group.PartitionHits) > 0 {
			item.PartitionHits = append([]string(nil), group.PartitionHits...)
		}
		out = append(out, item)
	}
	return out
}

func batchPartitionGroupStrategyApplied(groups []coalescedBatchGroup) bool {
	for _, group := range groups {
		if group.groupingStrategy == batchGroupingStrategyPartitionGroup {
			return true
		}
	}
	return false
}

func batchMergeAppliedKeys(inputKeys, coalescedKeys []batchPrimaryKeyDebugKey) []batchPrimaryKeyDebugKey {
	if len(inputKeys) == 0 || len(coalescedKeys) == 0 {
		return nil
	}

	type keyCount struct {
		key   batchPrimaryKeyDebugKey
		count int
	}

	counts := make(map[string]keyCount, len(inputKeys))
	for _, key := range inputKeys {
		lookup := key.Full
		if lookup == "" {
			continue
		}
		entry := counts[lookup]
		if entry.count == 0 {
			entry.key = key
		}
		entry.count++
		counts[lookup] = entry
	}
	for _, key := range coalescedKeys {
		lookup := key.Full
		if lookup == "" {
			continue
		}
		entry, ok := counts[lookup]
		if !ok {
			continue
		}
		entry.count--
		counts[lookup] = entry
	}

	out := make([]batchPrimaryKeyDebugKey, 0)
	seen := make(map[string]struct{}, len(inputKeys))
	for _, key := range inputKeys {
		lookup := key.Full
		if lookup == "" {
			continue
		}
		if _, ok := seen[lookup]; ok {
			continue
		}
		seen[lookup] = struct{}{}
		if entry, ok := counts[lookup]; ok && entry.count > 0 {
			out = append(out, entry.key)
		}
	}
	return out
}

func batchPrimaryKeyExprValueString(expr parser.Expression, params []interface{}, markerIndex map[parser.Expression]int) string {
	if expr == nil {
		return "null"
	}

	switch e := expr.(type) {
	case *parser.TypeHintExpr:
		return batchPrimaryKeyExprValueString(e.Expr, params, markerIndex)
	case *parser.CastExpr:
		return batchPrimaryKeyExprValueString(e.Expr, params, markerIndex)
	case *parser.BindMarker, *parser.NamedBindMarker:
		value, _, err := resolveBindMarkerValue(expr, params, markerIndex)
		if err != nil {
			return fmt.Sprintf("<bind_error:%v>", err)
		}
		return formatBatchPrimaryKeyValue(value)
	case *parser.NullLiteral:
		return "null"
	case *parser.StringLiteral:
		return formatBatchPrimaryKeyValue(e.Value)
	case *parser.IntegerLiteral:
		return formatBatchPrimaryKeyValue(e.Value)
	case *parser.FloatLiteral:
		if e.Raw != "" {
			return e.Raw
		}
		return formatBatchPrimaryKeyValue(e.Value)
	case *parser.BooleanLiteral:
		return e.Value
	case *parser.UUIDLiteral:
		return formatBatchPrimaryKeyValue(e.Value)
	case *parser.BlobLiteral:
		return formatBatchPrimaryKeyValue(e.Value)
	default:
		return expr.String()
	}
}

func formatBatchPrimaryKeyValue(value interface{}) string {
	if value == nil {
		return "null"
	}
	encoded, err := json.Marshal(value)
	if err != nil {
		return fmt.Sprintf("%v", value)
	}
	return string(encoded)
}

func batchTableLogName(table *parser.TableName) string {
	if table == nil {
		return "<unknown_table>"
	}
	if table.Schema == "" {
		return table.Name
	}
	return table.Schema + "." + table.Name
}