"""Wrap an Agentbyte agent as a workflow step."""

from __future__ import annotations

from typing import Any

from pydantic import BaseModel, Field

from agentbyte.agents import AgentResponse, BaseAgent
from agentbyte.context import AgentContext
from agentbyte.workflow.core.models import StepMetadata, WorkflowContext
from agentbyte.workflow.steps.step import BaseStep


class AgentbyteAgentInput(BaseModel):
    task: str = Field(
        description=(
            "Primary instruction for the agent step. Example: "
            "'Write a short article about solar energy benefits.'"
        )
    )
    additional_context: dict[str, Any] | None = Field(
        default=None,
        description=(
            "Optional JSON object with extra guidance such as audience, tone, "
            "constraints, or supporting notes."
        ),
    )
    agent_context: AgentContext | None = Field(
        default=None,
        description=(
            "Optional AgentContext JSON used to resume a prior agent conversation "
            "or seed the step with existing message history."
        ),
    )


class AgentbyteAgentOutput(BaseModel):
    response: str | None = None
    messages: list[dict[str, Any]] = Field(default_factory=list)
    usage: dict[str, Any] | None = None
    metadata: dict[str, Any] | None = None


class AgentStep(
    BaseStep[AgentbyteAgentInput, AgentbyteAgentOutput],
):
    """Executes an Agentbyte agent inside a workflow step."""

    def __init__(self, step_id: str, metadata: StepMetadata, agent: BaseAgent) -> None:
        super().__init__(step_id, metadata, AgentbyteAgentInput, AgentbyteAgentOutput)
        self.agent = agent

    async def execute(self, input_data: AgentbyteAgentInput, context: WorkflowContext) -> AgentbyteAgentOutput:
        # The workflow Context carries step state (a dict). AgentContext is the
        # agent's conversation history and is opt-in via input_data.agent_context.
        response: AgentResponse = await self.agent.run(
            task=input_data.task,
            context=input_data.agent_context,
        )

        last_assistant = response.final_content
        messages = [msg.model_dump() for msg in response.messages]
        usage = response.usage.model_dump()

        if input_data.additional_context:
            context.set(f"{self.step_id}_context", input_data.additional_context)

        return AgentbyteAgentOutput(
            response=last_assistant,
            messages=messages,
            usage=usage,
            metadata=input_data.additional_context,
        )


__all__ = [
    "AgentStep",
    "AgentbyteAgentInput",
    "AgentbyteAgentOutput",
]
