module transform

import v3.flat

// transform_for_body transforms transform for body data for transform.
fn (mut t Transformer) transform_for_body(id flat.NodeId, node flat.Node) []flat.NodeId {
	if node.children_count < 3 {
		return arr1(id)
	}
	// child 0: init statement
	init_id := t.a.child(&node, 0)
	mut new_init := init_id
	if int(init_id) >= 0 {
		expanded := t.transform_stmt(init_id)
		if expanded.len > 0 {
			new_init = expanded[0]
		}
	}
	// child 1: condition expression
	cond_id := t.a.child(&node, 1)
	cond_smartcasts := t.extract_all_is_exprs(cond_id)
	new_cond := t.transform_and_chain_smartcasts(cond_id)
	mut cond_prefix := []flat.NodeId{}
	t.drain_pending(mut cond_prefix)
	// child 2: post statement
	post_id := t.a.child(&node, 2)
	mut new_post := post_id
	if int(post_id) >= 0 {
		expanded := t.transform_stmt(post_id)
		if expanded.len > 0 {
			new_post = expanded[0]
		}
	}
	// children 3..n: body statements
	mut body_ids := []flat.NodeId{}
	for i in 3 .. node.children_count {
		body_ids << t.a.child(&node, i)
	}
	for info in cond_smartcasts {
		t.push_smartcast(info.expr_name, info.variant_name, info.sum_type_name)
	}
	new_body := t.transform_stmts(body_ids)
	for _ in cond_smartcasts {
		t.pop_smartcast()
	}
	mut loop_cond := new_cond
	mut loop_body := new_body.clone()
	if cond_prefix.len > 0 {
		mut guarded_body := []flat.NodeId{cap: cond_prefix.len + new_body.len + 1}
		for stmt in cond_prefix {
			guarded_body << stmt
		}
		not_cond := t.make_prefix(.not, t.make_paren(new_cond))
		break_block := t.make_block(arr1(t.a.add(.break_stmt)))
		guarded_body << t.make_if(not_cond, break_block, t.make_block([]flat.NodeId{}))
		for stmt in new_body {
			guarded_body << stmt
		}
		loop_cond = t.make_bool_literal(true)
		loop_body = guarded_body.clone()
	}
	// Rebuild the for_stmt with transformed children
	start := t.a.children.len
	t.a.children << new_init
	t.a.children << loop_cond
	t.a.children << new_post
	for bid in loop_body {
		t.a.children << bid
	}
	count := t.a.children.len - start
	new_id := t.a.add_node(flat.Node{
		kind:           .for_stmt
		op:             node.op
		children_start: start
		children_count: flat.child_count(count)
		pos:            node.pos
		value:          node.value
		typ:            node.typ
	})
	return arr1(new_id)
}

