import os
import v3.flat
import v3.parser
import v3.pref
import v3.transform
import v3.types

// parse_transform_source reads parse transform source input for v3 tests.
fn parse_transform_source(source string) &flat.FlatAst {
	src := os.join_path(os.temp_dir(), 'v3_transformer_parity_test.v')
	return parse_transform_file(src, source)
}

// parse_transform_file reads parse transform file input for v3 tests.
fn parse_transform_file(src string, source string) &flat.FlatAst {
	os.write_file(src, source) or { panic(err) }
	prefs := pref.new_preferences()
	mut p := parser.Parser.new(prefs)
	mut a := p.parse_file(src)
	mut tc := types.TypeChecker.new(a)
	tc.collect(a)
	tc.annotate_types()
	transform.transform(mut a, &tc)
	return a
}

// find_fn resolves find fn information for v3 tests.
fn find_fn(a &flat.FlatAst, name string) flat.Node {
	for node in a.nodes {
		if node.kind == .fn_decl && node.value == name {
			return node
		}
	}
	assert false
	return flat.Node{}
}

// first_decl_rhs returns first decl rhs data for v3 tests.
fn first_decl_rhs(a &flat.FlatAst, fn_name string) flat.Node {
	f := find_fn(a, fn_name)
	for i in 0 .. f.children_count {
		stmt := a.child_node(&f, i)
		if stmt.kind == .decl_assign && stmt.children_count == 2 {
			return *a.child_node(stmt, 1)
		}
	}
	assert false
	return flat.Node{}
}

// decl_rhs supports decl rhs handling for v3 tests.
fn decl_rhs(a &flat.FlatAst, fn_name string, name string) flat.Node {
	f := find_fn(a, fn_name)
	for i in 0 .. f.children_count {
		stmt := a.child_node(&f, i)
		if stmt.kind == .decl_assign && stmt.children_count == 2 {
			lhs := a.child_node(stmt, 0)
			if lhs.kind == .ident && lhs.value == name {
				return *a.child_node(stmt, 1)
			}
		}
	}
	assert false
	return flat.Node{}
}

// first_return_expr returns first return expr data for v3 tests.
fn first_return_expr(a &flat.FlatAst, fn_name string) flat.Node {
	f := find_fn(a, fn_name)
	for i in 0 .. f.children_count {
		stmt := a.child_node(&f, i)
		if stmt.kind == .return_stmt && stmt.children_count > 0 {
			return *a.child_node(stmt, 0)
		}
	}
	assert false
	return flat.Node{}
}

// count_kind supports count kind handling for v3 tests.
fn count_kind(a &flat.FlatAst, id flat.NodeId, kind flat.NodeKind) int {
	if int(id) < 0 {
		return 0
	}
	node := a.nodes[int(id)]
	mut total := if node.kind == kind { 1 } else { 0 }
	for i in 0 .. node.children_count {
		total += count_kind(a, a.child(&node, i), kind)
	}
	return total
}

// count_call_name supports count call name handling for v3 tests.
fn count_call_name(a &flat.FlatAst, id flat.NodeId, name string) int {
	if int(id) < 0 {
		return 0
	}
	node := a.nodes[int(id)]
	mut total := 0
	if node.kind == .call && node.children_count > 0 {
		fn_node := a.child_node(&node, 0)
		if fn_node.kind == .ident && fn_node.value == name {
			total++
		}
	}
	for i in 0 .. node.children_count {
		total += count_call_name(a, a.child(&node, i), name)
	}
	return total
}

// count_infix_op supports count infix op handling for v3 tests.
fn count_infix_op(a &flat.FlatAst, id flat.NodeId, op flat.Op) int {
	if int(id) < 0 {
		return 0
	}
	node := a.nodes[int(id)]
	mut total := if node.kind == .infix && node.op == op { 1 } else { 0 }
	for i in 0 .. node.children_count {
		total += count_infix_op(a, a.child(&node, i), op)
	}
	return total
}

// count_wide_decl_assigns supports count wide decl assigns handling for v3 tests.
fn count_wide_decl_assigns(a &flat.FlatAst, id flat.NodeId) int {
	if int(id) < 0 {
		return 0
	}
	node := a.nodes[int(id)]
	mut total := if node.kind == .decl_assign && node.children_count > 2 { 1 } else { 0 }
	for i in 0 .. node.children_count {
		total += count_wide_decl_assigns(a, a.child(&node, i))
	}
	return total
}

