"""End-to-end integration tests: assemble source, emulate, verify results.

Tests verify:
- or1-asm.AC9.1: CONST→ADD chain produces correct sum
- or1-asm.AC9.2: SM round-trip (write, deferred read) returns correct value
- or1-asm.AC9.3: Cross-PE routing delivers token to destination PE
- or1-asm.AC9.4: SWITCH routing sends data to taken path, trigger to not_taken
- or1-asm.AC9.5: Token stream mode produces identical results to direct mode
- or1-asm.AC10.5: Auto-placed (unplaced) programs assemble and execute correctly
"""

import pytest
import simpy

from asm import assemble, assemble_to_tokens
from emu import build_topology
from tokens import PELocalWriteToken, MonadToken, SMToken


def run_program_direct(source: str, until: int = 1000) -> dict:
    """Assemble source in direct mode, run through emulator.

    Args:
        source: dfasm source code as a string
        until: Simulation timeout in time units (default: 1000)

    Returns:
        Dict mapping PE ID to list of output tokens from that PE
    """
    result = assemble(source)
    env = simpy.Environment()
    sys = build_topology(env, result.pe_configs, result.sm_configs)

    # Inject setup tokens first (frame/IRAM initialization)
    for setup in result.setup_tokens:
        sys.inject(setup)

    # Then inject seed tokens
    for seed in result.seed_tokens:
        sys.inject(seed)

    env.run(until=until)

    # Collect output from each PE's output_log
    outputs = {}
    for pe_id, pe in sys.pes.items():
        outputs[pe_id] = list(pe.output_log)

    return outputs


def run_program_tokens(source: str, until: int = 1000) -> dict:
    """Assemble source to token stream mode, run through emulator.

    Builds topology normally, injects all tokens, runs simulation, and collects
    output from each PE's output_log.

    Args:
        source: dfasm source code as a string
        until: Simulation timeout in time units (default: 1000)

    Returns:
        Dict mapping PE ID to list of output tokens collected from that PE
    """
    tokens = assemble_to_tokens(source)
    env = simpy.Environment()

    # Extract PE and SM counts from tokens
    max_pe_id = 0
    max_sm_id = 0

    for token in tokens:
        if isinstance(token, SMToken):
            max_sm_id = max(max_sm_id, token.target)
        elif isinstance(token, PELocalWriteToken):
            max_pe_id = max(max_pe_id, token.target)
        elif isinstance(token, MonadToken):
            max_pe_id = max(max_pe_id, token.target)

    # Create minimal PE configs (empty IRAM - will be filled by PELocalWriteToken)
    from emu.types import PEConfig, SMConfig
    pe_configs = [PEConfig(i, {}) for i in range(max_pe_id + 1)]
    sm_configs = [SMConfig(i) for i in range(max_sm_id + 1)]

    sys = build_topology(env, pe_configs, sm_configs)

    # Inject tokens in order (do NOT modify route_table)
    for token in tokens:
        sys.inject(token)

    env.run(until=until)

    # Collect output from each PE's output_log
    outputs = {}
    for i in range(max_pe_id + 1):
        outputs[i] = list(sys.pes[i].output_log)

    return outputs


class TestAC91ConstToAddChain:
    """AC9.1: CONST→ADD chain produces correct sum."""

    def test_const_add_chain_direct(self):
        """Direct mode: two const nodes feed an add node, should produce sum (10)."""
        source = """
@system pe=2, sm=0
&c1|pe0 <| const, 3
&c2|pe0 <| const, 7
&result|pe0 <| add
&output|pe1 <| pass
&c1|pe0 |> &result|pe0:L
&c2|pe0 |> &result|pe0:R
&result|pe0 |> &output|pe1:L
"""
        outputs = run_program_direct(source)
        # Result PE produces the sum: 3 + 7 = 10
        result_outputs = outputs[0]
        assert any(t.data == 10 for t in result_outputs if hasattr(t, 'data')), \
            f"Expected result 10 in PE0 outputs, got {[t.data for t in result_outputs if hasattr(t, 'data')]}"

    def test_const_add_chain_tokens(self):
        """Token stream mode: const add chain should produce sum (10)."""
        source = """
@system pe=2, sm=0
&c1|pe0 <| const, 3
&c2|pe0 <| const, 7
&result|pe0 <| add
&output|pe1 <| pass
&c1|pe0 |> &result|pe0:L
&c2|pe0 |> &result|pe0:R
&result|pe0 |> &output|pe1:L
"""
        outputs = run_program_tokens(source)
        # Result PE produces the sum: 3 + 7 = 10
        result_outputs = outputs[0]
        assert any(t.data == 10 for t in result_outputs if hasattr(t, 'data')), \
            f"Expected result 10 in PE0 outputs, got {[t.data for t in result_outputs if hasattr(t, 'data')]}"


