// UpdateExpression Parser and Executor
// Supports: SET path = value [, path = value ...]
//           REMOVE path [, path ...]
//           ADD path value [, path value ...]    (numeric add / set add)
//           DELETE path value [, path value ...]  (set remove)
//
// Values can be:
//   :placeholder       → resolved from ExpressionAttributeValues
//   path + :placeholder → numeric addition
//   path - :placeholder → numeric subtraction
//   if_not_exists(path, :placeholder) → default value
//   list_append(operand, operand)     → list concatenation
package dynamodb

import "core:encoding/json"
import "core:fmt"
import "core:strconv"
import "core:strings"

// ============================================================================
// Update Plan — parsed representation of an UpdateExpression
// ============================================================================

Update_Action_Type :: enum {
	SET,
	REMOVE,
	ADD,
	DELETE,
}

Set_Value_Kind :: enum {
	Direct,           // SET x = :val
	Plus,             // SET x = x + :val  or  SET x = :val + x
	Minus,            // SET x = x - :val
	If_Not_Exists,    // SET x = if_not_exists(x, :val)
	List_Append,      // SET x = list_append(x, :val)
}

Set_Action :: struct {
	path:       string,
	value_kind: Set_Value_Kind,
	value:      Attribute_Value,          // primary value
	source:     string,                   // source path for Plus/Minus/If_Not_Exists/List_Append
	value2:     Maybe(Attribute_Value),   // second operand for list_append where both are values
}

Remove_Action :: struct {
	path: string,
}

Add_Action :: struct {
	path:  string,
	value: Attribute_Value,
}

Delete_Action :: struct {
	path:  string,
	value: Attribute_Value,
}

Update_Plan :: struct {
	sets:    [dynamic]Set_Action,
	removes: [dynamic]Remove_Action,
	adds:    [dynamic]Add_Action,
	deletes: [dynamic]Delete_Action,
}

update_plan_destroy :: proc(plan: ^Update_Plan) {
	for &s in plan.sets {
		attr_value_destroy(&s.value)
		if v2, ok := s.value2.?; ok {
			v2_copy := v2
			attr_value_destroy(&v2_copy)
		}
	}
	delete(plan.sets)

	delete(plan.removes)

	for &a in plan.adds {
		attr_value_destroy(&a.value)
	}
	delete(plan.adds)

	for &d in plan.deletes {
		attr_value_destroy(&d.value)
	}
	delete(plan.deletes)
}

// ============================================================================
// Parse UpdateExpression
//
// Grammar (simplified):
//   update_expr = clause { clause }
//   clause      = "SET" set_list | "REMOVE" remove_list | "ADD" add_list | "DELETE" delete_list
//   set_list    = set_entry { "," set_entry }
//   set_entry   = path "=" value_expr
//   value_expr  = :placeholder
//              | path "+" :placeholder
//              | path "-" :placeholder
//              | "if_not_exists" "(" path "," :placeholder ")"
//              | "list_append" "(" operand "," operand ")"
//   remove_list = path { "," path }
//   add_list    = add_entry { "," add_entry }
//   add_entry   = path :placeholder
//   delete_list = delete_entry { "," delete_entry }
//   delete_entry= path :placeholder
// ============================================================================

parse_update_expression :: proc(
	expression: string,
	attribute_names: Maybe(map[string]string),
	attribute_values: map[string]Attribute_Value,
) -> (plan: Update_Plan, ok: bool) {
	plan.sets    = make([dynamic]Set_Action)
	plan.removes = make([dynamic]Remove_Action)
	plan.adds    = make([dynamic]Add_Action)
	plan.deletes = make([dynamic]Delete_Action)

	t := tokenizer_init(expression)

	for {
		keyword_maybe := tokenizer_next(&t)
		keyword_str, has_keyword := keyword_maybe.?
		if !has_keyword {
			break // done
		}

		if strings.equal_fold(keyword_str, "SET") {
			if !parse_set_clause(&t, &plan, attribute_names, attribute_values) {
				update_plan_destroy(&plan)
				return {}, false
			}
		} else if strings.equal_fold(keyword_str, "REMOVE") {
			if !parse_remove_clause(&t, &plan, attribute_names) {
				update_plan_destroy(&plan)
				return {}, false
			}
		} else if strings.equal_fold(keyword_str, "ADD") {
			if !parse_add_clause(&t, &plan, attribute_names, attribute_values) {
				update_plan_destroy(&plan)
				return {}, false
			}
		} else if strings.equal_fold(keyword_str, "DELETE") {
			if !parse_delete_clause(&t, &plan, attribute_names, attribute_values) {
				update_plan_destroy(&plan)
				return {}, false
			}
		} else {
			update_plan_destroy(&plan)
			return {}, false
		}
	}

	return plan, true
}

