"""Tests for Enhancement 3: @ret wiring for macros (macro-enh.E3.*).

Tests verify:
- E3.1: Grammar accepts |> output list on macro_call_stmt
- E3.2: Lower stores output_dests on IRMacroCall
- E3.3: Expand rewrites @ret edges to concrete destinations (positional)
- E3.4: Expand rewrites @ret_name edges to concrete destinations (named)
- E3.5: Unmatched @ret marker produces MACRO error
- E3.6: Macro with @ret but no |> at call site produces MACRO error
- E3.7: Multiple @ret markers with mixed positional/named outputs
- E3.8: @ret in nested macro invocation
"""

from pathlib import Path

from lark import Lark

from asm.expand import expand
from asm.lower import lower
from asm.errors import ErrorCategory
from asm.ir import IRMacroCall


def _get_parser():
    grammar_path = Path(__file__).parent.parent / "dfasm.lark"
    return Lark(
        grammar_path.read_text(),
        parser="earley",
        propagate_positions=True,
    )


def parse_and_lower(source: str):
    parser = _get_parser()
    tree = parser.parse(source)
    return lower(tree)


def parse_lower_expand(source: str):
    graph = parse_and_lower(source)
    return expand(graph)


class TestE31_GrammarAcceptsOutputList:
    """E3.1: Grammar accepts |> output list on macro_call_stmt.

    Note: #macro |> &dest (no args, single positional output) is ambiguous
    with plain_edge syntax and parses as plain_edge. Disambiguation requires
    either named outputs (name=&dest) or at least one argument before |>.
    """

    def test_macro_call_with_named_output_no_args(self):
        """#macro |> name=&dest parses as macro_call_stmt (named output disambiguates)."""
        source = """
        @system pe=1, sm=1
        #simple |> {
            &g <| pass
            &g |> @ret_out
        }
        &sink <| add
        #simple |> out=&sink
        """
        graph = parse_and_lower(source)
        assert not graph.errors
        assert len(graph.macro_calls) == 1

    def test_macro_call_with_named_outputs(self):
        """#macro |> name=&dest, name2=&dest2 parses without error."""
        source = """
        @system pe=1, sm=1
        #dual |> {
            &g <| pass
            &g |> @ret_body
            &g |> @ret_exit:R
        }
        &body_sink <| add
        &exit_sink <| add
        #dual |> body=&body_sink, exit=&exit_sink
        """
        graph = parse_and_lower(source)
        assert not graph.errors

    def test_macro_call_with_args_and_positional_output(self):
        """#macro arg |> &dest parses as macro_call_stmt (args disambiguate)."""
        source = """
        @system pe=1, sm=1
        #with_arg val |> {
            &g <| const, ${val}
            &g |> @ret
        }
        &sink <| add
        #with_arg 42 |> &sink
        """
        graph = parse_and_lower(source)
        assert not graph.errors
        assert len(graph.macro_calls) == 1


class TestE32_LowerStoresOutputDests:
    """E3.2: Lower stores output_dests on IRMacroCall."""

    def test_output_dests_populated_with_arg(self):
        """IRMacroCall.output_dests contains output destinations (arg disambiguates)."""
        source = """
        @system pe=1, sm=1
        #simple val |> {
            &g <| const, ${val}
            &g |> @ret
        }
        &sink <| add
        #simple 1 |> &sink
        """
        graph = parse_and_lower(source)
        assert not graph.errors
        assert len(graph.macro_calls) == 1
        call = graph.macro_calls[0]
        assert len(call.output_dests) >= 1

    def test_named_output_dests(self):
        """Named outputs stored as {"name": ..., "ref": ...} dicts."""
        source = """
        @system pe=1, sm=1
        #dual |> {
            &g <| pass
        }
        &sink <| add
        #dual |> body=&sink
        """
        graph = parse_and_lower(source)
        assert not graph.errors
        call = graph.macro_calls[0]
        assert len(call.output_dests) >= 1
        output = call.output_dests[0]
        assert isinstance(output, dict)
        assert output.get("name") == "body"