// transform_for_in_body transforms transform for in body data for transform.
fn (mut t Transformer) transform_for_in_body(id flat.NodeId, node flat.Node) []flat.NodeId {
	header_count := node.value.int()
	if header_count < 3 || node.children_count < 3 {
		return arr1(id)
	}
	if t.cur_fn_is_generic {
		return arr1(id)
	}
	key_id := t.a.child(&node, 0) // loop var ident — pass through (do not transform a binding)
	val_id := t.a.child(&node, 1) // may be flat.empty_node (-1)
	container_id := t.a.child(&node, 2)
	iter_type := t.detect_for_in_type(node)
	has_index := int(val_id) >= 0
	container_is_range := if int(container_id) >= 0 {
		t.a.nodes[int(container_id)].kind == .range
	} else {
		false
	}
	if header_count == 4 {
		range_end_id := t.a.child(&node, 3)
		body_ids := t.a.children_of(&node)[header_count..].clone()
		return t.lower_range_for_in(id, node, key_id, container_id, range_end_id, body_ids)
	}
	if container_is_range {
		range_node := t.a.nodes[int(container_id)]
		if range_node.children_count >= 2 {
			body_ids := t.a.children_of(&node)[header_count..].clone()
			return t.lower_range_for_in(id, node, key_id, t.a.child(&range_node, 0), t.a.child(&range_node,
				1), body_ids)
		}
	}
	if iter_type.starts_with('map[') {
		return t.rebuild_for_in_stmt(id, node)
	}
	if pool_iter := t.pool_get_results_iter_type(container_id) {
		body_ids := t.a.children_of(&node)[header_count..].clone()
		return t.lower_indexed_for_in(id, node, key_id, val_id, container_id, pool_iter, has_index,
			body_ids)
	}
	if iter_elem_type := t.iterator_for_in_elem_type(iter_type) {
		body_ids := t.a.children_of(&node)[header_count..].clone()
		return t.lower_iterator_for_in(id, node, key_id, val_id, container_id, iter_type,
			iter_elem_type, has_index, body_ids)
	}
	effective_iter := if iter_type.starts_with('...') {
		'[]' + iter_type[3..]
	} else if iter_type.starts_with('&[]') {
		iter_type[1..]
	} else {
		iter_type
	}
	if effective_iter.starts_with('[]') || effective_iter == 'string'
		|| t.is_fixed_array_type(effective_iter) {
		body_ids := t.a.children_of(&node)[header_count..].clone()
		return t.lower_indexed_for_in(id, node, key_id, val_id, container_id, effective_iter,
			has_index, body_ids)
	}
	return t.rebuild_for_in_stmt(id, node)
}

// pool_get_results_iter_type supports pool get results iter type handling for Transformer.
fn (t &Transformer) pool_get_results_iter_type(container_id flat.NodeId) ?string {
	if int(container_id) < 0 {
		return none
	}
	node := t.a.nodes[int(container_id)]
	if node.kind != .call || node.children_count == 0 {
		return none
	}
	callee_id := t.a.child(&node, 0)
	if int(callee_id) < 0 {
		return none
	}
	callee := t.a.nodes[int(callee_id)]
	if callee.kind != .index || callee.children_count < 2 {
		return none
	}
	base := t.a.child_node(&callee, 0)
	if base.kind != .selector || base.value !in ['get_results', 'get_results_ref'] {
		return none
	}
	elem_type := t.generic_call_type_arg_name(t.a.child(&callee, 1))
	if elem_type.len == 0 {
		return none
	}
	if base.value == 'get_results_ref' {
		return '[]&${elem_type}'
	}
	return '[]${elem_type}'
}