// ============================================================================
// SET clause parsing
// ============================================================================

parse_set_clause :: proc(
	t: ^Tokenizer,
	plan: ^Update_Plan,
	names: Maybe(map[string]string),
	values: map[string]Attribute_Value,
) -> bool {
	saved_pos: int
	for {
		// Save position before reading so we can rewind if it's a clause keyword
		saved_pos = t.pos

		// Path
		path_tok, path_ok := next_token(t)
		if !path_ok {
			return false
		}

		// Check if this is actually a new clause keyword (SET/REMOVE/ADD/DELETE)
		if is_clause_keyword(path_tok) {
			t.pos = saved_pos
			return true
		}

		path, path_resolved := resolve_attribute_name(path_tok, names)
		if !path_resolved {
			return false
		}

		// "="
		eq_tok, eq_ok := next_token(t)
		if !eq_ok || eq_tok != "=" {
			return false
		}

		// Value expression
		action, act_ok := parse_set_value_expr(t, path, names, values)
		if !act_ok {
			return false
		}

		append(&plan.sets, action)

		// Check for comma (more entries) or end
		saved_pos = t.pos
		comma_maybe := tokenizer_next(t)
		if comma, has := comma_maybe.?; has {
			if comma == "," {
				continue
			}
			// Not a comma — put it back
			t.pos = saved_pos
		}
		break
	}
	return true
}

parse_set_value_expr :: proc(
	t: ^Tokenizer,
	path: string,
	names: Maybe(map[string]string),
	values: map[string]Attribute_Value,
) -> (action: Set_Action, ok: bool) {
	first_tok, first_ok := next_token(t)
	if !first_ok {
		return {}, false
	}

	// Check for if_not_exists(...)
	if strings.equal_fold(first_tok, "if_not_exists") {
		action, ok = parse_if_not_exists(t, path, names, values)
		return
	}

	// Check for list_append(...)
	if strings.equal_fold(first_tok, "list_append") {
		action, ok = parse_list_append(t, path, names, values)
		return
	}

	peek_pos: int

	// Check if first token is a :placeholder (direct value)
	if len(first_tok) > 0 && first_tok[0] == ':' {
		// Could be :val + path or :val - path or just :val
		peek_pos = t.pos
		op_maybe := tokenizer_next(t)
		if op, has_op := op_maybe.?; has_op && (op == "+" || op == "-") {
			// :val op path
			second_tok, sec_ok := next_token(t)
			if !sec_ok {
				return {}, false
			}
			source, source_resolved := resolve_attribute_name(second_tok, names)
			if !source_resolved {
				return {}, false
			}
			val, val_ok := resolve_attribute_value(first_tok, values)
			if !val_ok {
				return {}, false
			}
			kind := Set_Value_Kind.Plus if op == "+" else Set_Value_Kind.Minus
			return Set_Action{
				path       = path,
				value_kind = kind,
				value      = val,
				source     = source,
			}, true
		}
		// Just a direct value
		t.pos = peek_pos
		val, val_ok := resolve_attribute_value(first_tok, values)
		if !val_ok {
			return {}, false
		}
		return Set_Action{
			path       = path,
			value_kind = .Direct,
			value      = val,
		}, true
	}

	// First token is a path — check for path + :val or path - :val
	source, source_resolved := resolve_attribute_name(first_tok, names)
	if !source_resolved {
		return {}, false
	}
	peek_pos = t.pos
	op_maybe := tokenizer_next(t)
	if op, has_op := op_maybe.?; has_op && (op == "+" || op == "-") {
		val_tok, vt_ok := next_token(t)
		if !vt_ok {
			return {}, false
		}
		val, val_ok := resolve_attribute_value(val_tok, values)
		if !val_ok {
			return {}, false
		}
		kind := Set_Value_Kind.Plus if op == "+" else Set_Value_Kind.Minus
		return Set_Action{
			path       = path,
			value_kind = kind,
			value      = val,
			source     = source,
		}, true
	}
	// Just a path reference — treat as direct copy (SET a = b)
	t.pos = peek_pos
	return {}, false
}