class TestE33_ExpandRewritesPositionalRet:
    """E3.3: Expand rewrites @ret edges to concrete destinations (positional)."""

    def test_bare_ret_rewritten(self):
        """@ret in macro body becomes concrete &sink after expansion."""
        source = """
        @system pe=1, sm=1
        #simple val |> {
            &g <| const, ${val}
            &g |> @ret
        }
        &sink <| add
        #simple 1 |> &sink
        """
        graph = parse_lower_expand(source)
        assert not graph.errors
        # Find the edge from the expanded &g to &sink
        ret_edges = [
            e for e in graph.edges
            if e.dest == "&sink" and "#simple_0.&g" in e.source
        ]
        assert len(ret_edges) == 1

    def test_bare_ret_with_named_output(self):
        """@ret in macro body resolved via named output (no args needed)."""
        source = """
        @system pe=1, sm=1
        #simple |> {
            &g <| pass
            &g |> @ret_out
        }
        &sink <| add
        #simple |> out=&sink
        """
        graph = parse_lower_expand(source)
        assert not graph.errors
        ret_edges = [
            e for e in graph.edges
            if e.dest == "&sink" and "#simple_0.&g" in e.source
        ]
        assert len(ret_edges) == 1


class TestE34_ExpandRewritesNamedRet:
    """E3.4: Expand rewrites @ret_name edges to concrete destinations (named)."""

    def test_named_ret_body(self):
        """@ret_body in macro resolves to body=&body_sink output."""
        source = """
        @system pe=1, sm=1
        #dual |> {
            &g <| pass
            &g |> @ret_body
            &g |> @ret_exit:R
        }
        &body_sink <| add
        &exit_sink <| add
        #dual |> body=&body_sink, exit=&exit_sink
        """
        graph = parse_lower_expand(source)
        assert not graph.errors
        # Find edges from expanded #dual_0.&g
        g_edges = [e for e in graph.edges if "#dual_0.&g" in e.source]
        dests = {e.dest for e in g_edges}
        assert "&body_sink" in dests
        assert "&exit_sink" in dests

    def test_mixed_named_outputs(self):
        """Multiple @ret_name markers all resolve correctly."""
        source = """
        @system pe=1, sm=1
        #triple |> {
            &a <| pass
            &b <| pass
            &c <| pass
            &a |> @ret_x
            &b |> @ret_y
            &c |> @ret_z
        }
        &x_sink <| add
        &y_sink <| add
        &z_sink <| add
        #triple |> x=&x_sink, y=&y_sink, z=&z_sink
        """
        graph = parse_lower_expand(source)
        assert not graph.errors
        all_dests = {e.dest for e in graph.edges}
        assert "&x_sink" in all_dests
        assert "&y_sink" in all_dests
        assert "&z_sink" in all_dests


class TestE35_UnmatchedRetError:
    """E3.5: Unmatched @ret marker produces MACRO error."""

    def test_unmatched_named_ret(self):
        """@ret_missing has no matching named output -> MACRO error."""
        source = """
        @system pe=1, sm=1
        #bad |> {
            &g <| pass
            &g |> @ret_missing
        }
        &sink <| add
        #bad |> wrong_name=&sink
        """
        graph = parse_lower_expand(source)
        macro_errors = [e for e in graph.errors if e.category == ErrorCategory.MACRO]
        assert len(macro_errors) >= 1
        assert "@ret_missing" in macro_errors[0].message

    def test_extra_ret_marker(self):
        """Macro has @ret_body and @ret_exit but call only provides body output."""
        source = """
        @system pe=1, sm=1
        #dual |> {
            &a <| pass
            &b <| pass
            &a |> @ret_body
            &b |> @ret_exit
        }
        &sink <| add
        #dual |> body=&sink
        """
        graph = parse_lower_expand(source)
        macro_errors = [e for e in graph.errors if e.category == ErrorCategory.MACRO]
        assert len(macro_errors) >= 1
        assert "@ret_exit" in macro_errors[0].message


