# Copyright (c) 2025 JP Hutchins
# SPDX-License-Identifier: MIT

"""Test functional programming features with Func class."""

from typing import Any, NamedTuple, assert_type

from mahonia import Const, Contains, Expr, FilterExpr, Func, MapExpr, Not, Var, extract_vars


class FuncCtx(NamedTuple):
	x: int
	y: int
	z: float


def test_func_manual_construction() -> None:
	"""Test manual construction of Func objects."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	expr = x + y
	func = Func((x, y), expr)

	# Test string representation without context
	assert func.to_string() == "(x, y) -> (x + y)"

	# Test string representation with context
	ctx = FuncCtx(x=1, y=2, z=3.14)
	assert func.to_string(ctx) == "(1, 2) -> (x + y) -> 3"


def test_to_func_automatic_extraction() -> None:
	"""Test automatic variable extraction with to_func()."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	expr = x * y
	func = expr.to_func()

	# Should automatically extract (x, y) variables
	assert func.to_string() == "(x, y) -> (x * y)"

	# Test with context
	ctx = FuncCtx(x=3, y=4, z=1.0)
	assert func.to_string(ctx) == "(3, 4) -> (x * y) -> 12"


def test_single_variable_extraction() -> None:
	"""Test extraction with single variable."""
	x = Var[int, FuncCtx]("x")

	expr = x * 2  # x * Const(2)
	func = expr.to_func()

	# Should extract just x
	assert func.to_string() == "x -> (x * 2)"

	ctx = FuncCtx(x=5, y=0, z=0.0)
	assert func.to_string(ctx) == "5 -> (x * 2) -> 10"


def test_complex_expression_extraction() -> None:
	"""Test variable extraction from complex expressions."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")
	z = Var[float, FuncCtx]("z")

	# Complex expression: (x + y) * z
	expr = (x + y) * z
	func = expr.to_func()

	# Should extract all three variables in order
	assert func.to_string() == "(x, y, z) -> ((x + y) * z)"

	ctx = FuncCtx(x=2, y=3, z=1.5)
	expected = "(2, 3, 1.5) -> ((x + y) * z) -> 7.5"
	assert func.to_string(ctx) == expected


def test_constant_only_expression() -> None:
	"""Test expression with only constants."""
	expr = Const("answer", 42)
	func = expr.to_func()

	# No variables to extract
	assert func.to_string() == "() -> answer:42"

	ctx = FuncCtx(x=0, y=0, z=0.0)
	assert func.to_string(ctx) == "() -> answer:42 -> 42"


def test_mixed_variables_and_constants() -> None:
	"""Test extraction with mix of variables and constants."""
	x = Var[int, FuncCtx]("x")

	expr = x + Const("offset", 10)
	func = expr.to_func()

	# Should only extract variable x
	assert func.to_string() == "x -> (x + offset:10)"

	ctx = FuncCtx(x=5, y=0, z=0.0)
	assert func.to_string(ctx) == "5 -> (x + offset:10) -> 15"


def test_variable_order_preservation() -> None:
	"""Test that variable extraction preserves order of first occurrence."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")
	z = Var[float, FuncCtx]("z")

	# Use variables in different order: z, x, y
	expr = z + x + y
	func = expr.to_func()

	# Should preserve order of first occurrence: z, x, y
	assert func.to_string() == "(z, x, y) -> ((z + x) + y)"


def test_duplicate_variables() -> None:
	"""Test that duplicate variable usage doesn't create duplicates in args."""
	x = Var[int, FuncCtx]("x")

	# Use x multiple times
	expr = (x * x) + x
	func = expr.to_func()

	# Should only have x once in args
	assert func.to_string() == "x -> ((x * x) + x)"

	ctx = FuncCtx(x=3, y=0, z=0.0)
	assert func.to_string(ctx) == "3 -> ((x * x) + x) -> 12"  # 3*3 + 3 = 12


def test_nested_expressions() -> None:
	"""Test variable extraction from nested expressions."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	# Nested: (x + (y * 2))
	expr = x + (y * 2)
	func = expr.to_func()

	assert func.to_string() == "(x, y) -> (x + (y * 2))"

	ctx = FuncCtx(x=1, y=3, z=0.0)
	assert func.to_string(ctx) == "(1, 3) -> (x + (y * 2)) -> 7"  # 1 + (3*2) = 7


def test_extract_vars_helper_function() -> None:
	"""Test the extract_vars helper function directly."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	expr = x + y
	vars_tuple = extract_vars((), expr)

	assert len(vars_tuple) == 2
	assert vars_tuple[0].name == "x"
	assert vars_tuple[1].name == "y"


