#
# Copyright © 2025 Aitomatic, Inc.
#
# This source code is licensed under the license found in the LICENSE file in the root directory of this source tree
#
from typing import cast

import pytest

from dana.core.lang.ast import (
    Assignment,
    AttributeAccess,
    BinaryExpression,
    BinaryOperator,
    ForLoop,
    FStringExpression,
    FunctionCall,
    Identifier,
    ListLiteral,
    LiteralExpression,
    Statement,
    UnaryExpression,
)
from dana.core.lang.parser.transformer.expression_transformer import ExpressionTransformer
from dana.core.lang.parser.transformer.fstring_transformer import FStringTransformer
from dana.core.lang.parser.transformer.statement_transformer import StatementTransformer
from dana.core.lang.parser.transformer.variable_transformer import VariableTransformer

# 1. VariableTransformer tests


def make_token(type_, value):
    from lark import Token

    return Token(type_, value)


def test_simple_name():
    vt = VariableTransformer()
    token = make_token("NAME", "foo")
    result = vt.simple_name([token])
    assert isinstance(result, Identifier)
    assert result.name == "foo"  # No automatic local: prefix


def test_dotted_access():
    vt = VariableTransformer()
    token1 = make_token("NAME", "foo")
    token2 = make_token("NAME", "bar")
    result = vt.dotted_access([token1, token2])
    assert isinstance(result, AttributeAccess)
    assert isinstance(result.object, Identifier)
    assert result.object.name == "foo"  # No automatic local: prefix
    assert result.attribute == "bar"


def test_scoped_var():
    vt = VariableTransformer()
    scope = make_token("SCOPE", "private")
    var = make_token("NAME", "x")
    result = vt.scoped_var([scope, var])
    assert isinstance(result, Identifier)
    assert result.name == "private:x"


def test_identifier_utility():
    vt = VariableTransformer()
    token1 = make_token("NAME", "foo")
    token2 = make_token("NAME", "bar")
    result = vt.identifier([token1, token2])
    assert isinstance(result, Identifier)
    assert result.name == "foo.bar"  # No automatic local: prefix


def test_variable_scoped_and_dotted():
    vt = VariableTransformer()
    token_scope = make_token("SCOPE", "private")
    token_name = make_token("NAME", "foo")
    result_scoped = vt.scoped_var([token_scope, token_name])
    result_dotted = vt.dotted_access([token_name, token_name])
    assert result_scoped.name == "private:foo"
    assert isinstance(result_dotted, AttributeAccess)
    assert isinstance(result_dotted.object, Identifier)
    assert result_dotted.object.name == "foo"  # No automatic local: prefix
    assert result_dotted.attribute == "foo"  # Both tokens are "foo", not "bar"


# 2. FStringTransformer tests


def test_fstring_text_only():
    ft = FStringTransformer()
    token = make_token("F_STRING", 'f"Hello world!"')
    result = ft.fstring([token])
    assert isinstance(result, LiteralExpression)
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert fexpr.parts == ["Hello world!"]


def test_fstring_with_identifier():
    ft = FStringTransformer()
    token = make_token("F_STRING", 'f"Hello, {foo}"')
    result = ft.fstring([token])
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert fexpr.parts[0] == "Hello, "
    assert isinstance(fexpr.parts[1], Identifier)
    assert fexpr.parts[1].name == "foo"  # Same as regular expressions - no automatic prefix


def test_fstring_with_binary_expression():
    ft = FStringTransformer()
    token = make_token("F_STRING", 'f"Sum: {x + y}"')
    result = ft.fstring([token])
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert fexpr.parts[0] == "Sum: "
    bexpr = fexpr.parts[1]
    assert isinstance(bexpr, BinaryExpression)
    assert bexpr.operator == BinaryOperator.ADD
    assert isinstance(bexpr.left, Identifier)
    assert bexpr.left.name == "x"  # Same as regular expressions - no automatic prefix
    assert isinstance(bexpr.right, Identifier)
    assert bexpr.right.name == "y"  # Same as regular expressions - no automatic prefix


def test_fstring_unbalanced_braces():
    ft = FStringTransformer()
    token = make_token("F_STRING", 'f"Hello {foo"')
    with pytest.raises(ValueError):
        ft.fstring([token])