parse_if_not_exists :: proc(
	t: ^Tokenizer,
	path: string,
	names: Maybe(map[string]string),
	values: map[string]Attribute_Value,
) -> (action: Set_Action, ok: bool) {
	lparen, lp_ok := next_token(t)
	if !lp_ok || lparen != "(" {
		return {}, false
	}

	src_tok, src_ok := next_token(t)
	if !src_ok {
		return {}, false
	}
	source, source_resolved := resolve_attribute_name(src_tok, names)
	if !source_resolved {
		return {}, false
	}

	comma, comma_ok := next_token(t)
	if !comma_ok || comma != "," {
		return {}, false
	}

	val_tok, vt_ok := next_token(t)
	if !vt_ok {
		return {}, false
	}
	val, val_ok := resolve_attribute_value(val_tok, values)
	if !val_ok {
		return {}, false
	}

	rparen, rp_ok := next_token(t)
	if !rp_ok || rparen != ")" {
		attr_value_destroy(&val)
		return {}, false
	}

	return Set_Action{
		path       = path,
		value_kind = .If_Not_Exists,
		value      = val,
		source     = source,
	}, true
}

parse_list_append :: proc(
	t: ^Tokenizer,
	path: string,
	names: Maybe(map[string]string),
	values: map[string]Attribute_Value,
) -> (action: Set_Action, ok: bool) {
	lparen, lp_ok := next_token(t)
	if !lp_ok || lparen != "(" {
		return {}, false
	}

	// First operand — could be :val or path
	first_tok, first_ok := next_token(t)
	if !first_ok {
		return {}, false
	}

	comma, comma_ok := next_token(t)
	if !comma_ok || comma != "," {
		return {}, false
	}

	// Second operand
	second_tok, second_ok := next_token(t)
	if !second_ok {
		return {}, false
	}

	rparen, rp_ok := next_token(t)
	if !rp_ok || rparen != ")" {
		return {}, false
	}

	// Determine which is the path and which is the value
	// Common patterns: list_append(path, :val) or list_append(:val, path)
	source: string
	val: Attribute_Value
	resolved: bool

	if len(first_tok) > 0 && first_tok[0] == ':' {
		// list_append(:val, path)
		v, v_ok := resolve_attribute_value(first_tok, values)
		if !v_ok {
			return {}, false
		}
		val = v
		source, resolved = resolve_attribute_name(second_tok, names)
		if !resolved {
			return {}, false
		}
	} else if len(second_tok) > 0 && second_tok[0] == ':' {
		// list_append(path, :val)
		source, resolved = resolve_attribute_name(first_tok, names)
		if !resolved {
			return {}, false
		}
		v, v_ok := resolve_attribute_value(second_tok, values)
		if !v_ok {
			return {}, false
		}
		val = v
	} else {
		return {}, false
	}

	return Set_Action{
		path       = path,
		value_kind = .List_Append,
		value      = val,
		source     = source,
	}, true
}

// ============================================================================
// REMOVE clause parsing
// ============================================================================