def test_extract_vars_with_existing_vars() -> None:
	"""Test extract_vars when some variables already exist."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")
	z = Var[float, FuncCtx]("z")

	# Start with x already in vars
	existing_vars = (x,)

	expr = x + y + z  # x is already in existing_vars
	vars_tuple = extract_vars(existing_vars, expr)

	# Should have x, y, z (x was already there, y and z added)
	assert len(vars_tuple) == 3
	assert vars_tuple[0] is x  # Same object
	assert vars_tuple[1].name == "y"
	assert vars_tuple[2].name == "z"


def test_func_types() -> None:
	"""Test that Func maintains proper types."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	expr = x + y
	func = Func((x, y), expr)

	# Type checking
	assert_type(func, Func[int, FuncCtx])
	assert len(func.args) == 2
	assert func.expr == expr


def test_func_result_type_int() -> None:
	"""Verify Func preserves int result type."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	add_func = (x + y).to_func()
	assert_type(add_func, Func[int, FuncCtx])
	assert_type(add_func.expr, Expr[Any, FuncCtx, int])

	mul_func = (x * y).to_func()
	assert_type(mul_func, Func[int, FuncCtx])
	assert_type(mul_func.expr, Expr[Any, FuncCtx, int])

	sub_func = (x - y).to_func()
	assert_type(sub_func, Func[int, FuncCtx])
	assert_type(sub_func.expr, Expr[Any, FuncCtx, int])


def test_func_result_type_float() -> None:
	"""Verify Func preserves float result type."""
	z = Var[float, FuncCtx]("z")

	float_func = (z * 2.0).to_func()
	assert_type(float_func, Func[float, FuncCtx])
	assert_type(float_func.expr, Expr[Any, FuncCtx, float])


def test_func_result_type_bool() -> None:
	"""Verify Func preserves bool result type for comparisons."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	gt_func = (x > y).to_func()
	assert_type(gt_func, Func[bool, FuncCtx])
	assert_type(gt_func.expr, Expr[Any, FuncCtx, bool])

	eq_func = (x == y).to_func()
	assert_type(eq_func, Func[bool, FuncCtx])
	assert_type(eq_func.expr, Expr[Any, FuncCtx, bool])

	and_func = ((x > 0) & (y > 0)).to_func()
	assert_type(and_func, Func[bool, FuncCtx])
	assert_type(and_func.expr, Expr[Any, FuncCtx, bool])


def test_func_args_tuple_types() -> None:
	"""Verify Func.args tuple contains correctly typed Vars."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	func = Func((x, y), x + y)
	assert_type(func.args, tuple[Expr[Any, FuncCtx, Any], ...])

	single_func = Func((x,), x * 2)
	assert_type(single_func.args, tuple[Expr[Any, FuncCtx, Any], ...])


def test_func_from_const_any_context() -> None:
	"""Verify Func from Const has Any context type."""
	const_func = Const("answer", 42).to_func()
	assert_type(const_func, Func[int, Any])
	assert_type(const_func.expr, Expr[Any, Any, int])


def test_func_with_boolean_expressions() -> None:
	"""Test Func with boolean expressions."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	expr = x > y
	func = expr.to_func()

	assert func.to_string() == "(x, y) -> (x > y)"

	ctx = FuncCtx(x=5, y=3, z=0.0)
	assert func.to_string(ctx) == "(5, 3) -> (x > y) -> True"


def test_func_with_logical_operations() -> None:
	"""Test Func with logical operations."""
	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	expr = (x > 0) & (y < 10)
	func = expr.to_func()

	assert func.to_string() == "(x, y) -> ((x > 0) & (y < 10))"

	ctx = FuncCtx(x=5, y=3, z=0.0)
	assert func.to_string(ctx) == "(5, 3) -> ((x > 0) & (y < 10)) -> True"


def test_empty_args_func() -> None:
	"""Test Func with no arguments (constants only)."""
	expr = Const("pi", 3.14159)
	func = Func((), expr)

	assert func.to_string() == "() -> pi:3.14159"

	ctx = FuncCtx(x=0, y=0, z=0.0)
	assert func.to_string(ctx) == "() -> pi:3.14159 -> 3.14159"


