// Copyright (c) 2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

module cleanc

import strings
import v2.ast

fn (mut g Gen) gen_for_stmt(node ast.ForStmt) {
	if node.init is ast.ForInStmt {
		if g.gen_map_for_in_stmt(node, node.init) {
			return
		}
		if g.gen_array_for_in_stmt(node, node.init) {
			return
		}
	}
	if g.gen_transformed_untyped_map_for_in_stmt(node) {
		return
	}
	g.write_indent()
	has_init := !is_empty_stmt(node.init)
	has_cond := !is_empty_expr(node.cond)
	has_post := !is_empty_stmt(node.post)

	if has_init || has_post {
		// C-style for loop: for (init; cond; post)
		g.sb.write_string('for (')
		if has_init {
			g.gen_stmt_inline(node.init)
		}
		g.sb.write_string('; ')
		if has_cond {
			g.expr(node.cond)
		}
		g.sb.write_string('; ')
		if has_post {
			g.gen_stmt_inline(node.post)
		}
		g.sb.writeln(') {')
	} else if has_cond {
		// while-style: for cond {
		g.sb.write_string('while (')
		g.expr(node.cond)
		g.sb.writeln(') {')
	} else {
		// Infinite loop: for {
		g.sb.writeln('for (;;) {')
	}

	g.indent++
	// Save runtime_local_types: variables declared in the loop body
	// (e.g. _filter_it in nested map/filter) can shadow outer variables.
	// Without save/restore, the inner type overwrites the outer in the flat map.
	saved_local_types := g.runtime_local_types.clone()
	saved_decl_types := g.runtime_decl_types.clone()
	g.gen_stmts(node.stmts)
	g.runtime_local_types = saved_local_types.clone()
	g.runtime_decl_types = saved_decl_types.clone()
	g.not_local_var_cache.clear()
	g.indent--
	g.write_indent()
	g.sb.writeln('}')
}

fn (mut g Gen) gen_map_for_in_stmt(node ast.ForStmt, for_in ast.ForInStmt) bool {
	mut map_type := g.get_expr_type(for_in.expr).trim_space().trim_right('*')
	if (map_type == '' || map_type == 'int') && for_in.expr is ast.Ident {
		map_type = (g.get_local_var_c_type(for_in.expr.name) or { '' }).trim_space().trim_right('*')
	}
	if !map_type.starts_with('Map_') {
		return false
	}
	key_type, value_type := g.parse_map_kv_types(map_type['Map_'.len..])
	if key_type == '' || value_type == '' {
		return false
	}
	id := g.tmp_counter
	g.tmp_counter++
	map_tmp := '_map_iter_${id}'
	len_tmp := '_map_len_${id}'
	idx_tmp := '_map_idx_${id}'
	delta_tmp := '_map_delta_${id}'
	mut expr_sb := strings.new_builder(64)
	saved_sb := g.sb
	g.sb = expr_sb
	g.expr(for_in.expr)
	map_expr := g.sb.str()
	g.sb = saved_sb
	g.write_indent()
	g.sb.writeln('${map_type} ${map_tmp} = ${map_expr};')
	g.write_indent()
	g.sb.writeln('int ${len_tmp} = ${map_tmp}.key_values.len;')
	g.write_indent()
	g.sb.writeln('for (int ${idx_tmp} = 0; (${idx_tmp} < ${len_tmp}); ${idx_tmp} = (${idx_tmp} + 1)) {')
	g.indent++
	g.write_indent()
	g.sb.writeln('int ${delta_tmp} = (${map_tmp}.key_values.len - ${len_tmp});')
	g.write_indent()
	g.sb.writeln('${len_tmp} = ${map_tmp}.key_values.len;')
	g.write_indent()
	g.sb.writeln('if ((${delta_tmp} < 0)) {')
	g.indent++
	g.write_indent()
	g.sb.writeln('${idx_tmp} = -1;')
	g.write_indent()
	g.sb.writeln('continue;')
	g.indent--
	g.write_indent()
	g.sb.writeln('}')
	g.write_indent()
	g.sb.writeln('if (!DenseArray__has_index(&${map_tmp}.key_values, ${idx_tmp})) {')
	g.indent++
	g.write_indent()
	g.sb.writeln('continue;')
	g.indent--
	g.write_indent()
	g.sb.writeln('}')
	key_name := if for_in.key is ast.Ident { for_in.key.name } else { '' }
	value_name := if for_in.value is ast.Ident { for_in.value.name } else { '' }
	saved_local_types := g.runtime_local_types.clone()
	saved_decl_types := g.runtime_decl_types.clone()
	if key_name != '' && key_name != '_' {
		g.write_indent()
		g.sb.writeln('${key_type} ${key_name} = *((${key_type}*)DenseArray__key(&${map_tmp}.key_values, ${idx_tmp}));')
		if key_type == 'string' {
			g.write_indent()
			g.sb.writeln('${key_name} = string__clone(${key_name});')
		}
		g.remember_runtime_local_type(key_name, key_type)
	}
	if value_name != '' && value_name != '_' {
		g.write_indent()
		g.sb.writeln('${value_type} ${value_name} = *((${value_type}*)DenseArray__value(&${map_tmp}.key_values, ${idx_tmp}));')
		if value_type == 'string' {
			g.write_indent()
			g.sb.writeln('${value_name} = string__clone(${value_name});')
		}
		g.remember_runtime_local_type(value_name, value_type)
	}
	g.gen_stmts(node.stmts)
	g.runtime_local_types = saved_local_types.clone()
	g.runtime_decl_types = saved_decl_types.clone()
	g.not_local_var_cache.clear()
	g.indent--
	g.write_indent()
	g.sb.writeln('}')
	return true
}