parse_remove_clause :: proc(
	t: ^Tokenizer,
	plan: ^Update_Plan,
	names: Maybe(map[string]string),
) -> bool {
	saved_pos: int
	for {
		saved_pos = t.pos

		path_tok, path_ok := next_token(t)
		if !path_ok {
			return false
		}

		if is_clause_keyword(path_tok) {
			t.pos = saved_pos
			return true
		}

		path, path_resolved := resolve_attribute_name(path_tok, names)
		if !path_resolved {
			return false
		}

		append(&plan.removes, Remove_Action{path = path})

		saved_pos = t.pos
		comma_maybe := tokenizer_next(t)
		if comma, has := comma_maybe.?; has {
			if comma == "," {
				continue
			}
			t.pos = saved_pos
		}
		break
	}
	return true
}

// ============================================================================
// ADD clause parsing
// ============================================================================

parse_add_clause :: proc(
	t: ^Tokenizer,
	plan: ^Update_Plan,
	names: Maybe(map[string]string),
	values: map[string]Attribute_Value,
) -> bool {
	saved_pos: int
	for {
		saved_pos = t.pos

		path_tok, path_ok := next_token(t)
		if !path_ok {
			return false
		}

		if is_clause_keyword(path_tok) {
			t.pos = saved_pos
			return true
		}

		path, path_resolved := resolve_attribute_name(path_tok, names)
		if !path_resolved {
			return false
		}

		val_tok, vt_ok := next_token(t)
		if !vt_ok {
			return false
		}
		val, val_ok := resolve_attribute_value(val_tok, values)
		if !val_ok {
			return false
		}

		append(&plan.adds, Add_Action{path = path, value = val})

		saved_pos = t.pos
		comma_maybe := tokenizer_next(t)
		if comma, has := comma_maybe.?; has {
			if comma == "," {
				continue
			}
			t.pos = saved_pos
		}
		break
	}
	return true
}

// ============================================================================
// DELETE clause parsing
// ============================================================================

parse_delete_clause :: proc(
	t: ^Tokenizer,
	plan: ^Update_Plan,
	names: Maybe(map[string]string),
	values: map[string]Attribute_Value,
) -> bool {
	saved_pos: int
	for {
		saved_pos = t.pos

		path_tok, path_ok := next_token(t)
		if !path_ok {
			return false
		}

		if is_clause_keyword(path_tok) {
			t.pos = saved_pos
			return true
		}

		path, path_resolved := resolve_attribute_name(path_tok, names)
		if !path_resolved {
			return false
		}

		val_tok, vt_ok := next_token(t)
		if !vt_ok {
			return false
		}
		val, val_ok := resolve_attribute_value(val_tok, values)
		if !val_ok {
			return false
		}

		append(&plan.deletes, Delete_Action{path = path, value = val})

		saved_pos = t.pos
		comma_maybe := tokenizer_next(t)
		if comma, has := comma_maybe.?; has {
			if comma == "," {
				continue
			}
			t.pos = saved_pos
		}
		break
	}
	return true
}

// ============================================================================
// Helpers
// ============================================================================

is_clause_keyword :: proc(tok: string) -> bool {
	return strings.equal_fold(tok, "SET") ||
	       strings.equal_fold(tok, "REMOVE") ||
	       strings.equal_fold(tok, "ADD") ||
	       strings.equal_fold(tok, "DELETE")
}

// ============================================================================
// Execute Update Plan — apply mutations to an Item (in-place)
// ============================================================================