def test_func_integration_with_existing_features() -> None:
	"""Test that Func works with other Mahonia features."""
	from mahonia import Predicate

	x = Var[int, FuncCtx]("x")
	y = Var[int, FuncCtx]("y")

	# Create a predicate and convert to function
	expr = (x + y) > 10
	predicate = Predicate("Sum greater than 10", expr)
	func = expr.to_func()

	assert func.to_string() == "(x, y) -> ((x + y) > 10)"

	ctx_true = FuncCtx(x=7, y=8, z=0.0)
	ctx_false = FuncCtx(x=2, y=3, z=0.0)

	assert "True" in func.to_string(ctx_true)
	assert "False" in func.to_string(ctx_false)

	# Verify predicate still works
	assert predicate.unwrap(ctx_true) is True
	assert predicate.unwrap(ctx_false) is False


class TestExtractVarsResult:
	def test_pow_extracts_var(self) -> None:
		x = Var[int, FuncCtx]("x")
		expr = x**2
		assert extract_vars((), expr) == (x,)

	def test_nested_result_extracts_all(self) -> None:
		x = Var[int, FuncCtx]("x")
		y = Var[int, FuncCtx]("y")
		expr = (x + y) ** 2
		vars = extract_vars((), expr)
		assert len(vars) == 2
		assert x in vars
		assert y in vars

	def test_mod_extracts_both(self) -> None:
		x = Var[int, FuncCtx]("x")
		y = Var[int, FuncCtx]("y")
		expr = x % y
		vars = extract_vars((), expr)
		assert len(vars) == 2


def test_extract_vars_from_not() -> None:
	x = Var[bool, FuncCtx]("x")
	expr: Not[FuncCtx] = ~x
	vars = extract_vars((), expr)
	assert len(vars) == 1


def test_extract_vars_from_contains() -> None:
	from mahonia import SizedIterable

	class SearchCtx(NamedTuple):
		target: int
		values: list[int]

	target_var = Var[int, SearchCtx]("target")
	values_var = Var[SizedIterable[int], SearchCtx]("values")
	expr = Contains(target_var, values_var)
	vars = extract_vars((), expr)
	assert len(vars) == 2


def test_extract_vars_from_map_expr() -> None:
	from mahonia import SizedIterable

	class MapCtx(NamedTuple):
		n: int
		nums: list[int]

	n = Var[int, MapCtx]("n")
	nums = Var[SizedIterable[int], MapCtx]("nums")
	map_expr = (n > 0).map(nums)
	assert isinstance(map_expr, MapExpr)
	vars = extract_vars((), map_expr)
	assert len(vars) == 2


def test_extract_vars_from_filter_expr() -> None:
	from mahonia import SizedIterable

	class FilterCtx(NamedTuple):
		n: int
		nums: list[int]

	n = Var[int, FilterCtx]("n")
	nums = Var[SizedIterable[int], FilterCtx]("nums")
	predicate = (n > 0).to_func()
	expr = FilterExpr(predicate, nums)
	vars = extract_vars((), expr)
	assert len(vars) == 2


def test_extract_vars_from_any_all_min_max() -> None:
	from mahonia import Add, AllExpr, AnyExpr, FoldLExpr, MaxExpr, MinExpr, SizedIterable

	class Ctx(NamedTuple):
		values: list[int]
		flags: list[bool]

	values = Var[SizedIterable[int], Ctx]("values")
	flags = Var[SizedIterable[bool], Ctx]("flags")

	assert extract_vars((), AnyExpr(flags)) == (flags,)
	assert extract_vars((), AllExpr(flags)) == (flags,)
	assert extract_vars((), MinExpr(values)) == (values,)
	assert extract_vars((), MaxExpr(values)) == (values,)
	assert extract_vars((), FoldLExpr(Add, values)) == (values,)


def test_extract_vars_map_duplicate_var() -> None:
	from mahonia import SizedIterable

	class MapCtx(NamedTuple):
		n: int
		nums: list[int]

	n = Var[int, MapCtx]("n")
	nums = Var[SizedIterable[int], MapCtx]("nums")
	map_expr = (n > 0).map(nums)
	assert isinstance(map_expr, MapExpr)
	vars_tuple = extract_vars((n,), map_expr)
	assert len(vars_tuple) == 2
	assert vars_tuple[0] is n
	assert vars_tuple[1] is nums


def test_extract_vars_filter_duplicate_var() -> None:
	from mahonia import SizedIterable

	class FilterCtx(NamedTuple):
		n: int
		nums: list[int]

	n = Var[int, FilterCtx]("n")
	nums = Var[SizedIterable[int], FilterCtx]("nums")
	predicate = (n > 0).to_func()
	expr = FilterExpr(predicate, nums)
	vars_tuple = extract_vars((n,), expr)
	assert len(vars_tuple) == 2
	assert vars_tuple[0] is n
	assert vars_tuple[1] is nums