// count_wide_assigns supports count wide assigns handling for v3 tests.
fn count_wide_assigns(a &flat.FlatAst, id flat.NodeId) int {
	if int(id) < 0 {
		return 0
	}
	node := a.nodes[int(id)]
	mut total := if node.kind == .assign && node.children_count > 2 { 1 } else { 0 }
	for i in 0 .. node.children_count {
		total += count_wide_assigns(a, a.child(&node, i))
	}
	return total
}

// count_selector_value supports count selector value handling for v3 tests.
fn count_selector_value(a &flat.FlatAst, id flat.NodeId, value string) int {
	if int(id) < 0 {
		return 0
	}
	node := a.nodes[int(id)]
	mut total := if node.kind == .selector && node.value == value { 1 } else { 0 }
	for i in 0 .. node.children_count {
		total += count_selector_value(a, a.child(&node, i), value)
	}
	return total
}

// test_typeof_expression_lowers_to_string_literal validates this v3 regression case.
fn test_typeof_expression_lowers_to_string_literal() {
	a := parse_transform_source('
fn main() {
	name := typeof(123)
}
')
	main_fn := find_fn(a, 'main')
	for i in 0 .. main_fn.children_count {
		stmt := a.child_node(&main_fn, i)
		if stmt.kind == .decl_assign && stmt.children_count == 2 {
			rhs := a.child_node(stmt, 1)
			assert rhs.kind == .string_literal
			assert rhs.value == 'int literal'
			return
		}
	}
	assert false
}

// test_vmodroot_lowers_to_nearest_vmod_dir validates this v3 regression case.
fn test_vmodroot_lowers_to_nearest_vmod_dir() {
	root := os.join_path(os.temp_dir(), 'v3_vmodroot_transformer_parity')
	os.mkdir_all(root) or { panic(err) }
	os.write_file(os.join_path(root, 'v.mod'), 'Module { name: "parity" }') or { panic(err) }
	src := os.join_path(root, 'main.v')
	a := parse_transform_file(src, '
fn main() {
	root := @VMODROOT
}
')
	rhs := first_decl_rhs(a, 'main')
	assert rhs.kind == .string_literal
	assert rhs.value == root
}

// test_return_match_lowers_to_explicit_branch_returns validates this v3 regression case.
fn test_return_match_lowers_to_explicit_branch_returns() {
	a := parse_transform_source('
fn choose(x int) int {
	return match x {
		0 {
			10
		}
		else {
			20
		}
	}
}
')
	choose_fn := find_fn(a, 'choose')
	mut match_count := 0
	mut return_count := 0
	for i in 0 .. choose_fn.children_count {
		child_id := a.child(&choose_fn, i)
		match_count += count_kind(a, child_id, .match_stmt)
		return_count += count_kind(a, child_id, .return_stmt)
	}
	assert match_count == 0
	assert return_count == 2
}

// test_return_if_keeps_or_lowering_inside_branch validates this v3 regression case.
fn test_return_if_keeps_or_lowering_inside_branch() {
	a := parse_transform_source('
fn maybe_int() ?int {
	return 8
}

fn choose(flag bool) int {
	return if flag {
		maybe_int() or {
			1
		}
	} else {
		2
	}
}
')
	choose_fn := find_fn(a, 'choose')
	mut body_ids := []flat.NodeId{}
	for i in 0 .. choose_fn.children_count {
		child_id := a.child(&choose_fn, i)
		child := a.nodes[int(child_id)]
		if child.kind != .param {
			body_ids << child_id
		}
	}
	assert body_ids.len == 1
	assert a.nodes[int(body_ids[0])].kind == .if_expr
}

// test_if_expr_value_lowers_to_temp_and_branch_assigns validates this v3 regression case.
fn test_if_expr_value_lowers_to_temp_and_branch_assigns() {
	a := parse_transform_source('
fn main() {
	x := if true {
		1
	} else {
		2
	}
}
')
	rhs := decl_rhs(a, 'main', 'x')
	assert rhs.kind == .ident
	assert rhs.value.starts_with('__if_val_')
	main_fn := find_fn(a, 'main')
	mut if_count := 0
	mut assign_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		if_count += count_kind(a, child_id, .if_expr)
		assign_count += count_kind(a, child_id, .assign)
	}
	assert if_count == 1
	assert assign_count == 2
}

// test_if_expr_call_arg_lowers_before_call validates this v3 regression case.
fn test_if_expr_call_arg_lowers_before_call() {
	a := parse_transform_source('
fn use(x int) int {
	return x
}

fn main() {
	x := use(if true { 1 } else { 2 })
}
')
	rhs := decl_rhs(a, 'main', 'x')
	assert rhs.kind == .call
	assert rhs.children_count == 2
	arg := a.child_node(&rhs, 1)
	assert arg.kind == .ident
	assert arg.value.starts_with('__if_val_')
}

// test_nested_if_expr_branch_lowers_to_outer_temp_assignment validates this v3 regression case.
fn test_nested_if_expr_branch_lowers_to_outer_temp_assignment() {
	a := parse_transform_source('
fn main() {
	x := if true {
		if false {
			1
		} else {
			2
		}
	} else {
		3
	}
}
')
	rhs := decl_rhs(a, 'main', 'x')
	assert rhs.kind == .ident
	assert rhs.value.starts_with('__if_val_')
	main_fn := find_fn(a, 'main')
	mut if_count := 0
	mut assign_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		if_count += count_kind(a, child_id, .if_expr)
		assign_count += count_kind(a, child_id, .assign)
	}
	assert if_count == 2
	assert assign_count == 4
}

// test_multi_return_decl_lowers_to_temp_field_decls validates this v3 regression case.
fn test_multi_return_decl_lowers_to_temp_field_decls() {
	a := parse_transform_source("
fn pair() (int, string) {
	return 1, 'ok'
}

fn main() {
	a, b := pair()
}
")
	main_fn := find_fn(a, 'main')
	mut wide_decl_count := 0
	for i in 0 .. main_fn.children_count {
		wide_decl_count += count_wide_decl_assigns(a, a.child(&main_fn, i))
	}
	assert wide_decl_count == 0
	a_rhs := decl_rhs(a, 'main', 'a')
	b_rhs := decl_rhs(a, 'main', 'b')
	assert a_rhs.kind == .selector
	assert a_rhs.value == 'arg0'
	assert b_rhs.kind == .selector
	assert b_rhs.value == 'arg1'
}

// test_multi_return_assign_lowers_to_temp_field_assigns validates this v3 regression case.
fn test_multi_return_assign_lowers_to_temp_field_assigns() {
	a := parse_transform_source("
fn pair() (int, string) {
	return 1, 'ok'
}

fn main() {
	mut a := 0
	mut b := ''
	a, b = pair()
}
")
	main_fn := find_fn(a, 'main')
	mut wide_assign_count := 0
	mut arg0_count := 0
	mut arg1_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		wide_assign_count += count_wide_assigns(a, child_id)
		arg0_count += count_selector_value(a, child_id, 'arg0')
		arg1_count += count_selector_value(a, child_id, 'arg1')
	}
	assert wide_assign_count == 0
	assert arg0_count == 1
	assert arg1_count == 1
}

// test_assoc_expr_lowers_to_temp_and_field_assigns validates this v3 regression case.
fn test_assoc_expr_lowers_to_temp_and_field_assigns() {
	a := parse_transform_source('
struct Point {
	x int
	y int
}

fn main() {
	p := Point{x: 1, y: 2}
	q := Point{...p, y: 3}
}
')
	rhs := decl_rhs(a, 'main', 'q')
	assert rhs.kind == .ident
	assert rhs.value.starts_with('__assoc_')
	main_fn := find_fn(a, 'main')
	mut assoc_count := 0
	mut assign_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		assoc_count += count_kind(a, child_id, .assoc)
		assign_count += count_kind(a, child_id, .assign)
	}
	assert assoc_count == 0
	assert assign_count == 1
}

// test_return_assoc_expr_lowers_before_return validates this v3 regression case.
fn test_return_assoc_expr_lowers_before_return() {
	a := parse_transform_source('
struct Point {
	x int
	y int
}

fn moved(p Point) Point {
	return Point{...p, x: 4}
}
')
	ret := first_return_expr(a, 'moved')
	assert ret.kind == .ident
	assert ret.value.starts_with('__assoc_')
	moved_fn := find_fn(a, 'moved')
	mut assoc_count := 0
	for i in 0 .. moved_fn.children_count {
		assoc_count += count_kind(a, a.child(&moved_fn, i), .assoc)
	}
	assert assoc_count == 0
}

// test_array_append_stmt_lowers_to_runtime_push validates this v3 regression case.
fn test_array_append_stmt_lowers_to_runtime_push() {
	a := parse_transform_source('
fn main() {
	mut xs := []int{}
	xs << 3
}
')
	main_fn := find_fn(a, 'main')
	mut push_count := 0
	mut left_shift_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		push_count += count_call_name(a, child_id, 'array_push')
		left_shift_count += count_infix_op(a, child_id, .left_shift)
	}
	assert push_count == 1
	assert left_shift_count == 0
}

// test_array_append_many_stmt_lowers_to_runtime_push_many validates this v3 regression case.
fn test_array_append_many_stmt_lowers_to_runtime_push_many() {
	a := parse_transform_source('
fn main() {
	mut xs := []int{}
	ys := []int{}
	xs << ys
}
')
	main_fn := find_fn(a, 'main')
	mut push_many_count := 0
	mut left_shift_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		push_many_count += count_call_name(a, child_id, 'array__push_many')
		left_shift_count += count_infix_op(a, child_id, .left_shift)
	}
	assert push_many_count == 1
	assert left_shift_count == 0
}

fn test_array_push_many_method_lowers_to_runtime_ptr_call() {
	a := parse_transform_source('
fn main() {
	mut xs := []u8{}
	mut b := u8(7)
	unsafe {
		xs.push_many(&b, 1)
	}
}
')
	main_fn := find_fn(a, 'main')
	mut push_many_ptr_count := 0
	mut push_many_selector_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		push_many_ptr_count += count_call_name(a, child_id, 'array_push_many_ptr')
		push_many_selector_count += count_selector_value(a, child_id, 'push_many')
	}
	assert push_many_ptr_count == 1
	assert push_many_selector_count == 0
}

fn test_array_push_many_type_marker_lowers_to_runtime_ptr_call() {
	a := parse_transform_source('
struct PushManyItem {
	value int
}

fn main() {
	mut xs := []PushManyItem{}
	item := PushManyItem{value: 7}
	xs.push_many(item, PushManyItem)
}
')
	main_fn := find_fn(a, 'main')
	mut push_many_ptr_count := 0
	mut push_many_selector_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		push_many_ptr_count += count_call_name(a, child_id, 'array_push_many_ptr')
		push_many_selector_count += count_selector_value(a, child_id, 'push_many')
	}
	assert push_many_ptr_count == 1
	assert push_many_selector_count == 0
}

fn test_sql_expr_lowers_to_success_result_default() {
	a := parse_transform_source('
struct User {
	id int
}

fn main() {
	db := 0
	users := sql db {
		select from User
	}
	count := sql db {
		select count from User
	}
	created := sql db {
		create table User
	}
}
')
	main_fn := find_fn(a, 'main')
	mut sql_expr_count := 0
	for i in 0 .. main_fn.children_count {
		sql_expr_count += count_kind(a, a.child(&main_fn, i), .sql_expr)
	}
	assert sql_expr_count == 0
	users := decl_rhs(a, 'main', 'users')
	assert users.kind == .struct_init
	assert users.value == '![]User'
	count := decl_rhs(a, 'main', 'count')
	assert count.kind == .struct_init
	assert count.value == '!int'
	created := decl_rhs(a, 'main', 'created')
	assert created.kind == .struct_init
	assert created.value == '!void'
}

// test_or_expr_lowers_to_temp_and_if validates this v3 regression case.
fn test_or_expr_lowers_to_temp_and_if() {
	a := parse_transform_source('
fn maybe_int() ?int {
	return 3
}

fn main() {
	x := maybe_int() or {
		7
	}
}
')
	main_fn := find_fn(a, 'main')
	mut or_count := 0
	mut if_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		or_count += count_kind(a, child_id, .or_expr)
		if_count += count_kind(a, child_id, .if_expr)
	}
	assert or_count == 0
	assert if_count == 1
}

// test_map_index_or_lowers_to_get_check validates this v3 regression case.
fn test_map_index_or_lowers_to_get_check() {
	a := parse_transform_source('
fn main() {
	mut m := map[string]int{}
	x := m["a"] or {
		7
	}
}
')
	main_fn := find_fn(a, 'main')
	mut or_count := 0
	mut get_check_count := 0
	for i in 0 .. main_fn.children_count {
		child_id := a.child(&main_fn, i)
		or_count += count_kind(a, child_id, .or_expr)
		get_check_count += count_call_name(a, child_id, 'map__get_check')
	}
	assert or_count == 0
	assert get_check_count == 1
}