fn cleanc_for_in_ident_name(expr ast.Expr) (string, bool) {
	if expr is ast.Ident {
		return expr.name, false
	}
	if expr is ast.ModifierExpr {
		if expr.expr is ast.Ident {
			return expr.expr.name, expr.kind == .key_mut
		}
	}
	return '', false
}

fn (mut g Gen) gen_array_for_in_stmt(node ast.ForStmt, for_in ast.ForInStmt) bool {
	mut array_type := g.get_expr_type(for_in.expr).trim_space().trim_right('*')
	if (array_type == '' || array_type == 'int') && for_in.expr is ast.Ident {
		array_type =
			(g.get_local_var_c_type(for_in.expr.name) or { '' }).trim_space().trim_right('*')
	}
	if array_type == '' {
		return false
	}
	mut array_decl_type := array_type
	if !c_type_is_array_value(array_decl_type) {
		alias_base := g.array_alias_base_type(array_decl_type)
		if alias_base != '' {
			array_decl_type = alias_base
		}
	}
	if !c_type_is_array_value(array_decl_type) {
		return false
	}
	elem_type := g.array_alias_elem_type_from_c_type(array_decl_type)
	if elem_type == '' {
		return false
	}
	id := g.tmp_counter
	g.tmp_counter++
	array_tmp := '_arr_iter_${id}'
	idx_tmp := '_arr_idx_${id}'
	mut expr_sb := strings.new_builder(64)
	saved_sb := g.sb
	g.sb = expr_sb
	g.expr(for_in.expr)
	array_expr := g.sb.str()
	g.sb = saved_sb
	g.write_indent()
	g.sb.writeln('${array_decl_type} ${array_tmp} = ${array_expr};')
	g.write_indent()
	g.sb.writeln('for (int ${idx_tmp} = 0; (${idx_tmp} < ${array_tmp}.len); ${idx_tmp} = (${idx_tmp} + 1)) {')
	g.indent++
	key_name, _ := cleanc_for_in_ident_name(for_in.key)
	value_name, value_is_mut := cleanc_for_in_ident_name(for_in.value)
	saved_local_types := g.runtime_local_types.clone()
	saved_decl_types := g.runtime_decl_types.clone()
	if key_name != '' && key_name != '_' {
		g.write_indent()
		g.sb.writeln('int ${key_name} = ${idx_tmp};')
		g.remember_runtime_local_type(key_name, 'int')
	}
	if value_name != '' && value_name != '_' {
		g.write_indent()
		if value_is_mut {
			g.sb.writeln('${elem_type}* ${value_name} = &(((${elem_type}*)${array_tmp}.data)[${idx_tmp}]);')
			g.remember_runtime_local_type(value_name, '${elem_type}*')
		} else {
			g.sb.writeln('${elem_type} ${value_name} = ((${elem_type}*)${array_tmp}.data)[${idx_tmp}];')
			if elem_type == 'string' {
				g.write_indent()
				g.sb.writeln('${value_name} = string__clone(${value_name});')
			}
			g.remember_runtime_local_type(value_name, elem_type)
		}
	}
	g.gen_stmts(node.stmts)
	g.runtime_local_types = saved_local_types.clone()
	g.runtime_decl_types = saved_decl_types.clone()
	g.not_local_var_cache.clear()
	g.indent--
	g.write_indent()
	g.sb.writeln('}')
	return true
}