class TestAC92SMMRoundTrip:
    """AC9.2: SM round-trip (write, deferred read) returns correct value."""

    def test_sm_read_deferred_direct(self):
        """Direct mode: SM write+read round-trip returns stored value 0x42."""
        source = """
@system pe=3, sm=1
@val|sm0:5 = 0x42
&trigger|pe0 <| const, 1
&reader|pe0 <| read, 5
&relay|pe1 <| pass
&sink|pe2 <| pass
&trigger|pe0 |> &reader|pe0:L
&reader|pe0 |> &relay|pe1:L
&relay|pe1 |> &sink|pe2:L
"""
        outputs = run_program_direct(source)
        relay_outputs = [t.data for t in outputs[1] if hasattr(t, 'data')]
        assert 66 in relay_outputs, \
            f"Expected SM read value 66 (0x42) in PE1 outputs, got {relay_outputs}"

    def test_sm_read_deferred_tokens(self):
        """Token stream mode: SM write+read round-trip returns stored value 0x42."""
        source = """
@system pe=3, sm=1
@val|sm0:5 = 0x42
&trigger|pe0 <| const, 1
&reader|pe0 <| read, 5
&relay|pe1 <| pass
&sink|pe2 <| pass
&trigger|pe0 |> &reader|pe0:L
&reader|pe0 |> &relay|pe1:L
&relay|pe1 |> &sink|pe2:L
"""
        outputs = run_program_tokens(source)
        relay_outputs = [t.data for t in outputs[1] if hasattr(t, 'data')]
        assert 66 in relay_outputs, \
            f"Expected SM read value 66 (0x42) in PE1 outputs, got {relay_outputs}"


class TestAC93CrossPERouting:
    """AC9.3: Cross-PE routing delivers token to destination PE."""

    def test_cross_pe_routing_direct(self):
        """Direct mode: cross-PE routing assembles and PE0 emits token to PE1."""
        source = """
@system pe=3, sm=0
&source|pe0 <| const, 99
&dest|pe1 <| pass
&output|pe2 <| pass
&source|pe0 |> &dest|pe1:L
&dest|pe1 |> &output|pe2:L
"""
        outputs = run_program_direct(source)
        # PE0 should emit the constant value 99 to PE1
        source_outputs = outputs[0]
        assert any(t.data == 99 for t in source_outputs if hasattr(t, 'data')), \
            f"Expected value 99 in PE0 outputs, got {[t.data for t in source_outputs if hasattr(t, 'data')]}"

    def test_cross_pe_routing_tokens(self):
        """Token stream mode: cross-PE routing assembles and PE0 emits token to PE1."""
        source = """
@system pe=3, sm=0
&source|pe0 <| const, 99
&dest|pe1 <| pass
&output|pe2 <| pass
&source|pe0 |> &dest|pe1:L
&dest|pe1 |> &output|pe2:L
"""
        outputs = run_program_tokens(source)
        # PE0 should emit the constant value 99 to PE1
        source_outputs = outputs[0]
        assert any(t.data == 99 for t in source_outputs if hasattr(t, 'data')), \
            f"Expected value 99 in PE0 outputs, got {[t.data for t in source_outputs if hasattr(t, 'data')]}"


