"""Structured output for CLIN."""

from typing import Any, Dict, List

from pydantic import BaseModel, Field

from agential.agents.base.output import BaseAgentOutput
from agential.core.llm import Response


class CLINReActStepOutput(BaseModel):
    """CLIN ReAct Step Pydantic output class.

    Attributes:
        thought (str): The thought process of the agent.
        action_type (str): The type of action performed by the agent.
        query (str): The query requested by the agent.
        observation (str): The observation made by the agent.
        answer (str): The answer generated by the agent.
        external_tool_info (Dict[str, Any]): The external tool outputs.
        is_correct (bool): Indicates if the action was correct.
        thought_response (Response): Thought response.
        action_response (Response): Action response.
    """

    thought: str = Field(..., description="The thought process of the agent.")
    action_type: str = Field(
        ..., description="The type of action performed by the agent."
    )
    query: str = Field(..., description="The query requested by the agent.")
    observation: str = Field(..., description="The observation made by the agent.")
    answer: str = Field(..., description="The answer generated by the agent.")
    external_tool_info: Dict[str, Any] = Field(
        ..., description="The external tool outputs."
    )
    is_correct: bool = Field(..., description="Indicates if the action was correct.")
    thought_response: Response = Field(
        ..., description="Prompt response for the thought."
    )
    action_response: Response = Field(
        ..., description="Prompt response for the thought."
    )


class CLINStepOutput(BaseModel):
    """CLIN Step Pydantic output class.

    Attributes:
        steps (List[CLINReActStepOutput]): The steps of the agent.
        summaries (str): The summaries of the agent.
        summaries_response (Response): Prompt response for the summaries.
        meta_summaries (str): The meta summaries of the agent.
        previous_trials (str): The previous trials of the agent.
    """

    steps: List[CLINReActStepOutput] = Field(..., description="The steps of the agent.")
    summaries: str = Field(..., description="The summaries of the agent.")
    summaries_response: Response = Field(
        ..., description="Prompt response for the summaries."
    )
    meta_summaries: str = Field(..., description="The meta summaries of the agent.")
    previous_trials: str = Field(..., description="The previous trials of the agent.")


class CLINOutput(BaseAgentOutput):
    """Structured output for CLIN.

    Attributes:
        additional_info (List[CLINStepOutput]): The answer to the question.
    """

    additional_info: List[CLINStepOutput] = Field(
        description="The answer to the question."
    )