// rebuild_for_in_stmt supports rebuild for in stmt handling for Transformer.
fn (mut t Transformer) rebuild_for_in_stmt(_id flat.NodeId, node flat.Node) []flat.NodeId {
	header_count := node.value.int()
	key_id := t.a.child(&node, 0)
	val_id := t.a.child(&node, 1)
	container_id := t.a.child(&node, 2)
	new_container := t.transform_expr(container_id)
	iter_type := t.detect_for_in_type(node)
	has_index := int(val_id) >= 0
	container_is_range := if int(container_id) >= 0 {
		t.a.nodes[int(container_id)].kind == .range
	} else {
		false
	}
	if header_count == 4 || container_is_range {
		// range `for i in 0 .. n`: single loop var (child0) follows the lower bound
		if int(key_id) >= 0 {
			key_name := t.a.nodes[int(key_id)].value
			if key_name.len > 0 {
				low_id := if container_is_range {
					range_node := t.a.nodes[int(container_id)]
					t.a.child(&range_node, 0)
				} else {
					container_id
				}
				t.set_var_type(key_name, t.range_loop_var_type_name(low_id))
			}
		}
	} else if has_index {
		// two loop vars: child0 = key/index, child1 = value/element
		key_name := if int(key_id) >= 0 { t.a.nodes[int(key_id)].value } else { '' }
		val_name := if int(val_id) >= 0 { t.a.nodes[int(val_id)].value } else { '' }
		if iter_type.starts_with('map[') {
			// map[K]V: child0 (key) -> key type, child1 (val) -> value type
			bracket_end := iter_type.index(']') or { 0 }
			if key_name.len > 0 && bracket_end > 4 {
				t.set_var_type(key_name, iter_type[4..bracket_end])
			}
		} else if iter_type.starts_with('[]') || iter_type == 'string' {
			// []E: child0 (index) -> 'int'
			if key_name.len > 0 {
				t.set_var_type(key_name, 'int')
			}
		}
		if val_name.len > 0 {
			elem_type := t.infer_for_in_elem_type(iter_type, node)
			if elem_type.len > 0 {
				t.set_var_type(val_name, elem_type)
			}
		}
	} else {
		// single var, child0 is the element
		if int(key_id) >= 0 {
			key_name := t.a.nodes[int(key_id)].value
			if key_name.len > 0 {
				elem_type := t.infer_for_in_elem_type(iter_type, node)
				if elem_type.len > 0 {
					t.set_var_type(key_name, elem_type)
				}
			}
		}
	}

	mut ids := []flat.NodeId{}
	ids << key_id
	ids << val_id
	ids << new_container
	if header_count == 4 {
		range_end_id := t.a.child(&node, 3)
		ids << t.transform_expr(range_end_id)
	}
	body_ids := t.a.children_of(&node)[header_count..].clone()
	new_body := t.transform_stmts(body_ids)
	for bid in new_body {
		ids << bid
	}
	start := t.a.children.len
	for cid in ids {
		t.a.children << cid
	}
	return arr1(t.a.add_node(flat.Node{
		kind:           .for_in_stmt
		op:             node.op
		children_start: start
		children_count: flat.child_count(ids.len)
		pos:            node.pos
		value:          node.value
		typ:            node.typ
	}))
}

fn (t &Transformer) iterator_for_in_elem_type(iter_type string) ?string {
	mut clean := iter_type.trim_space()
	if clean.starts_with('&') {
		clean = clean[1..]
	}
	if clean == 'RunesIterator' || clean == 'builtin.RunesIterator' {
		return 'rune'
	}
	return none
}

// lower_range_for_in builds lower range for in data for transform.
fn (mut t Transformer) lower_range_for_in(id flat.NodeId, node flat.Node, key_id flat.NodeId, low_id flat.NodeId, high_id flat.NodeId, body_ids []flat.NodeId) []flat.NodeId {
	if int(key_id) < 0 {
		return arr1(id)
	}
	key := t.a.nodes[int(key_id)]
	if key.kind != .ident || key.value.len == 0 {
		return arr1(id)
	}
	range_type := t.range_loop_var_type_name(low_id)
	t.set_var_type(key.value, range_type)
	low := t.transform_expr(low_id)
	high := t.stable_expr_for_reuse(high_id)
	mut prefix := []flat.NodeId{}
	t.drain_pending(mut prefix)
	init := t.make_decl_assign_typed(key.value, low, range_type)
	cond := t.make_infix(.lt, t.make_ident(key.value), high)
	post := t.make_expr_stmt(t.make_postfix(t.make_ident(key.value), .inc))
	new_body := t.transform_stmts(body_ids)
	prefix << t.make_for_stmt(init, cond, post, new_body, node)
	return prefix
}

fn (t &Transformer) range_loop_var_type_name(low_id flat.NodeId) string {
	typ := t.node_type(low_id)
	if typ.len == 0 {
		return 'int'
	}
	clean := t.normalize_type_alias(typ)
	if t.is_integer_type_name(clean) {
		return typ
	}
	return 'int'
}