execute_update_plan :: proc(item: ^Item, plan: ^Update_Plan) -> bool {
	// Execute SET actions
	for &action in plan.sets {
		switch action.value_kind {
		case .Direct:
			// Remove old value if exists
			if old, found := item[action.path]; found {
				old_copy := old
				attr_value_destroy(&old_copy)
				delete_key(item, action.path)
			}
			item[strings.clone(action.path)] = attr_value_deep_copy(action.value)

		case .Plus:
			// Numeric addition: path = source + value  or  path = value + source
			existing: Attribute_Value
			if src, found := item[action.source]; found {
				existing = src
			} else {
				return false // source attribute not found
			}
			result, add_ok := numeric_add(existing, action.value)
			if !add_ok {
				return false
			}
			if old, found := item[action.path]; found {
				old_copy := old
				attr_value_destroy(&old_copy)
				delete_key(item, action.path)
			}
			item[strings.clone(action.path)] = result

		case .Minus:
			// Numeric subtraction: path = source - value
			existing: Attribute_Value
			if src, found := item[action.source]; found {
				existing = src
			} else {
				return false
			}
			result, sub_ok := numeric_subtract(existing, action.value)
			if !sub_ok {
				return false
			}
			if old, found := item[action.path]; found {
				old_copy := old
				attr_value_destroy(&old_copy)
				delete_key(item, action.path)
			}
			item[strings.clone(action.path)] = result

		case .If_Not_Exists:
			// Only set if attribute doesn't exist
			if _, found := item[action.source]; !found {
				if old, found2 := item[action.path]; found2 {
					old_copy := old
					attr_value_destroy(&old_copy)
					delete_key(item, action.path)
				}
				item[strings.clone(action.path)] = attr_value_deep_copy(action.value)
			}
			// If attribute exists, do nothing (keep current value)

		case .List_Append:
			// Append to list
			existing_list: []Attribute_Value
			if src, found := item[action.source]; found {
				if l, is_list := src.(List); is_list {
					existing_list = ([]Attribute_Value)(l)
				} else {
					return false
				}
			} else {
				existing_list = {}
			}

			append_list: []Attribute_Value
			if l, is_list := action.value.(List); is_list {
				append_list = ([]Attribute_Value)(l)
			} else {
				return false
			}

			new_list := make([]Attribute_Value, len(existing_list) + len(append_list))
			for item_val, i in existing_list {
				new_list[i] = attr_value_deep_copy(item_val)
			}
			for item_val, i in append_list {
				new_list[len(existing_list) + i] = attr_value_deep_copy(item_val)
			}

			if old, found := item[action.path]; found {
				old_copy := old
				attr_value_destroy(&old_copy)
				delete_key(item, action.path)
			}
			item[strings.clone(action.path)] = List(new_list)
		}
	}

	// Execute REMOVE actions
	for &action in plan.removes {
		if old, found := item[action.path]; found {
			old_copy := old
			attr_value_destroy(&old_copy)
			delete_key(item, action.path)
		}
	}

	// Execute ADD actions
	for &action in plan.adds {
		if existing, found := item[action.path]; found {
			// If existing is a number, add numerically
			#partial switch v in existing {
			case Number:
				result, add_ok := numeric_add(existing, action.value)
				if !add_ok {
					return false
				}
				old_copy := existing
				attr_value_destroy(&old_copy)
				delete_key(item, action.path)
				item[strings.clone(action.path)] = result

			case String_Set:
				// Add elements to string set
				if new_ss, is_ss := action.value.(String_Set); is_ss {
					merged := set_union_strings(([]string)(v), ([]string)(new_ss))
					old_copy := existing
					attr_value_destroy(&old_copy)
					delete_key(item, action.path)
					item[strings.clone(action.path)] = String_Set(merged)
				} else {
					return false
				}

			case Number_Set:
				if new_ns, is_ns := action.value.(Number_Set); is_ns {
					merged := set_union_strings(([]string)(v), ([]string)(new_ns))
					old_copy := existing
					attr_value_destroy(&old_copy)
					delete_key(item, action.path)
					item[strings.clone(action.path)] = Number_Set(merged)
				} else {
					return false
				}

			case:
				return false
			}
		} else {
			// Attribute doesn't exist — create it
			item[strings.clone(action.path)] = attr_value_deep_copy(action.value)
		}
	}

	// Execute DELETE actions (remove elements from sets)
	for &action in plan.deletes {
		if existing, found := item[action.path]; found {
			#partial switch v in existing {
			case String_Set:
				if del_ss, is_ss := action.value.(String_Set); is_ss {
					result := set_difference_strings(([]string)(v), ([]string)(del_ss))
					old_copy := existing
					attr_value_destroy(&old_copy)
					delete_key(item, action.path)
					if len(result) > 0 {
						item[strings.clone(action.path)] = String_Set(result)
					} else {
						delete(result)
					}
				}

			case Number_Set:
				if del_ns, is_ns := action.value.(Number_Set); is_ns {
					result := set_difference_strings(([]string)(v), ([]string)(del_ns))
					old_copy := existing
					attr_value_destroy(&old_copy)
					delete_key(item, action.path)
					if len(result) > 0 {
						item[strings.clone(action.path)] = Number_Set(result)
					} else {
						delete(result)
					}
				}

			case:
				// DELETE on non-set type is a no-op in DynamoDB
			}
		}
	}

	return true
}

