from __future__ import annotations

from types import SimpleNamespace

from agentbyte.presets.workflow import get_workflow
from agentbyte.presets.workflow import _reviewer_to_writer_fn
from agentbyte.presets.workflow import _writer_to_reviewer_fn
from agentbyte.workflow.core.models import WorkflowContext
from agentbyte.workflow.steps.agentbyte_agent import AgentbyteAgentOutput


def _dummy_client():
    return SimpleNamespace(model="test-model")


def test_get_workflow_returns_linear_research_writer_reviewer_workflow():
    workflow = get_workflow(model_client=_dummy_client())

    validation = workflow.validate_workflow()

    assert validation.is_valid
    assert workflow.start_step_id == "researcher"
    assert workflow.end_step_ids == ["writer_final"]
    assert [(edge.from_step, edge.to_step) for edge in workflow.edges] == [
        ("researcher", "research_to_writer"),
        ("research_to_writer", "writer"),
        ("writer", "writer_to_reviewer"),
        ("writer_to_reviewer", "reviewer"),
        ("reviewer", "reviewer_to_writer"),
        ("reviewer_to_writer", "writer_final"),
    ]
    assert [
        step_id for step_id in workflow.steps
        if step_id in {"researcher", "writer", "reviewer", "writer_final"}
    ] == [
        "researcher",
        "writer",
        "reviewer",
        "writer_final",
    ]


def test_writer_to_reviewer_stores_draft_in_workflow_state():
    context = WorkflowContext(state={})
    output = AgentbyteAgentOutput(
        response="Draft article",
        messages=[{"role": "assistant", "content": "Draft article"}],
        usage={"total_tokens": 42},
    )

    transformed = _writer_to_reviewer_fn(output, context)

    assert context.get("writer_draft") == "Draft article"
    assert transformed["task"] == "Draft article"
    assert transformed["additional_context"] == {
        "messages": [{"role": "assistant", "content": "Draft article"}],
        "usage": {"total_tokens": 42},
    }


def test_reviewer_to_writer_approved_uses_stored_draft_for_final_polish():
    context = WorkflowContext(state={"writer_draft": "Draft v1"})
    output = AgentbyteAgentOutput(
        response="approved",
        messages=[{"role": "assistant", "content": "approved"}],
        usage={"total_tokens": 10},
    )

    transformed = _reviewer_to_writer_fn(output, context)

    assert "Output the final polished version" in transformed["task"]
    assert "Draft v1" in transformed["task"]
    assert transformed["additional_context"] == {
        "messages": [{"role": "assistant", "content": "approved"}],
        "usage": {"total_tokens": 10},
    }


def test_reviewer_to_writer_non_approved_requests_single_revision():
    context = WorkflowContext(state={"writer_draft": "Draft v1"})
    output = AgentbyteAgentOutput(
        response="Needs stronger structure and citations.",
        messages=[
            {
                "role": "assistant",
                "content": "Needs stronger structure and citations.",
            }
        ],
        usage={"total_tokens": 12},
    )

    transformed = _reviewer_to_writer_fn(output, context)

    assert "Revise your draft based on the reviewer's feedback." in transformed["task"]
    assert "Draft v1" in transformed["task"]
    assert "Needs stronger structure and citations." in transformed["task"]
