package executor

import (
	"strings"

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

func applyTimeConversionsForDML(stmt parser.Statement, params []interface{}, columnTypes map[string]string) (parser.Statement, []interface{}, bool, error) {
	if stmt == nil {
		return stmt, params, false, nil
	}
	markerIndex := buildBindMarkerIndex(stmt)
	lookup := func(name string) (string, bool) {
		if columnTypes == nil {
			return "", false
		}
		typ, ok := columnTypes[strings.ToLower(name)]
		if !ok {
			return "", false
		}
		return normalizeCastType(typ), true
	}

	changed := false
	switch s := stmt.(type) {
	case *parser.InsertStatement:
		out := *s
		if s.JSON {
			if s.JSONValue != nil {
				expr, ch, _, err := rewriteTimeFunctionsInExpr(s.JSONValue, params, markerIndex, lookup, false)
				if err != nil {
					return nil, nil, false, err
				}
				if ch {
					out.JSONValue = expr
					changed = true
				}
			}
		} else if len(s.Values) > 0 {
			values := make([][]parser.Expression, len(s.Values))
			for i, row := range s.Values {
				if len(row) == 0 {
					continue
				}
				newRow := make([]parser.Expression, len(row))
				rowChanged := false
				for j, expr := range row {
					updated, ch, _, err := rewriteTimeFunctionsInExpr(expr, params, markerIndex, lookup, false)
					if err != nil {
						return nil, nil, false, err
					}
					if ch {
						rowChanged = true
					}
					newRow[j] = updated
				}
				if rowChanged {
					changed = true
					values[i] = newRow
				} else {
					values[i] = row
				}
			}
			if changed {
				out.Values = values
			}
		}
		if s.TTL != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.TTL, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.TTL = expr
				changed = true
			}
		}
		if s.Timestamp != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.Timestamp, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.Timestamp = expr
				changed = true
			}
		}
		if changed {
			return &out, params, true, nil
		}
		return stmt, params, false, nil
	case *parser.UpdateStatement:
		out := *s
		if s.TTL != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.TTL, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.TTL = expr
				changed = true
			}
		}
		if s.Timestamp != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.Timestamp, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.Timestamp = expr
				changed = true
			}
		}
		if len(s.Assigments) > 0 {
			assignments := make([]*parser.Assignment, len(s.Assigments))
			for i, assignment := range s.Assigments {
				if assignment == nil {
					continue
				}
				target, targetChanged, _, err := rewriteTimeFunctionsInExpr(assignment.Target, params, markerIndex, lookup, false)
				if err != nil {
					return nil, nil, false, err
				}
				value, valueChanged, _, err := rewriteTimeFunctionsInExpr(assignment.Value, params, markerIndex, lookup, false)
				if err != nil {
					return nil, nil, false, err
				}
				if targetChanged || valueChanged {
					assignments[i] = &parser.Assignment{
						Target: target,
						Column: assignment.Column,
						Value:  value,
						Op:     assignment.Op,
					}
					changed = true
				} else {
					assignments[i] = assignment
				}
			}
			if changed {
				out.Assigments = assignments
			}
		}
		if s.Where != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.Where, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.Where = expr
				changed = true
			}
		}
		if s.IfCondition != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.IfCondition, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.IfCondition = expr
				changed = true
			}
		}
		if changed {
			return &out, params, true, nil
		}
		return stmt, params, false, nil
	case *parser.DeleteStatement:
		out := *s
		if len(s.Targets) > 0 {
			targets := make([]parser.Expression, len(s.Targets))
			for i, target := range s.Targets {
				updated, ch, _, err := rewriteTimeFunctionsInExpr(target, params, markerIndex, lookup, false)
				if err != nil {
					return nil, nil, false, err
				}
				if ch {
					changed = true
				}
				targets[i] = updated
			}
			if changed {
				out.Targets = targets
			}
		}
		if s.Where != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.Where, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.Where = expr
				changed = true
			}
		}
		if s.IfCondition != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.IfCondition, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.IfCondition = expr
				changed = true
			}
		}
		if s.Timestamp != nil {
			expr, ch, _, err := rewriteTimeFunctionsInExpr(s.Timestamp, params, markerIndex, lookup, false)
			if err != nil {
				return nil, nil, false, err
			}
			if ch {
				out.Timestamp = expr
				changed = true
			}
		}
		if changed {
			return &out, params, true, nil
		}
		return stmt, params, false, nil
	default:
		return stmt, params, false, nil
	}
}