fn (mut t Transformer) lower_iterator_for_in(id flat.NodeId, node flat.Node, key_id flat.NodeId, val_id flat.NodeId, container_id flat.NodeId, iter_type string, elem_type string, has_index bool, body_ids []flat.NodeId) []flat.NodeId {
	if int(key_id) < 0 {
		return arr1(id)
	}
	key := t.a.nodes[int(key_id)]
	if key.kind != .ident || key.value.len == 0 {
		return arr1(id)
	}
	mut elem_name := key.value
	if has_index {
		if int(val_id) < 0 {
			return arr1(id)
		}
		val := t.a.nodes[int(val_id)]
		if val.kind != .ident || val.value.len == 0 {
			return arr1(id)
		}
		elem_name = val.value
	}
	iter_name := t.new_temp('iter')
	next_name := t.new_temp('iter_next')
	iter_expr := t.transform_expr(container_id)
	mut prefix := []flat.NodeId{}
	t.drain_pending(mut prefix)
	prefix << t.make_decl_assign_typed(iter_name, iter_expr, iter_type)
	init := if has_index {
		t.set_var_type(key.value, 'int')
		t.make_decl_assign_typed(key.value, t.make_int_literal(0), 'int')
	} else {
		t.make_empty()
	}
	t.set_var_type(elem_name, elem_type)
	next_call := t.make_call_typed('RunesIterator.next', arr1(t.make_prefix(.amp,
		t.make_ident(iter_name))), '?${elem_type}')
	next_decl := t.make_decl_assign_typed(next_name, next_call, '?${elem_type}')
	no_value := t.make_prefix(.not, t.make_selector(t.make_ident(next_name), 'ok', 'bool'))
	break_if_done := t.make_if(no_value, t.make_block(arr1(t.a.add(.break_stmt))), t.make_empty())
	elem_decl := t.make_decl_assign_typed(elem_name, t.make_selector(t.make_ident(next_name),
		'value', elem_type), elem_type)
	mut loop_body := []flat.NodeId{}
	loop_body << next_decl
	loop_body << break_if_done
	loop_body << elem_decl
	loop_body << t.transform_stmts(body_ids)
	post := if has_index {
		t.make_expr_stmt(t.make_postfix(t.make_ident(key.value), .inc))
	} else {
		t.make_empty()
	}
	prefix << t.make_for_stmt(init, t.make_bool_literal(true), post, loop_body, node)
	return prefix
}

// lower_indexed_for_in builds lower indexed for in data for transform.
fn (mut t Transformer) lower_indexed_for_in(id flat.NodeId, node flat.Node, key_id flat.NodeId, val_id flat.NodeId, container_id flat.NodeId, iter_type string, has_index bool, body_ids []flat.NodeId) []flat.NodeId {
	if int(key_id) < 0 {
		return arr1(id)
	}
	key := t.a.nodes[int(key_id)]
	if key.kind != .ident || key.value.len == 0 {
		return arr1(id)
	}
	mut container := t.stable_expr_for_reuse(container_id)
	mut prefix := []flat.NodeId{}
	t.drain_pending(mut prefix)
	mut actual_iter_type := iter_type
	container_type := t.node_type(container)
	if container_type.len > 0 && container_type !in ['array', 'map', 'unknown']
		&& for_iter_type_is_container(container_type)
		&& !for_iter_type_has_generic_placeholder(actual_iter_type)
		&& !(t.is_fixed_array_type(actual_iter_type) && !t.is_fixed_array_type(container_type)) {
		actual_iter_type = container_type
	}
	if actual_iter_type.starts_with('&[]') {
		container = t.make_prefix(.mul, container)
		t.a.nodes[int(container)].typ = actual_iter_type[1..]
		actual_iter_type = actual_iter_type[1..]
	}
	elem_type := t.infer_for_in_elem_type(actual_iter_type, node)
	if elem_type.len == 0 {
		return arr1(id)
	}
	mut idx_name := key.value
	if !has_index {
		idx_name = t.new_temp('for_idx')
	}
	mut elem_name := key.value
	if has_index {
		if int(val_id) < 0 {
			return arr1(id)
		}
		val := t.a.nodes[int(val_id)]
		if val.kind != .ident || val.value.len == 0 {
			return arr1(id)
		}
		elem_name = val.value
	}
	t.set_var_type(idx_name, 'int')
	elem_is_mut := node.op == .amp && actual_iter_type != 'string'
	elem_needs_ref := elem_is_mut && !elem_type.starts_with('&')
	elem_var_type := if elem_needs_ref { '&${elem_type}' } else { elem_type }
	t.set_var_type(elem_name, elem_var_type)
	len_expr := if t.is_fixed_array_type(actual_iter_type) {
		t.make_for_in_fixed_array_len_expr(actual_iter_type)
	} else {
		t.make_selector(container, 'len', 'int')
	}
	init := t.make_decl_assign_typed(idx_name, t.make_int_literal(0), 'int')
	cond := t.make_infix(.lt, t.make_ident(idx_name), len_expr)
	post := t.make_expr_stmt(t.make_postfix(t.make_ident(idx_name), .inc))
	elem_expr := if elem_needs_ref && actual_iter_type.starts_with('[]') {
		t.array_get_ptr(container, t.make_ident(idx_name), elem_type)
	} else if elem_needs_ref {
		t.make_prefix(.amp, t.make_index(container, t.make_ident(idx_name), elem_type))
	} else if actual_iter_type.starts_with('[]') {
		t.array_get_value(container, t.make_ident(idx_name), elem_type)
	} else {
		t.make_index(container, t.make_ident(idx_name), elem_type)
	}
	elem_decl := t.make_decl_assign_typed(elem_name, elem_expr, elem_var_type)
	mut transformed_body := []flat.NodeId{}
	if elem_needs_ref {
		had_pointer_value_lvalue := t.pointer_value_lvalues[elem_name] or { false }
		t.pointer_value_lvalues[elem_name] = true
		transformed_body = t.transform_stmts(body_ids)
		if had_pointer_value_lvalue {
			t.pointer_value_lvalues[elem_name] = true
		} else {
			t.pointer_value_lvalues.delete(elem_name)
		}
	} else {
		transformed_body = t.transform_stmts(body_ids)
	}
	mut new_body := []flat.NodeId{}
	new_body << elem_decl
	new_body << transformed_body
	prefix << t.make_for_stmt(init, cond, post, new_body, node)
	return prefix
}

