"""SyntheticDataGenerator component for creating synthetic data using LLM-based generation."""

from __future__ import annotations

from typing import ClassVar

from lfx.components.agentics.constants import ERROR_AGENTICS_NOT_INSTALLED
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 DataFrameInput, IntInput, MessageTextInput, Output
from lfx.schema.dataframe import DataFrame


class AgenerateComponent(BaseAgenticComponent):
    """Generate synthetic data using either example data or a defined schema.

    This component creates realistic synthetic data by either:
    1. Learning from an input DataFrame and generating similar rows, or
    2. Following a user-defined schema to create data from scratch.

    """

    code_class_base_inheritance: ClassVar[str] = "Component"
    display_name = "aGenerate"
    description = (
        "Generate mock data for user defined schema. If a dataframe is provided, "
        "the component will generate similar rows."
    )
    documentation: str = "https://docs.langflow.org/bundles-agentics"
    icon = "Agentics"

    inputs = [
        *get_model_provider_inputs(),
        get_generated_fields_input(
            name="schema",
            display_name="Schema",
            info=(
                "Define the structure of data to generate. Specify column names, "
                "descriptions, and types. Used only when input DataFrame is not provided."
            ),
            required=False,
        ),
        DataFrameInput(
            name="source",
            display_name="Input Table",
            info=(
                "Provide example DataFrame to learn from and generate similar data. "
                "Only the first 50 rows will be used as examples."
            ),
            required=False,
            advanced=False,
            value=None,
        ),
        MessageTextInput(
            name="instructions",
            display_name="Instructions",
            info="Optional natural language instructions to guide the synthetic data generation process.",
            value="",
            required=False,
            advanced=True,
        ),
        IntInput(
            name="batch_size",
            display_name="Number of Rows to Generate",
            info=(
                "Number of new synthetic rows to generate. When an Input Table is provided, "
                "the generated rows are appended to the original data."
            ),
            value=10,
            advanced=False,
        ),
    ]

    outputs = [
        Output(
            name="states",
            display_name="Output Table",
            info="Synthetic DataFrame generated by the LLM based on the schema or example data.",
            method="aGenerate",
            tool_mode=True,
        ),
    ]

    async def aGenerate(self) -> DataFrame:  # noqa: N802
        """Generate synthetic data using LLM-based generation.

        Returns:
            DataFrame containing the generated synthetic data.
        """
        try:
            from agentics import AG
            from agentics.core.atype import create_pydantic_model
            from agentics.core.transducible_functions import generate_prototypical_instances
        except ImportError as e:
            raise ImportError(ERROR_AGENTICS_NOT_INSTALLED) from e

        llm = prepare_llm_from_component(self)

        if self.source:
            source = AG.from_dataframe(DataFrame(self.source))
            atype = source.atype
            instructions = (
                str(self.instructions) if self.instructions else "Generate similar data based on the examples provided."
            )
            instructions += "\nHere are examples to take inspiration from:\n" + str(source.states[:50])
        elif self.schema != []:
            schema_fields = build_schema_fields(self.schema)
            atype = create_pydantic_model(schema_fields, name="GeneratedData")
            instructions = (
                str(self.instructions)
                if self.instructions
                else "Generate realistic synthetic data following the provided schema."
            )
        else:
            msg = "Synthetic data generation requires either a sample DataFrame or schema definition (but not both)."
            raise ValueError(msg)

        output_states = await generate_prototypical_instances(
            atype,
            n_instances=self.batch_size,
            llm=llm,
            instructions=instructions,
        )
        # Ensure output_states is a list, not None
        if output_states is None:
            output_states = []

        if self.source:
            output_states = source.states + output_states

        output = AG(atype=atype, states=output_states)

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