"""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, WorkflowContext
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 _writer_to_reviewer_fn(
    output: AgentbyteAgentOutput, context: WorkflowContext
) -> dict:
    """Transform writer output into reviewer input; store draft for later retrieval."""
    context.set("writer_draft", output.response or "")
    return {
        "task": output.response or "",
        "additional_context": {"messages": output.messages, "usage": output.usage},
    }


def _reviewer_to_writer_fn(
    output: AgentbyteAgentOutput, context: WorkflowContext
) -> dict:
    """Transform reviewer output into writer_final input.

    Passes the original draft plus reviewer feedback so the writer can
    produce an improved final version. If the reviewer approved, the writer
    is asked to confirm and present the polished version.
    """
    draft = context.get("writer_draft") or ""
    feedback = (output.response or "").strip()
    approved = feedback.upper() == "APPROVED"
    if approved:
        task = f"Your draft was approved by the reviewer. Output the final polished version:\n\n{draft}"
    else:
        task = (
            f"Revise your draft based on the reviewer's feedback.\n\n"
            f"Your draft:\n{draft}\n\nReviewer feedback:\n{feedback}"
        )
    return {
        "task": task,
        "additional_context": {"messages": output.messages, "usage": output.usage},
    }


def _create_linear_research_writer_reviewer_workflow(
    *,
    name: str,
    description: str | None,
    researcher,
    writer,
    reviewer,
) -> Workflow:
    """Create the default researcher -> writer -> reviewer -> writer_final 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=_writer_to_reviewer_fn,
    )
    reviewer_to_writer = FunctionStep(
        step_id="reviewer_to_writer",
        metadata=StepMetadata(name="reviewer_to_writer"),
        input_type=AgentbyteAgentOutput,
        output_type=AgentbyteAgentInput,
        func=_reviewer_to_writer_fn,
    )
    writer_final_step = AgentStep(
        step_id="writer_final",
        metadata=StepMetadata(name="writer_final"),
        agent=writer,
    )
    workflow.chain(
        researcher_step,
        research_to_writer,
        writer_step,
        writer_to_reviewer,
        reviewer_step,
        reviewer_to_writer,
        writer_final_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"]
