"""Default preset workflow builders."""

from __future__ import annotations

from typing import Any

from agentbyte.llm import BaseChatCompletionClient
from agentbyte.workflow import Workflow, WorkflowConfig
from agentbyte.workflow.core.models import StepMetadata
from agentbyte.workflow.steps import AgentStep, FunctionStep
from agentbyte.workflow.steps.agentbyte_agent import (
    AgentbyteAgentInput,
    AgentbyteAgentOutput,
)

from .agents import get_researcher, get_reviewer, get_writer
from .clients import build_chat_client


def _create_linear_research_writer_reviewer_workflow(
    *,
    name: str,
    description: str | None,
    researcher,
    writer,
    reviewer,
) -> Workflow:
    """Create the default researcher -> writer -> reviewer workflow."""
    workflow = Workflow(WorkflowConfig(name=name, description=description))
    researcher_step = AgentStep(
        step_id="researcher",
        metadata=StepMetadata(name="researcher"),
        agent=researcher,
    )
    writer_step = AgentStep(
        step_id="writer",
        metadata=StepMetadata(name="writer"),
        agent=writer,
    )
    research_to_writer = FunctionStep(
        step_id="research_to_writer",
        metadata=StepMetadata(name="research_to_writer"),
        input_type=AgentbyteAgentOutput,
        output_type=AgentbyteAgentInput,
        func=lambda output, _context: {
            "task": output.response or "",
            "additional_context": {"messages": output.messages, "usage": output.usage},
        },
    )
    reviewer_step = AgentStep(
        step_id="reviewer",
        metadata=StepMetadata(name="reviewer"),
        agent=reviewer,
    )
    writer_to_reviewer = FunctionStep(
        step_id="writer_to_reviewer",
        metadata=StepMetadata(name="writer_to_reviewer"),
        input_type=AgentbyteAgentOutput,
        output_type=AgentbyteAgentInput,
        func=lambda output, _context: {
            "task": output.response or "",
            "additional_context": {"messages": output.messages, "usage": output.usage},
        },
    )
    workflow.chain(
        researcher_step,
        research_to_writer,
        writer_step,
        writer_to_reviewer,
        reviewer_step,
    )
    return workflow


def get_workflow(
    model_client: BaseChatCompletionClient | None = None,
    *,
    provider: str | None = None,
    model: str | None = None,
    config: dict[str, Any] | None = None,
    name: str = "research_writer_reviewer",
    description: str = "Default linear researcher -> writer -> reviewer workflow preset.",
    **kwargs: Any,
) -> Workflow:
    """Build the default workflow preset."""
    del kwargs
    client = model_client or build_chat_client(
        provider=provider,
        model=model,
        config=config,
    )
    workflow = _create_linear_research_writer_reviewer_workflow(
        researcher=get_researcher(client),
        writer=get_writer(client),
        reviewer=get_reviewer(client),
        name=name,
        description=description,
    )
    workflow.config.metadata.setdefault(
        "example_tasks",
        ["Write a brief article about solar energy benefits."],
    )
    return workflow


__all__ = ["get_workflow"]