// ============================================================================
// Numeric helpers
// ============================================================================

numeric_add :: proc(a: Attribute_Value, b: Attribute_Value) -> (Attribute_Value, bool) {
	a_num, a_ok := a.(DDB_Number)
	b_num, b_ok := b.(DDB_Number)
	if !a_ok || !b_ok {
		return nil, false
	}

	result, result_ok := add_ddb_numbers(a_num, b_num)
	if !result_ok {
		return nil, false
	}
	return result, true
}

numeric_subtract :: proc(a: Attribute_Value, b: Attribute_Value) -> (Attribute_Value, bool) {
	a_num, a_ok := a.(Number)
	b_num, b_ok := b.(Number)
	if !a_ok || !b_ok {
		return nil, false
	}

	a_val, a_parse := strconv.parse_f64(string(a_num))
	b_val, b_parse := strconv.parse_f64(string(b_num))
	if !a_parse || !b_parse {
		return nil, false
	}

	result := a_val - b_val
	result_str := format_number(result)
	return Number(result_str), true
}

format_number :: proc(val: f64) -> string {
	// If it's an integer, format without decimal point
	int_val := i64(val)
	if f64(int_val) == val {
		return fmt.aprintf("%d", int_val)
	}
	return fmt.aprintf("%g", val)
}

// ============================================================================
// Set helpers
// ============================================================================

set_union_strings :: proc(a: []string, b: []string) -> []string {
	seen := make(map[string]bool, allocator = context.temp_allocator)
	for s in a {
		seen[s] = true
	}
	for s in b {
		seen[s] = true
	}

	result := make([]string, len(seen))
	i := 0
	for s in seen {
		result[i] = strings.clone(s)
		i += 1
	}
	return result
}

set_difference_strings :: proc(a: []string, b: []string) -> []string {
	to_remove := make(map[string]bool, allocator = context.temp_allocator)
	for s in b {
		to_remove[s] = true
	}

	result := make([dynamic]string)
	for s in a {
		if !(s in to_remove) {
			append(&result, strings.clone(s))
		}
	}
	return result[:]
}

// ============================================================================
// Request Parsing Helper
// ============================================================================

parse_update_expression_string :: proc(request_body: []byte) -> (expr: string, ok: bool) {
	data, parse_err := json.parse(request_body, allocator = context.temp_allocator)
	if parse_err != nil {
		return
	}
	defer json.destroy_value(data)

	root, root_ok := data.(json.Object)
	if !root_ok {
		return
	}

	ue_val, found := root["UpdateExpression"]
	if !found {
		return
	}

	ue_str, str_ok := ue_val.(json.String)
	if !str_ok {
		return
	}

	expr = string(ue_str)
	ok = true
	return
}

// Parse ReturnValues from request body
parse_return_values :: proc(request_body: []byte) -> string {
	data, parse_err := json.parse(request_body, allocator = context.temp_allocator)
	if parse_err != nil {
		return "NONE"
	}
	defer json.destroy_value(data)

	root, root_ok := data.(json.Object)
	if !root_ok {
		return "NONE"
	}

	rv_val, found := root["ReturnValues"]
	if !found {
		return "NONE"
	}

	rv_str, str_ok := rv_val.(json.String)
	if !str_ok {
		return "NONE"
	}

	return string(rv_str)
}