class TestAC94SwitchRouting:
    """AC9.4: SWITCH routing sends data to taken path, trigger to not_taken."""

    def test_switch_equal_inputs_direct(self):
        """Direct mode: SWITCH correctly routes data to taken and trigger to not_taken."""
        source = """
@system pe=3, sm=0
&val|pe0 <| const, 5
&cmp|pe0 <| const, 5
&branch|pe0 <| sweq
&taken|pe1 <| pass
&not_taken|pe1 <| pass
&output|pe2 <| pass
&val|pe0 |> &branch|pe0:L
&cmp|pe0 |> &branch|pe0:R
&branch|pe0:L |> &taken|pe1:L
&branch|pe0:R |> &not_taken|pe1:L
&taken|pe1 |> &output|pe2:L
&not_taken|pe1 |> &output|pe2:R
"""
        outputs = run_program_direct(source)
        # PE0 should emit data (5) to taken and trigger (0) to not_taken
        pe0_outputs = [t.data for t in outputs[0] if hasattr(t, 'data')]
        assert 5 in pe0_outputs, f"Expected data value 5 emitted from PE0, got {pe0_outputs}"
        assert 0 in pe0_outputs, f"Expected trigger value 0 emitted from PE0, got {pe0_outputs}"

    def test_switch_equal_inputs_tokens(self):
        """Token stream mode: SWITCH correctly routes data to taken and trigger to not_taken."""
        source = """
@system pe=3, sm=0
&val|pe0 <| const, 5
&cmp|pe0 <| const, 5
&branch|pe0 <| sweq
&taken|pe1 <| pass
&not_taken|pe1 <| pass
&output|pe2 <| pass
&val|pe0 |> &branch|pe0:L
&cmp|pe0 |> &branch|pe0:R
&branch|pe0:L |> &taken|pe1:L
&branch|pe0:R |> &not_taken|pe1:L
&taken|pe1 |> &output|pe2:L
&not_taken|pe1 |> &output|pe2:R
"""
        outputs = run_program_tokens(source)
        # PE0 should emit data (5) to taken and trigger (0) to not_taken
        pe0_outputs = [t.data for t in outputs[0] if hasattr(t, 'data')]
        assert 5 in pe0_outputs, f"Expected data value 5 emitted from PE0, got {pe0_outputs}"
        assert 0 in pe0_outputs, f"Expected trigger value 0 emitted from PE0, got {pe0_outputs}"


class TestAC95ModeEquivalence:
    """AC9.5: Both output modes (direct and token stream) produce identical results."""

    def test_mode_equivalence_complex_graph(self):
        """Complex program produces same result (30) in both direct and token modes."""
        source = """
@system pe=3, sm=0
&a|pe0 <| const, 10
&b|pe0 <| const, 20
&sum|pe0 <| add
&out|pe1 <| pass
&ext|pe2 <| pass
&a|pe0 |> &sum|pe0:L
&b|pe0 |> &sum|pe0:R
&sum|pe0 |> &out|pe1:L
&out|pe1 |> &ext|pe2:L
"""
        # Both modes should produce the same result: 30 (10 + 20)
        direct_outputs = run_program_direct(source)
        token_outputs = run_program_tokens(source)

        # Get result from PE0 in both modes (where sum is computed and emitted)
        direct_result = [t.data for t in direct_outputs[0] if hasattr(t, 'data')]
        token_result = [t.data for t in token_outputs[0] if hasattr(t, 'data')]

        # Both should produce 30 (10 + 20)
        assert 30 in direct_result, f"Direct mode: expected 30 in PE0, got {direct_result}"
        assert 30 in token_result, f"Token mode: expected 30 in PE0, got {token_result}"


class TestAC105AutoPlacedE2E:
    """AC10.5: Auto-placed (unplaced) programs assemble and execute correctly."""

    def test_autoplaced_const_add_chain(self):
        """Unplaced const-add program auto-places and produces correct sum."""
        source = """
@system pe=3, sm=0
&c1 <| const, 3
&c2 <| const, 7
&result <| add
&output <| pass
&c1 |> &result:L
&c2 |> &result:R
&result |> &output:L
"""
        outputs = run_program_direct(source)
        # Find which PE has the output by checking all outputs
        all_values = []
        for pe_outputs in outputs.values():
            all_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])
        assert 10 in all_values, f"Expected sum 10 in any PE output, got {all_values}"

    def test_autoplaced_cross_pe_routing(self):
        """Unplaced cross-PE routing auto-places and produces 99 in both modes."""
        source = """
@system pe=3, sm=0
&source <| const, 99
&dest <| pass
&output <| pass
&source |> &dest:L
&dest |> &output:L
"""
        # Both modes should produce 99 somewhere
        direct_outputs = run_program_direct(source)
        token_outputs = run_program_tokens(source)

        # Check direct mode - source node should emit 99
        direct_values = []
        for pe_outputs in direct_outputs.values():
            direct_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])
        assert 99 in direct_values, f"Direct mode: expected 99, got {direct_values}"

        # Check token mode - source node should emit 99
        token_values = []
        for pe_outputs in token_outputs.values():
            token_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])
        assert 99 in token_values, f"Token mode: expected 99, got {token_values}"

    def test_autoplaced_vs_explicit_equivalence(self):
        """Auto-placed program produces same result (8) as explicitly-placed version."""
        explicit = """
@system pe=3, sm=0
&c1|pe0 <| const, 5
&c2|pe0 <| const, 3
&result|pe1 <| add
&output|pe2 <| pass
&c1|pe0 |> &result|pe1:L
&c2|pe0 |> &result|pe1:R
&result|pe1 |> &output|pe2:L
"""
        autoplaced = """
@system pe=3, sm=0
&c1 <| const, 5
&c2 <| const, 3
&result <| add
&output <| pass
&c1 |> &result:L
&c2 |> &result:R
&result |> &output:L
"""
        # Both should produce 8 (5 + 3) in both modes
        explicit_direct = run_program_direct(explicit)
        explicit_tokens = run_program_tokens(explicit)
        autoplaced_direct = run_program_direct(autoplaced)
        autoplaced_tokens = run_program_tokens(autoplaced)

        # Verify all modes produce 8
        for mode_name, outputs in [
            ("explicit_direct", explicit_direct),
            ("explicit_tokens", explicit_tokens),
            ("autoplaced_direct", autoplaced_direct),
            ("autoplaced_tokens", autoplaced_tokens),
        ]:
            all_values = []
            for pe_outputs in outputs.values():
                all_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])
            assert 8 in all_values, f"{mode_name}: expected 8, got {all_values}"


