package executor

import (
	"fmt"
	"strings"

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

type batchFoldEnvironment struct {
	resolvePK             batchUpdatePKResolver
	resolveColumnOrder    func(table *parser.TableName) ([]string, error)
	lookupColumnTypes     func(table *parser.TableName) (map[string]string, error)
	shouldInsertIncrement func(table *parser.TableName, column string) bool
}

type batchFoldColumnOpKind uint8

const (
	batchFoldColumnOpSet batchFoldColumnOpKind = iota
	batchFoldColumnOpDelta
)

type batchFoldColumnOp struct {
	kind  batchFoldColumnOpKind
	value interface{}
}

type batchFoldMutation struct {
	kind     batchMergeKind
	table    *parser.TableName
	pkCols   []string
	pkKey    string
	pkValues map[string]interface{}
	ops      map[string]batchFoldColumnOp
	opNames  map[string]string
	opOrder  []string
}

type batchFoldState struct {
	kind      batchMergeKind
	table     *parser.TableName
	pkCols    []string
	pkKey     string
	pkValues  map[string]interface{}
	ops       map[string]batchFoldColumnOp
	opNames   map[string]string
	opOrder   []string
	originals []BatchEntry
}

func foldBatchEntriesWithEnv(entries []BatchEntry, allowUpdateMerge bool, env batchFoldEnvironment) ([]BatchEntry, error) {
	if len(entries) == 0 {
		return nil, nil
	}

	out := make([]BatchEntry, 0, len(entries))
	pending := make(map[string]*batchFoldState)
	order := make([]string, 0, len(entries))
	var pendingTable batchTableKey
	havePendingTable := false

	flushPending := func() error {
		if len(order) == 0 {
			pending = make(map[string]*batchFoldState)
			havePendingTable = false
			return nil
		}
		for _, pkKey := range order {
			state := pending[pkKey]
			if state == nil {
				continue
			}
			entry, ok, err := materializeBatchFoldState(state, env)
			if err != nil {
				return err
			}
			if ok {
				out = append(out, entry)
			}
		}
		pending = make(map[string]*batchFoldState)
		order = order[:0]
		havePendingTable = false
		return nil
	}

	for _, entry := range entries {
		mutation, ok, barrier, noOp, err := normalizeBatchEntryForFold(entry, allowUpdateMerge, env)
		if err != nil {
			return nil, err
		}
		if noOp {
			continue
		}
		if barrier || !ok {
			if err := flushPending(); err != nil {
				return nil, err
			}
			out = append(out, annotateBatchEntryIncrementColumns(shallowBatchEntry(entry), env))
			continue
		}

		tableKey := batchTableKey{schema: mutation.table.Schema, table: mutation.table.Name}
		if havePendingTable && pendingTable != tableKey {
			if err := flushPending(); err != nil {
				return nil, err
			}
		}
		pendingTable = tableKey
		havePendingTable = true

		state := pending[mutation.pkKey]
		if state == nil {
			state = &batchFoldState{
				kind:      mutation.kind,
				table:     cloneTableName(mutation.table),
				pkCols:    append([]string(nil), mutation.pkCols...),
				pkKey:     mutation.pkKey,
				pkValues:  copyBatchFoldValueMap(mutation.pkValues),
				ops:       make(map[string]batchFoldColumnOp, len(mutation.ops)),
				opNames:   make(map[string]string, len(mutation.opNames)),
				originals: make([]BatchEntry, 0, 1),
			}
			pending[mutation.pkKey] = state
			order = append(order, mutation.pkKey)
		}
		if err := applyBatchFoldMutation(state, mutation, batchEntryOriginals(entry)); err != nil {
			if err := flushPending(); err != nil {
				return nil, err
			}
			out = append(out, annotateBatchEntryIncrementColumns(shallowBatchEntry(entry), env))
			continue
		}
	}

	if err := flushPending(); err != nil {
		return nil, err
	}
	return out, nil
}

func foldBatchEntriesGloballyWithEnv(entries []BatchEntry, allowUpdateMerge bool, env batchFoldEnvironment) ([]BatchEntry, batchMergeKind, bool, string, error) {
	if len(entries) == 0 {
		return nil, batchMergeKindInsert, true, "", nil
	}

	pending := make(map[string]*batchFoldState)
	order := make([]string, 0, len(entries))
	var (
		kindSet bool
		kind    batchMergeKind
	)

	for _, entry := range entries {
		mutation, ok, barrier, noOp, err := normalizeBatchEntryForFold(entry, allowUpdateMerge, env)
		if err != nil {
			return nil, 0, false, "", err
		}
		if noOp {
			continue
		}
		if barrier {
			return nil, 0, false, batchGlobalFoldFallbackBarrierStatement, nil
		}
		if !ok {
			return nil, 0, false, batchGlobalFoldFallbackNonFoldableStatement, nil
		}
		if !kindSet {
			kind = mutation.kind
			kindSet = true
		} else if kind != mutation.kind {
			return nil, 0, false, batchGlobalFoldFallbackMixedStatementKinds, nil
		}

		state := pending[mutation.pkKey]
		if state == nil {
			state = &batchFoldState{
				kind:      mutation.kind,
				table:     cloneTableName(mutation.table),
				pkCols:    append([]string(nil), mutation.pkCols...),
				pkKey:     mutation.pkKey,
				pkValues:  copyBatchFoldValueMap(mutation.pkValues),
				ops:       make(map[string]batchFoldColumnOp, len(mutation.ops)),
				opNames:   make(map[string]string, len(mutation.opNames)),
				originals: make([]BatchEntry, 0, 1),
			}
			pending[mutation.pkKey] = state
			order = append(order, mutation.pkKey)
		}
		if err := applyBatchFoldMutation(state, mutation, batchEntryOriginals(entry)); err != nil {
			return nil, 0, false, batchGlobalFoldFallbackFoldConflict, nil
		}
	}

	if !kindSet {
		return nil, batchMergeKindInsert, true, "", nil
	}

	out := make([]BatchEntry, 0, len(order))
	for _, pkKey := range order {
		state := pending[pkKey]
		if state == nil {
			continue
		}
		entry, ok, err := materializeBatchFoldState(state, env)
		if err != nil {
			return nil, 0, false, "", err
		}
		if ok {
			out = append(out, entry)
		}
	}
	return out, kind, true, "", nil
}

func annotateBatchEntryIncrementColumns(entry BatchEntry, env batchFoldEnvironment) BatchEntry {
	insert, ok := entry.Statement.(*parser.InsertStatement)
	if !ok || insert == nil || env.shouldInsertIncrement == nil {
		return entry
	}
	incrementCols := make(map[string]struct{})
	for _, col := range insert.Columns {
		if env.shouldInsertIncrement(insert.Table, col) {
			incrementCols[normalizeBatchColumnName(col)] = struct{}{}
		}
	}
	if len(incrementCols) == 0 {
		return entry
	}
	entry.incrementCols = incrementCols
	return entry
}

func normalizeBatchEntryForFold(entry BatchEntry, allowUpdateMerge bool, env batchFoldEnvironment) (batchFoldMutation, bool, bool, bool, error) {
	stmt := entry.Statement
	params := append([]interface{}(nil), entry.Params...)

	columnTypes, err := lookupBatchFoldColumnTypes(stmt, env)
	if err != nil {
		return batchFoldMutation{}, false, false, false, err
	}
	stmt, params, _, err = applyCastConversionsForDML(stmt, params, columnTypes)
	if err != nil {
		return batchFoldMutation{}, false, false, false, err
	}
	stmt, params, _, err = applyTimeConversionsForDML(stmt, params, columnTypes)
	if err != nil {
		return batchFoldMutation{}, false, false, false, err
	}

	switch s := stmt.(type) {
	case *parser.InsertStatement:
		if s == nil || len(s.Values) != 1 {
			return batchFoldMutation{}, false, true, false, nil
		}
		if _, mergeable := mergeableBatchInsert(stmt); !mergeable {
			return batchFoldMutation{}, false, true, false, nil
		}
		if hasConfiguredInsertDuplicateIncrement(s, env) {
			return batchFoldMutation{}, false, true, false, nil
		}
		if len(params) > 0 && hasUnset(params) {
			if env.resolvePK == nil {
				return batchFoldMutation{}, false, true, false, nil
			}
			pkCols, err := env.resolvePK(s.Table)
			if err != nil {
				return batchFoldMutation{}, false, true, false, nil
			}
			filtered, remapped, _, noOp, err := filterInsertStatementForUnset(s, params, buildBindMarkerIndex(s), pkCols)
			if err != nil {
				return batchFoldMutation{}, false, false, false, err
			}
			if noOp {
				return batchFoldMutation{}, false, false, true, nil
			}
			if filtered == nil || len(filtered.Values) != 1 {
				return batchFoldMutation{}, false, true, false, nil
			}
			s = filtered
			stmt = filtered
			params = remapped
		}
		return buildBatchFoldMutationFromInsert(s, params, env, batchMergeKindInsert, false)
	case *parser.UpdateStatement:
		if !allowUpdateMerge {
			return batchFoldMutation{}, false, true, false, nil
		}
		if len(params) > 0 && hasUnset(params) {
			filtered, remapped, _, noOp, err := filterUpdateStatementForUnset(s, params, buildBindMarkerIndex(s))
			if err != nil {
				return batchFoldMutation{}, false, false, false, err
			}
			if noOp {
				return batchFoldMutation{}, false, false, true, nil
			}
			if filtered == nil {
				return batchFoldMutation{}, false, true, false, nil
			}
			s = filtered
			stmt = filtered
			params = remapped
		}
		converted, ok, err := mergeableBatchUpdateAsInsert(stmt, env.resolvePK)
		if err != nil {
			return batchFoldMutation{}, false, false, false, err
		}
		if !ok {
			return batchFoldMutation{}, false, true, false, nil
		}
		return buildBatchFoldMutationFromInsert(converted, params, env, batchMergeKindUpdate, true)
	default:
		return batchFoldMutation{}, false, true, false, nil
	}
}

func lookupBatchFoldColumnTypes(stmt parser.Statement, env batchFoldEnvironment) (map[string]string, error) {
	if env.lookupColumnTypes == nil {
		return nil, nil
	}
	switch s := stmt.(type) {
	case *parser.InsertStatement:
		return env.lookupColumnTypes(s.Table)
	case *parser.UpdateStatement:
		return env.lookupColumnTypes(s.Table)
	case *parser.DeleteStatement:
		return env.lookupColumnTypes(s.From)
	default:
		return nil, nil
	}
}

func hasConfiguredInsertDuplicateIncrement(stmt *parser.InsertStatement, env batchFoldEnvironment) bool {
	if stmt == nil || env.shouldInsertIncrement == nil {
		return false
	}
	for _, col := range stmt.Columns {
		if env.shouldInsertIncrement(stmt.Table, col) {
			return true
		}
	}
	return false
}

func buildBatchFoldMutationFromInsert(stmt *parser.InsertStatement, params []interface{}, env batchFoldEnvironment, kind batchMergeKind, useStatementIncrementHints bool) (batchFoldMutation, bool, bool, bool, error) {
	if stmt == nil || stmt.Table == nil || len(stmt.Values) != 1 || env.resolvePK == nil {
		return batchFoldMutation{}, false, true, false, nil
	}

	pkCols, err := env.resolvePK(stmt.Table)
	if err != nil {
		return batchFoldMutation{}, false, true, false, nil
	}
	if len(pkCols) == 0 {
		return batchFoldMutation{}, false, true, false, nil
	}

	row := stmt.Values[0]
	if len(row) != len(stmt.Columns) {
		return batchFoldMutation{}, false, true, false, nil
	}

	columnIndex := make(map[string]int, len(stmt.Columns))
	columnNames := make(map[string]string, len(stmt.Columns))
	for i, col := range stmt.Columns {
		key := normalizeBatchColumnName(col)
		columnIndex[key] = i
		columnNames[key] = col
	}

	markerIndex := buildBindMarkerIndex(stmt)
	pkValues := make(map[string]interface{}, len(pkCols))
	for _, pk := range pkCols {
		idx, ok := columnIndex[normalizeBatchColumnName(pk)]
		if !ok || idx >= len(row) {
			return batchFoldMutation{}, false, true, false, nil
		}
		value, err := resolveScalarExpr(row[idx], params, markerIndex)
		if err != nil {
			return batchFoldMutation{}, false, true, false, nil
		}
		pkValues[normalizeBatchColumnName(pk)] = value
	}

	ops := make(map[string]batchFoldColumnOp)
	opNames := make(map[string]string)
	opOrder := make([]string, 0, len(stmt.Columns))
	pkSet := make(map[string]struct{}, len(pkCols))
	for _, pk := range pkCols {
		pkSet[normalizeBatchColumnName(pk)] = struct{}{}
	}
	for i, col := range stmt.Columns {
		key := normalizeBatchColumnName(col)
		if _, ok := pkSet[key]; ok {
			continue
		}
		value, err := resolveScalarExpr(row[i], params, markerIndex)
		if err != nil {
			return batchFoldMutation{}, false, true, false, nil
		}
		opKind := batchFoldColumnOpSet
		if useStatementIncrementHints {
			if _, ok := stmt.OnDuplicateIncrementColumns[key]; ok {
				opKind = batchFoldColumnOpDelta
			}
		}
		ops[key] = batchFoldColumnOp{kind: opKind, value: value}
		opNames[key] = columnNames[key]
		opOrder = append(opOrder, key)
	}

	return batchFoldMutation{
		kind:     kind,
		table:    cloneTableName(stmt.Table),
		pkCols:   append([]string(nil), pkCols...),
		pkKey:    batchFoldPrimaryKeyKey(stmt.Table, pkCols, pkValues),
		pkValues: pkValues,
		ops:      ops,
		opNames:  opNames,
		opOrder:  opOrder,
	}, true, false, false, nil
}

func applyBatchFoldMutation(state *batchFoldState, mutation batchFoldMutation, originals []BatchEntry) error {
	if state == nil {
		return fmt.Errorf("batch fold state is required")
	}
	if mutation.kind == batchMergeKindUpdate {
		state.kind = batchMergeKindUpdate
	}
	for _, key := range mutation.opOrder {
		next, ok := mutation.ops[key]
		if !ok {
			continue
		}
		if _, exists := state.ops[key]; !exists {
			state.opOrder = append(state.opOrder, key)
			if name := mutation.opNames[key]; name != "" {
				state.opNames[key] = name
			}
			state.ops[key] = next
			continue
		}

		prev := state.ops[key]
		switch next.kind {
		case batchFoldColumnOpSet:
			state.ops[key] = next
		case batchFoldColumnOpDelta:
			if prev.kind == batchFoldColumnOpDelta {
				value, err := accumulateBatchIncrementValue(prev.value, next.value)
				if err != nil {
					return err
				}
				state.ops[key] = batchFoldColumnOp{kind: batchFoldColumnOpDelta, value: value}
				continue
			}
			value, err := accumulateBatchIncrementValue(prev.value, next.value)
			if err != nil {
				return err
			}
			state.ops[key] = batchFoldColumnOp{kind: batchFoldColumnOpSet, value: value}
		default:
			return fmt.Errorf("unsupported batch fold operation kind %d", next.kind)
		}
	}
	state.originals = append(state.originals, originals...)
	return nil
}

func materializeBatchFoldState(state *batchFoldState, env batchFoldEnvironment) (BatchEntry, bool, error) {
	if state == nil || len(state.ops) == 0 {
		return BatchEntry{}, false, nil
	}

	var schemaOrder []string
	if env.resolveColumnOrder != nil {
		cols, err := env.resolveColumnOrder(state.table)
		if err == nil {
			schemaOrder = cols
		}
	}

	columns := make([]string, 0, len(state.ops)+len(state.pkCols))
	values := make([]parser.Expression, 0, len(state.ops)+len(state.pkCols))
	params := make([]interface{}, 0, len(state.ops)+len(state.pkCols))
	incrementCols := make(map[string]struct{})
	seen := make(map[string]struct{}, len(state.ops))
	pkSet := make(map[string]struct{}, len(state.pkCols))
	for _, pk := range state.pkCols {
		pkSet[normalizeBatchColumnName(pk)] = struct{}{}
	}

	for _, col := range schemaOrder {
		key := normalizeBatchColumnName(col)
		op, ok := state.ops[key]
		if !ok {
			continue
		}
		if _, isPK := pkSet[key]; isPK {
			continue
		}
		columns = append(columns, col)
		values = append(values, &parser.BindMarker{})
		params = append(params, op.value)
		if op.kind == batchFoldColumnOpDelta {
			incrementCols[key] = struct{}{}
		}
		seen[key] = struct{}{}
	}
	for _, key := range state.opOrder {
		if _, ok := seen[key]; ok {
			continue
		}
		op, ok := state.ops[key]
		if !ok {
			continue
		}
		if _, isPK := pkSet[key]; isPK {
			continue
		}
		name := state.opNames[key]
		if name == "" {
			name = key
		}
		columns = append(columns, name)
		values = append(values, &parser.BindMarker{})
		params = append(params, op.value)
		if op.kind == batchFoldColumnOpDelta {
			incrementCols[key] = struct{}{}
		}
	}
	for _, pk := range state.pkCols {
		columns = append(columns, pk)
		values = append(values, &parser.BindMarker{})
		params = append(params, state.pkValues[normalizeBatchColumnName(pk)])
	}
	if len(columns) == len(state.pkCols) {
		return BatchEntry{}, false, nil
	}

	insert := &parser.InsertStatement{
		Table:   cloneTableName(state.table),
		Columns: columns,
		Values:  [][]parser.Expression{values},
	}
	if len(incrementCols) > 0 {
		insert.OnDuplicateIncrementColumns = incrementCols
	}
	return BatchEntry{
		Statement:     insert,
		Params:        params,
		originals:     append([]BatchEntry(nil), state.originals...),
		mergeKind:     state.kind,
		folded:        true,
		incrementCols: cloneBatchColumnSet(insert.OnDuplicateIncrementColumns),
	}, true, nil
}

func copyBatchFoldValueMap(src map[string]interface{}) map[string]interface{} {
	if len(src) == 0 {
		return nil
	}
	cloned := make(map[string]interface{}, len(src))
	for key, value := range src {
		cloned[key] = value
	}
	return cloned
}

func cloneTableName(table *parser.TableName) *parser.TableName {
	if table == nil {
		return nil
	}
	cloned := *table
	return &cloned
}

func batchFoldPrimaryKeyKey(table *parser.TableName, pkCols []string, pkValues map[string]interface{}) string {
	parts := make([]string, 0, len(pkCols))
	for _, pk := range pkCols {
		parts = append(parts, normalizeBatchColumnName(pk)+"="+formatBatchPrimaryKeyValue(pkValues[normalizeBatchColumnName(pk)]))
	}
	return batchTableLogName(table) + "|" + strings.Join(parts, "|")
}