def test_fstring_multiple_and_expression():
    ft = FStringTransformer()
    token = make_token("F_STRING", 'f"A: {a}, B: {b + 1}"')
    result = ft.fstring([token])
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert len(fexpr.parts) == 4
    assert isinstance(fexpr.parts[1], Identifier)
    assert isinstance(fexpr.parts[3], BinaryExpression)


def test_fstring_complex_expressions():
    ft = FStringTransformer()
    # Test nested expressions
    token = make_token("F_STRING", 'f"Nested: {x * (y + 2)}"')
    result = ft.fstring([token])
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert isinstance(fexpr.parts[1], BinaryExpression)
    assert fexpr.parts[1].operator == BinaryOperator.MULTIPLY

    # Test multiple operations
    token = make_token("F_STRING", 'f"Math: {a + b * c}"')
    result = ft.fstring([token])
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert isinstance(fexpr.parts[1], BinaryExpression)

    # Test with comparison
    token = make_token("F_STRING", 'f"Check: {x > 10}"')
    result = ft.fstring([token])
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert isinstance(fexpr.parts[1], BinaryExpression)
    assert fexpr.parts[1].operator == BinaryOperator.GREATER_THAN


def test_fstring_starts_with_expression():
    # Test an f-string that starts immediately with an expression
    ft = FStringTransformer()
    token = make_token("F_STRING", 'f"{a}"')
    result = ft.fstring([token])
    fexpr = result.value
    assert isinstance(fexpr, FStringExpression)
    assert isinstance(fexpr.parts[0], Identifier)
    assert fexpr.parts[0].name == "a"  # Same as regular expressions - no automatic prefix

    # Test with expression followed by text
    token2 = make_token("F_STRING", 'f"{a} text"')
    result2 = ft.fstring([token2])
    fexpr2 = result2.value
    assert isinstance(fexpr2, FStringExpression)
    assert isinstance(fexpr2.parts[0], Identifier)
    assert fexpr2.parts[0].name == "a"  # Same as regular expressions - no automatic prefix
    assert fexpr2.parts[1] == " text"


# 3. StatementTransformer tests


def test_statement_assignment():
    st = StatementTransformer()
    target = Identifier(name="x")  # No automatic local: prefix
    value = LiteralExpression(value=1)
    stmt = st.assignment([target, value])
    assert isinstance(stmt.value, LiteralExpression)
    assert stmt.value.value == 1


def test_statement_conditional():
    st = StatementTransformer()
    cond = BinaryExpression(
        left=Identifier("x"), operator=BinaryOperator.GREATER_THAN, right=LiteralExpression(0)
    )  # No automatic local: prefix
    if_body = [Assignment(target=Identifier("x"), value=LiteralExpression(1))]  # No automatic local: prefix
    else_body = [Assignment(target=Identifier("x"), value=LiteralExpression(2))]  # No automatic local: prefix
    conditional = st.conditional([[cond] + if_body, else_body])
    assert conditional.condition.operator == BinaryOperator.GREATER_THAN
    assert len(conditional.body) == 1
    assert len(conditional.else_body) == 1


def test_statement_while_loop():
    st = StatementTransformer()
    cond = BinaryExpression(
        left=Identifier("x"), operator=BinaryOperator.LESS_THAN, right=LiteralExpression(10)
    )  # No automatic local: prefix
    body = [
        Assignment(
            target=Identifier("x"),  # No automatic local: prefix
            value=BinaryExpression(
                left=Identifier("x"), operator=BinaryOperator.ADD, right=LiteralExpression(1)
            ),  # No automatic local: prefix
        )
    ]
    loop = st.while_stmt([cond] + body)
    assert loop.condition.operator == BinaryOperator.LESS_THAN
    assert len(loop.body) == 1


def test_statement_for_loop():
    target = Identifier(name="i")
    iterable = Identifier(name="xs")  # No automatic local: prefix
    body: list[Statement] = [Assignment(target=Identifier("i"), value=LiteralExpression(1))]  # No automatic local: prefix
    loop = ForLoop(target=target, iterable=iterable, body=body)
    assert loop.target.name == "i"
    assert loop.iterable.name == "xs"  # No automatic local: prefix
    assert len(loop.body) == 1


def test_statement_function_def():
    st = StatementTransformer()
    name = make_token("NAME", "foo")
    parameters = [Identifier("x")]  # No automatic local: prefix
    body = [Assignment(target=Identifier("x"), value=LiteralExpression(1))]  # No automatic local: prefix
    func = st.function_def([name, parameters, body])
    assert func.name.name == "foo"
    assert func.parameters[0].name == "x"  # No automatic local: prefix
    assert len(func.body) == 1