// make_for_stmt builds make for stmt data for transform.
fn (mut t Transformer) make_for_stmt(init flat.NodeId, cond flat.NodeId, post flat.NodeId, body []flat.NodeId, src flat.Node) flat.NodeId {
	start := t.a.children.len
	t.a.children << init
	t.a.children << cond
	t.a.children << post
	for id in body {
		t.a.children << id
	}
	return t.a.add_node(flat.Node{
		kind:           .for_stmt
		op:             src.op
		children_start: start
		children_count: flat.child_count(3 + body.len)
		pos:            src.pos
		typ:            src.typ
	})
}

// detect_for_in_type resolves detect for in type information for transform.
fn (mut t Transformer) detect_for_in_type(node flat.Node) string {
	header_count := node.value.int()
	container_idx := if header_count >= 3 { header_count - 1 } else { 2 }
	if node.children_count > container_idx {
		iter_id := t.a.child(&node, container_idx)
		if fixed_array_type := t.detect_for_in_global_fixed_array_type(iter_id) {
			return fixed_array_type
		}
		iter_type := t.node_type(iter_id)
		if iter_type.len > 0
			&& (node.typ.len == 0 || for_iter_type_has_generic_placeholder(node.typ)
			|| for_iter_type_is_container(iter_type)) {
			return iter_type
		}
	}
	if node.typ.len > 0 {
		return node.typ
	}
	return ''
}

fn for_iter_type_is_container(iter_type string) bool {
	clean := iter_type.trim_space()
	return clean == 'string' || clean.starts_with('[]') || clean.starts_with('&[]')
		|| clean.starts_with('...') || clean.starts_with('map[') || clean.starts_with('&map[')
		|| clean.starts_with('[')
}

