"""SemanticAggregator component for aggregating and summarizing input data using LLM-based semantic analysis."""

from __future__ import annotations

from typing import ClassVar

from pydantic import create_model

from lfx.components.agentics.constants import (
    ERROR_AGENTICS_NOT_INSTALLED,
    ERROR_INPUT_SCHEMA_REQUIRED,
    TRANSDUCTION_AREDUCE,
)
from lfx.components.agentics.helpers import (
    build_schema_fields,
    prepare_llm_from_component,
)
from lfx.components.agentics.inputs import (
    get_generated_fields_input,
    get_model_provider_inputs,
)
from lfx.components.agentics.inputs.base_component import BaseAgenticComponent
from lfx.io import (
    BoolInput,
    DataFrameInput,
    MessageTextInput,
    Output,
)
from lfx.schema.dataframe import DataFrame


class AreduceComponent(BaseAgenticComponent):
    """Aggregate or summarize entire input data using natural language instructions and a defined output schema.

    This component processes all rows of input data collectively to produce aggregated results,
    such as summaries, statistics, or consolidated information based on LLM analysis.
    """

    code_class_base_inheritance: ClassVar[str] = "Component"
    display_name = "aReduce"
    description = (
        "Analyze the entire input dataframe at once and generate a new dataframe "
        "following the instruction and the required schema"
    )
    documentation: str = "https://docs.langflow.org/bundles-agentics"
    icon = "Agentics"

    inputs = [
        *get_model_provider_inputs(),
        DataFrameInput(
            name="source",
            display_name="Input Table",
            info="Input DataFrame to aggregate. The schema is automatically inferred from column names and types.",
            required=True,
        ),
        get_generated_fields_input(),
        BoolInput(
            name="return_multiple_instances",
            display_name="As List",
            info="If True, generate a list of instances of the provided schema.",
            advanced=False,
            value=False,
        ),
        MessageTextInput(
            name="instructions",
            display_name="Instructions",
            info="Natural language instructions describing how to aggregate the input data into the output schema.",
            advanced=False,
            value="",
            required=False,
        ),
    ]

    outputs = [
        Output(
            name="states",
            method="aReduce",
            display_name="Output Table",
            info="Aggregated DataFrame generated by the LLM following the specified output schema.",
            tool_mode=True,
        ),
    ]

    async def aReduce(self) -> DataFrame:  # noqa: N802
        """Aggregate input data using LLM-based semantic analysis.

        Returns:
            DataFrame containing the aggregated results following the output schema.
        """
        try:
            from agentics import AG
            from agentics.core.atype import create_pydantic_model
        except ImportError as e:
            raise ImportError(ERROR_AGENTICS_NOT_INSTALLED) from e

        llm = prepare_llm_from_component(self)

        if self.source and self.schema != []:
            source = AG.from_dataframe(DataFrame(self.source))

            schema_fields = build_schema_fields(self.schema)
            atype = create_pydantic_model(schema_fields, name="Target")
            if self.return_multiple_instances:
                final_atype = create_model("ListOfTarget", items=(list[atype], ...))  # type: ignore[valid-type]
            else:
                final_atype = atype

            target = AG(
                atype=final_atype,
                transduction_type=TRANSDUCTION_AREDUCE,
                instructions=self.instructions
                if not self.return_multiple_instances
                else "\nGenerate a list of instances of the target type following those instructions : ."
                + self.instructions,
                llm=llm,
                areduce_batch_size=100,
            )

            output = await (target << source)
            if self.return_multiple_instances:
                appended_states = [item_state for state in output for item_state in state.items]
                output = AG(atype=atype, states=appended_states)

            return DataFrame(output.to_dataframe().to_dict(orient="records"))
        raise ValueError(ERROR_INPUT_SCHEMA_REQUIRED)