def test_statement_return_break_continue_pass():
    st = StatementTransformer()
    ret = st.return_stmt([LiteralExpression(1)]).value
    assert isinstance(ret, LiteralExpression)
    assert ret.value == 1
    assert st.break_stmt([]) is not None
    assert st.continue_stmt([]) is not None
    assert st.pass_stmt([]) is not None


# 4. ExpressionTransformer tests


def test_expression_arithmetic():
    et = ExpressionTransformer()
    a = LiteralExpression(2)
    b = LiteralExpression(3)
    expr_add = et.sum_expr([a, "+", b])
    expr_mul = et.term([a, "*", b])
    assert isinstance(expr_add, BinaryExpression)
    expr_add = cast(BinaryExpression, expr_add)
    assert expr_add.operator == BinaryOperator.ADD
    assert isinstance(expr_mul, BinaryExpression)
    expr_mul = cast(BinaryExpression, expr_mul)
    assert expr_mul.operator == BinaryOperator.MULTIPLY


def test_expression_logical():
    et = ExpressionTransformer()
    a = Identifier("a")  # No automatic local: prefix
    b = Identifier("b")  # No automatic local: prefix
    c = Identifier("c")  # No automatic local: prefix
    expr = et.or_expr([et.and_expr([a, b]), c])
    assert isinstance(expr, BinaryExpression)
    assert expr.operator == BinaryOperator.OR
    assert isinstance(expr.left, BinaryExpression)
    assert expr.left.operator == BinaryOperator.AND


def test_expression_comparison():
    et = ExpressionTransformer()
    x = Identifier("x")  # No automatic local: prefix
    zero = LiteralExpression(0)
    expr = et.comparison([x, ">", zero])
    assert isinstance(expr, BinaryExpression)
    assert expr.operator == BinaryOperator.GREATER_THAN
    assert expr.left.name == "x"  # No automatic local: prefix
    assert expr.right.value == 0


def test_expression_unary_not():
    et = ExpressionTransformer()
    x = Identifier("x")  # No automatic local: prefix
    expr = et.not_expr(["not", x])
    assert isinstance(expr, UnaryExpression)
    assert expr.operator == "not"
    assert expr.operand.name == "x"  # No automatic local: prefix


def test_expression_arithmetic_ops():
    et = ExpressionTransformer()
    a = LiteralExpression(2)
    b = LiteralExpression(3)
    expr_add = et.sum_expr([a, "+", b])
    expr_mul = et.term([a, "*", b])
    expr_add = cast(BinaryExpression, expr_add)
    assert expr_add.operator == BinaryOperator.ADD
    expr_mul = cast(BinaryExpression, expr_mul)
    assert expr_mul.operator == BinaryOperator.MULTIPLY


def test_expression_comparison_and_logical():
    et = ExpressionTransformer()
    a = Identifier("a")  # No automatic local: prefix
    b = Identifier("b")  # No automatic local: prefix
    expr_eq = et.comparison([a, "==", b])
    expr_and = et.and_expr([a, "and", b])
    expr_not = et.not_expr(["not", a])
    assert isinstance(expr_eq, BinaryExpression)
    expr_eq = cast(BinaryExpression, expr_eq)
    assert expr_eq.operator == BinaryOperator.EQUALS
    assert isinstance(expr_and, BinaryExpression)
    expr_and = cast(BinaryExpression, expr_and)
    assert expr_and.operator == BinaryOperator.AND
    assert isinstance(expr_not, UnaryExpression)
    assert expr_not.operator == "not"


def test_expression_function_call_and_attr():
    call = FunctionCall(name="foo", args={"__positional": [LiteralExpression(1), LiteralExpression(2)]})  # No automatic local: prefix
    assert call.name == "foo"  # No automatic local: prefix
    assert len(call.args["__positional"]) == 2
    assert call.args["__positional"][0].value == 1
    assert call.args["__positional"][1].value == 2


def test_expression_literals_and_collections():
    et = ExpressionTransformer()
    assert cast(LiteralExpression, et.literal([LiteralExpression(True)])).value is True
    assert cast(LiteralExpression, et.literal([LiteralExpression(None)])).value is None
    list_literal = cast(ListLiteral, et.list([LiteralExpression(1), LiteralExpression(2)]))
    assert isinstance(list_literal, ListLiteral)
    assert len(list_literal.items) == 2