fn (t &Transformer) detect_for_in_global_fixed_array_type(id flat.NodeId) ?string {
	if int(id) < 0 {
		return none
	}
	node := t.a.nodes[int(id)]
	if node.kind != .ident || node.value.len == 0 {
		return none
	}
	if t.var_type(node.value).len > 0 {
		return none
	}
	mut candidates := []string{}
	if t.cur_module.len > 0 && t.cur_module != 'main' && t.cur_module != 'builtin' {
		candidates << '${t.cur_module}.${node.value}'
	}
	candidates << node.value
	for candidate in candidates {
		if typ := t.globals[candidate] {
			normalized := t.normalize_type_alias(typ)
			if t.is_fixed_array_type(normalized) {
				return normalized
			}
		}
	}
	if !isnil(t.tc) {
		checker_type := t.normalize_type_alias(t.tc.resolve_type(id).name())
		if t.is_fixed_array_type(checker_type) {
			return checker_type
		}
	}
	return none
}

fn for_iter_type_has_generic_placeholder(iter_type string) bool {
	clean := iter_type.trim_space()
	if clean.len == 0 {
		return false
	}
	if is_generic_placeholder_type_name(clean) {
		return true
	}
	if clean.starts_with('&') {
		return for_iter_type_has_generic_placeholder(clean[1..])
	}
	if clean.starts_with('[]') {
		return for_iter_type_has_generic_placeholder(clean[2..])
	}
	if clean.starts_with('map[') {
		bracket_end := clean.index(']') or { return false }
		return for_iter_type_has_generic_placeholder(clean[4..bracket_end])
			|| for_iter_type_has_generic_placeholder(clean[bracket_end + 1..])
	}
	if clean.starts_with('[') {
		bracket_end := clean.index(']') or { return false }
		return for_iter_type_has_generic_placeholder(clean[bracket_end + 1..])
	}
	return false
}

// Infer the element type for the loop variable from the iterable type.
fn (t &Transformer) infer_for_in_elem_type(iter_type string, node flat.Node) string {
	if iter_type.starts_with('&[]') {
		return iter_type[3..]
	}
	if iter_type.starts_with('[]') {
		return iter_type[2..]
	}
	if iter_type.starts_with('map[') {
		// map[K]V -> value type is everything after the closing ']'
		bracket_end := iter_type.index(']') or { return '' }
		if bracket_end + 1 < iter_type.len {
			return iter_type[bracket_end + 1..]
		}
		return ''
	}
	if iter_type == 'string' {
		return 'u8'
	}
	if t.is_fixed_array_type(iter_type) {
		return for_in_fixed_array_elem_type(iter_type)
	}
	// Check if the iterable is a range expression
	if node.children_count > 0 {
		iter_id := t.a.child(&node, 0)
		iter_node := t.a.nodes[int(iter_id)]
		if iter_node.kind == .range {
			return 'int'
		}
	}
	return ''
}

fn (mut t Transformer) make_for_in_fixed_array_len_expr(s string) flat.NodeId {
	len_text := for_in_fixed_array_len_text(s)
	if is_decimal_text(len_text) {
		return t.make_int_literal(len_text.int())
	}
	if !isnil(t.tc) {
		if v := t.tc.const_int_value(len_text, []string{}) {
			return t.make_int_literal(v)
		}
	}
	if len_text.contains('.') {
		base := len_text.all_before_last('.')
		field := len_text.all_after_last('.')
		return t.make_selector(t.make_ident(base), field, 'int')
	}
	return t.make_ident(len_text)
}

fn for_in_fixed_array_elem_type(s string) string {
	clean := s.trim_space()
	if clean.starts_with('[') {
		return clean.all_after(']')
	}
	open := clean.last_index_u8(`[`)
	if open < 0 {
		return ''
	}
	return clean[..open]
}

fn for_in_fixed_array_len_text(s string) string {
	clean := s.trim_space()
	if clean.starts_with('[') {
		return clean.all_after('[').all_before(']').trim_space()
	}
	open := clean.last_index_u8(`[`)
	if open < 0 {
		return ''
	}
	close_rel := clean[open + 1..].index_u8(`]`)
	if close_rel < 0 {
		return ''
	}
	return clean[open + 1..open + 1 + close_rel].trim_space()
}