"""Test for aReduce bug: 'As List' toggle loses items from chunks beyond the first."""

from __future__ import annotations

import pytest

try:
    import agentics  # noqa: F401
except ImportError:
    pytest.skip("agentics-py not installed", allow_module_level=True)

from agentics import AG
from agentics.core.atype import create_pydantic_model
from lfx.components.agentics.helpers.schema_builder import build_schema_fields
from pydantic import create_model


def _extract_list_items(output: AG, atype: type) -> AG:
    """Replicate the aReduce 'As List' post-processing logic from AreduceComponent.

    This function mirrors the exact code path in areduce_component.py
    for when return_multiple_instances=True.
    """
    # Current (buggy) code: output = AG(atype=atype, states=output[0].items)
    # Fixed code should iterate all states like aMap does
    from lfx.components.agentics.areduce_component import AreduceComponent  # noqa: F401

    # We replicate what the component does — after the fix, this should iterate all states
    appended_states = [item_state for state in output for item_state in state.items]
    return AG(atype=atype, states=appended_states)


@pytest.mark.unit
class TestAreduceAsList:
    """Tests for the aReduce 'As List' post-processing logic."""

    def test_should_collect_items_from_all_output_states_not_just_first(self):
        """Verify aReduce collects items from ALL output states, not just the first.

        When aReduce processes data in multiple batches (areduce_batch_size), each batch
        produces a separate output state with a ListOfTarget wrapper. The post-processing
        must collect items from ALL states, not just the first one.

        Bug: output[0].items only gets items from the first batch, losing data.
        Fix: iterate all states like aMap does.
        """
        # Arrange
        output_fields = build_schema_fields(
            [
                {"name": "total_orders", "description": "Total orders", "type": "int", "multiple": False},
                {"name": "top_product", "description": "Top product", "type": "str", "multiple": False},
            ]
        )
        atype = create_pydantic_model(output_fields, name="Target")
        final_atype = create_model("ListOfTarget", items=(list[atype], ...))

        # Simulate multiple output states (as happens with multiple areduce batches)
        state1 = final_atype(
            items=[
                atype(total_orders=100, top_product="Widget A"),
                atype(total_orders=50, top_product="Widget B"),
            ]
        )
        state2 = final_atype(
            items=[
                atype(total_orders=25, top_product="Widget C"),
                atype(total_orders=10, top_product="Widget D"),
            ]
        )
        output = AG(atype=final_atype, states=[state1, state2])

        # Act
        result = _extract_list_items(output, atype)
        result_df = result.to_dataframe()

        # Assert — ALL 4 items from both states must be present
        assert len(result_df) == 4
        assert set(result_df["top_product"].tolist()) == {"Widget A", "Widget B", "Widget C", "Widget D"}

    def test_should_work_with_single_output_state(self):
        """Single output state (common case: <100 rows) should also work correctly."""
        # Arrange
        output_fields = build_schema_fields(
            [
                {"name": "category", "description": "Category", "type": "str", "multiple": False},
                {"name": "count", "description": "Count", "type": "int", "multiple": False},
            ]
        )
        atype = create_pydantic_model(output_fields, name="Target")
        final_atype = create_model("ListOfTarget", items=(list[atype], ...))

        state = final_atype(
            items=[
                atype(category="Electronics", count=42),
                atype(category="Books", count=15),
                atype(category="Clothing", count=8),
            ]
        )
        output = AG(atype=final_atype, states=[state])

        # Act
        result = _extract_list_items(output, atype)
        result_df = result.to_dataframe()

        # Assert
        assert len(result_df) == 3
        assert set(result_df["category"].tolist()) == {"Electronics", "Books", "Clothing"}