from __future__ import annotations

import importlib.util
from pathlib import Path
from types import SimpleNamespace

import pytest

from agentbyte.agents.types import AgentResponse
from agentbyte.context import AgentContext
from agentbyte.messages import AssistantMessage
from agentbyte.orchestration import AIOrchestrator, PlanBasedOrchestrator, PlanStep, RoundRobinOrchestrator
from agentbyte.presets.orchestration import (
    get_ai_orchestrator,
    get_orchestrator,
    get_plan_based_orchestrator,
    get_round_robin_orchestrator,
)
from agentbyte.termination import CompositeTermination, MaxMessageTermination


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


def _entity_name(entity) -> str | None:
    name = getattr(entity, "name", None)
    if name is not None:
        return name
    config = getattr(entity, "config", None)
    return getattr(config, "name", None)


def test_presets_webui_build_entities_includes_plan_based_team(monkeypatch):
    module_path = (
        Path(__file__).resolve().parents[2] / "examples" / "webui" / "presets_webui.py"
    )
    spec = importlib.util.spec_from_file_location("presets_webui", module_path)
    assert spec and spec.loader
    presets_webui = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(presets_webui)

    monkeypatch.setattr(presets_webui, "build_chat_client", _dummy_client)

    entities = presets_webui.build_entities()
    entities_by_name = {_entity_name(entity): entity for entity in entities}

    assert list(entities_by_name) == [
        "query_rewriter",
        "researcher",
        "writer",
        "reviewer",
        "round_robin_team",
        "ai_team",
        "plan_based_team",
        "research_writer_reviewer_workflow",
    ]
    assert isinstance(entities_by_name["plan_based_team"], PlanBasedOrchestrator)
    assert "approval gate" in entities_by_name["plan_based_team"].description.lower()


def test_get_round_robin_orchestrator_uses_researcher_and_writer():
    orchestrator = get_round_robin_orchestrator(model_client=_dummy_client())

    assert isinstance(orchestrator, RoundRobinOrchestrator)
    assert [agent.name for agent in orchestrator.agents] == ["researcher", "writer"]
    assert isinstance(orchestrator.termination, CompositeTermination)
    assert orchestrator.history_policy.default.window == "full"
    assert orchestrator.history_policy.default.format == "text"
    assert "TERMINATE" in orchestrator.agents[1].instructions


def test_get_ai_orchestrator_uses_default_roster():
    orchestrator = get_ai_orchestrator(model_client=_dummy_client())

    assert isinstance(orchestrator, AIOrchestrator)
    assert [agent.name for agent in orchestrator.agents] == [
        "researcher",
        "writer",
        "reviewer",
    ]
    assert isinstance(orchestrator.termination, CompositeTermination)
    assert orchestrator.history_policy.default.window == "full"
    assert orchestrator.history_policy.default.format == "text"
    assert orchestrator.selector_policy.window == "recent"
    assert orchestrator.selector_policy.message_limit == 12
    assert orchestrator.selector_policy.format == "summary"
    assert "TERMINATE" not in orchestrator.agents[1].instructions
    assert "APPROVED" in orchestrator.agents[2].instructions


def test_get_plan_based_orchestrator_uses_default_roster():
    orchestrator = get_plan_based_orchestrator(model_client=_dummy_client())

    assert isinstance(orchestrator, PlanBasedOrchestrator)
    assert [agent.name for agent in orchestrator.agents] == [
        "researcher",
        "writer",
        "reviewer",
    ]
    assert isinstance(orchestrator.termination, MaxMessageTermination)
    assert orchestrator.history_policy.default.window == "recent"
    assert orchestrator.history_policy.default.message_limit == 5
    assert orchestrator.history_policy.default.format == "messages"
    assert orchestrator.review_gate_policy.enabled is True
    assert orchestrator.review_gate_policy.reviewer_agent_name == "reviewer"
    assert orchestrator.review_gate_policy.revision_agent_name == "writer"
    assert "TERMINATE" not in orchestrator.agents[1].instructions
    assert "APPROVED" in orchestrator.agents[2].instructions


def test_get_orchestrator_dispatches_to_requested_pattern():
    round_robin = get_orchestrator("round_robin", model_client=_dummy_client())
    ai = get_orchestrator("ai", model_client=_dummy_client())
    plan_based = get_orchestrator("plan_based", model_client=_dummy_client())

    assert isinstance(round_robin, RoundRobinOrchestrator)
    assert isinstance(ai, AIOrchestrator)
    assert isinstance(plan_based, PlanBasedOrchestrator)


def test_get_orchestrator_rejects_unknown_pattern():
    with pytest.raises(ValueError, match="Unsupported orchestrator pattern"):
        get_orchestrator("unknown", model_client=_dummy_client())  # type: ignore[arg-type]


@pytest.mark.asyncio
async def test_plan_based_preset_requires_explicit_reviewer_approval():
    orchestrator = get_plan_based_orchestrator(model_client=_dummy_client())
    review_step = PlanStep(
        task="Review the final draft",
        agent_name="reviewer",
        reasoning="Quality gate",
    )
    context = AgentContext(
        messages=[
            AssistantMessage(
                content="The article is clear but needs one more improvement.",
                source="reviewer",
            )
        ]
    )
    result = AgentResponse(
        source="reviewer",
        context=context,
        finish_reason="stop",
    )

    evaluation = await orchestrator.evaluate_step_progress(review_step, result)

    assert evaluation.step_completed is False
    assert "did not explicitly approve" in evaluation.failure_reason


@pytest.mark.asyncio
async def test_plan_based_preset_inserts_writer_revision_after_reviewer_feedback():
    orchestrator = get_plan_based_orchestrator(model_client=_dummy_client())
    orchestrator.execution_plan = type(
        "ExecutionPlanLike",
        (),
        {
            "steps": [
                PlanStep(
                    task="Review the final draft",
                    agent_name="reviewer",
                    reasoning="Quality gate",
                )
            ]
        },
    )()
    orchestrator.current_step_index = 0
    orchestrator._last_context_size = 0
    orchestrator._last_context_was_text = False

    context = AgentContext(
        messages=[
            AssistantMessage(
                content="Needs a clearer evidence-backed revision before approval.",
                source="reviewer",
            )
        ]
    )
    result = AgentResponse(
        source="reviewer",
        context=context,
        finish_reason="stop",
    )

    await orchestrator.update_shared_state(result)

    assert orchestrator.current_step_index == 1
    assert orchestrator.execution_plan.steps[1].agent_name == "writer"
    assert orchestrator.execution_plan.steps[2].agent_name == "reviewer"


def test_plan_based_preset_final_result_prefers_latest_writer_output():
    orchestrator = get_plan_based_orchestrator(model_client=_dummy_client())
    orchestrator.shared_messages = [
        AssistantMessage(content="Initial research notes", source="researcher"),
        AssistantMessage(content="Final approved article draft", source="writer"),
        AssistantMessage(content="APPROVED", source="reviewer"),
    ]

    assert orchestrator._generate_final_result() == "Final approved article draft"