fn (mut g Gen) gen_transformed_untyped_map_for_in_stmt(node ast.ForStmt) bool {
	if node.init !is ast.AssignStmt || node.stmts.len == 0 {
		return false
	}
	init := node.init as ast.AssignStmt
	if init.lhs.len == 0 || init.lhs[0] !is ast.Ident {
		return false
	}
	key_name := (init.lhs[0] as ast.Ident).name
	first_stmt := node.stmts[0]
	if first_stmt !is ast.AssignStmt {
		return false
	}
	value_assign := first_stmt as ast.AssignStmt
	if value_assign.lhs.len == 0 || value_assign.rhs.len == 0 || value_assign.lhs[0] !is ast.Ident
		|| value_assign.rhs[0] !is ast.IndexExpr {
		return false
	}
	value_name := (value_assign.lhs[0] as ast.Ident).name
	index_expr := value_assign.rhs[0] as ast.IndexExpr
	mut map_type := g.get_expr_type(index_expr.lhs).trim_space().trim_right('*')
	if (map_type == '' || map_type == 'int') && index_expr.lhs is ast.Ident {
		map_type =
			(g.get_local_var_c_type(index_expr.lhs.name) or { '' }).trim_space().trim_right('*')
	}
	if !map_type.starts_with('Map_') {
		return false
	}
	map_node := ast.ForStmt{
		stmts: node.stmts[1..]
	}
	return g.gen_map_for_in_stmt(map_node, ast.ForInStmt{
		key:   ast.Expr(ast.Ident{
			name: key_name
		})
		value: ast.Expr(ast.Ident{
			name: value_name
		})
		expr:  index_expr.lhs
	})
}

fn (mut g Gen) gen_stmt_inline(node ast.Stmt) {
	if !stmt_has_valid_data(node) {
		return
	}
	match node {
		ast.AssignStmt {
			lhs := node.lhs[0]
			rhs := node.rhs[0]
			if node.op == .decl_assign {
				mut name := ''
				if lhs is ast.Ident {
					name = lhs.name
				}
				typ := g.get_expr_type(rhs)
				g.sb.write_string('${typ} ${name} = ')
				g.expr(rhs)
				// Register the for-loop init variable so assignments inside
				// the loop body don't re-declare it.
				if name != '' && typ != '' {
					g.remember_runtime_local_type(name, typ)
				}
			} else {
				g.expr(lhs)
				op_str := match node.op {
					.assign { '=' }
					.plus_assign { '+=' }
					.minus_assign { '-=' }
					.mul_assign { '*=' }
					.div_assign { '/=' }
					.mod_assign { '%=' }
					.and_assign { '&=' }
					.or_assign { '|=' }
					.xor_assign { '^=' }
					.left_shift_assign { '<<=' }
					.right_shift_assign { '>>=' }
					else { '=' }
				}

				g.sb.write_string(' ${op_str} ')
				g.expr(rhs)
			}
		}
		ast.ExprStmt {
			g.expr(node.expr)
		}
		else {}
	}
}