class TestMacroE2E:
    """End-to-end tests for macro expansion through full pipeline."""

    def test_const_pass_macro_direct(self):
        """Direct mode: macro expands and executes correctly through full pipeline.

        Defines a macro with const→pass pipeline, invokes it, and verifies the value
        flows through: lower → expand → resolve → place → allocate → codegen → emulator.
        Uses scoped references within the macro to connect the pipeline.
        """
        source = """
@system pe=1, sm=0

#const_pass |> {
    &const_node <| const, 42
    &const_node |> &sink:L
    &sink <| pass
}

#const_pass
"""
        outputs = run_program_direct(source)
        # Check all PE outputs for the constant value 42
        all_values = []
        for pe_outputs in outputs.values():
            all_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])
        assert 42 in all_values, \
            f"Expected value 42 in any PE output from macro expansion, got {all_values}"

    def test_const_pass_macro_tokens(self):
        """Token stream mode: macro expansion produces correct output."""
        source = """
@system pe=1, sm=0

#const_pass |> {
    &const_node <| const, 42
    &const_node |> &sink:L
    &sink <| pass
}

#const_pass
"""
        outputs = run_program_tokens(source)
        # Check all PE outputs for the constant value 42
        all_values = []
        for pe_outputs in outputs.values():
            all_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])
        assert 42 in all_values, \
            f"Expected value 42 in any PE output from macro expansion, got {all_values}"

    def test_macro_with_multiple_invocations(self):
        """Multiple invocations of the same macro each get unique scopes.

        Verifies that two invocations of the same macro create separate
        scope-qualified nodes (#macro_0, #macro_1) that execute independently.
        Each invocation produces output independently.
        """
        source = """
@system pe=1, sm=0

#const_pipeline |> {
    &const_node <| const, 15
    &const_node |> &out:L
    &out <| pass
}

#const_pipeline

#const_pipeline
"""
        outputs = run_program_direct(source)
        # Both macro invocations create const→pass pipelines that emit 15
        all_values = []
        for pe_outputs in outputs.values():
            all_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])
        # Should have at least two 15s (one from each macro invocation)
        count_15 = all_values.count(15)
        assert count_15 >= 2, \
            f"Expected at least two 15s in outputs (from two macro invocations), got {all_values}"


class TestAC48FunctionCalls:
    """AC4.8: Function call wiring works correctly end-to-end."""

    def test_function_call_basic_direct(self):
        """Direct mode: simple function call with argument and return.

        Defines a function that adds two inputs and returns the result,
        then calls it with two constants and verifies the output.
        """
        source = """
@system pe=1, sm=0

$adder |> {
    &a <| pass
    &b <| pass
    &sum <| add
    &a |> &sum:L
    &b |> &sum:R
    &sum |> @ret
}

&three <| const, 3
&seven <| const, 7
&result <| pass
$adder a=&three, b=&seven |> &result
"""
        outputs = run_program_direct(source)
        all_values = []
        for pe_outputs in outputs.values():
            all_values.extend([t.data for t in pe_outputs if hasattr(t, 'data')])

        assert 10 in all_values, \
            f"Expected result 10 from function call, got {all_values}"