class TestE36_NoOutputWiringError:
    """E3.6: Macro with @ret but no |> at call site produces MACRO error."""

    def test_ret_without_output_wiring(self):
        """Macro body uses @ret but call has no |> -> error."""
        source = """
        @system pe=1, sm=1
        #needs_output |> {
            &g <| pass
            &g |> @ret
        }
        #needs_output
        """
        graph = parse_lower_expand(source)
        macro_errors = [e for e in graph.errors if e.category == ErrorCategory.MACRO]
        assert len(macro_errors) >= 1
        assert "output" in macro_errors[0].message.lower() or "@ret" in macro_errors[0].message


    def test_bare_ret_with_only_named_outputs(self):
        """Bare @ret with only named outputs at call site -> MACRO error."""
        source = """
        @system pe=1, sm=1
        #test val |> {
            &g <| const, ${val}
            &g |> @ret
        }
        &sink <| add
        #test 1 |> out=&sink
        """
        graph = parse_lower_expand(source)
        macro_errors = [e for e in graph.errors if e.category == ErrorCategory.MACRO]
        assert len(macro_errors) >= 1
        assert "@ret" in macro_errors[0].message


class TestE37_MultipleRetMarkers:
    """E3.7: Multiple @ret markers with mixed positional/named outputs."""

    def test_named_and_positional_ret(self):
        """Macro with @ret_main and @ret_extra resolves to named outputs."""
        source = """
        @system pe=1, sm=1
        #mixed |> {
            &a <| pass
            &b <| pass
            &a |> @ret_main
            &b |> @ret_extra
        }
        &main_sink <| add
        &extra_sink <| add
        #mixed |> main=&main_sink, extra=&extra_sink
        """
        graph = parse_lower_expand(source)
        assert not graph.errors
        all_dests = {e.dest for e in graph.edges}
        assert "&main_sink" in all_dests
        assert "&extra_sink" in all_dests


class TestE38_NestedMacroRet:
    """E3.8: @ret in nested macro invocation."""

    def test_nested_macro_ret_resolves(self):
        """Inner macro's @ret resolves independently from outer macro's @ret."""
        source = """
        @system pe=1, sm=1
        #inner val |> {
            &i <| const, ${val}
            &i |> @ret
        }
        #outer val |> {
            &o <| pass
            #inner ${val} |> &o
            &o |> @ret_out
        }
        &final_sink <| add
        #outer 1 |> out=&final_sink
        """
        graph = parse_lower_expand(source)
        assert not graph.errors
        # The outer @ret_out should resolve to &final_sink
        # The inner @ret should resolve to #outer_0.&o
        outer_ret_edges = [
            e for e in graph.edges
            if e.dest == "&final_sink"
        ]
        assert len(outer_ret_edges) >= 1

    def test_inner_ret_failure_does_not_cascade_to_outer(self):
        """Inner macro @ret failure should not cause spurious outer errors.

        Uses named output (@ret_out / out=&sink) to disambiguate from plain_edge.
        """
        source = """
        @system pe=1, sm=1
        #inner val |> {
            &i <| const, ${val}
            &i |> @ret
        }
        #outer |> {
            &o <| pass
            #inner 1
            &o |> @ret_out
        }
        &sink <| add
        #outer |> out=&sink
        """
        graph = parse_lower_expand(source)
        macro_errors = [e for e in graph.errors if e.category == ErrorCategory.MACRO]
        assert len(macro_errors) == 1, f"Expected 1 error (inner #inner), got {len(macro_errors)}: {[e.message for e in macro_errors]}"
        assert "#inner" in macro_errors[0].message
        sink_edges = [e for e in graph.edges if e.dest == "&sink"]
        assert len(sink_edges) >= 1