from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any, Optional, Union

import pytest

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, Usage
from agentbyte.context import AgentContext
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import AssistantMessage, Message, UserMessage
from agentbyte.orchestration import (
    AgentHistoryPolicy,
    ExecutionPlan,
    FinalResultPolicy,
    HistoryContextPolicy,
    PlanBasedOrchestrator,
    PlanStep,
    ReviewGatePolicy,
    StepProgressEvaluation,
)
from agentbyte.termination import MaxMessageTermination
from agentbyte.types import OrchestrationEvent, OrchestrationResponse


class FakePlanClientConfig(BaseChatCompletionClientConfig):
    pass


class FakePlanClient(BaseChatCompletionClient):
    def __init__(
        self,
        responses: list[ChatCompletionResult],
        fail_on_calls: Optional[set[int]] = None,
    ) -> None:
        super().__init__(model="plan-fake", client=object())
        self._responses = responses
        self._index = 0
        self._fail_on_calls = fail_on_calls or set()

    async def create(
        self,
        messages: list[Message],
        tools: Optional[list[dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        del messages, tools, output_format, kwargs
        self._index += 1
        if self._index in self._fail_on_calls:
            raise RuntimeError("synthetic client failure")
        return self._responses[self._index - 1]

    async def create_stream(
        self,
        messages: list[Message],
        tools: Optional[list[dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        del messages, tools, output_format, kwargs
        yield ChatCompletionChunk(content="", model="plan-fake", is_complete=True)

    def _to_config(self) -> BaseChatCompletionClientConfig:
        return FakePlanClientConfig(model=self.model, config=self.config)

    @classmethod
    def _from_config(cls, config: BaseChatCompletionClientConfig) -> "FakePlanClient":
        del config
        return cls(responses=[])


class MockAgent(BaseAgent):
    def __init__(self, name: str, response_texts: list[str]) -> None:
        self.name = name
        self.description = f"{name} description"
        self.tools = []
        self._response_texts = response_texts
        self._index = 0
        self.received_tasks: list[Union[UserMessage, list[Message]]] = []

    async def run(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[Any] = None,
    ) -> AgentResponse:
        del task, context, cancellation_token
        raise NotImplementedError("Use run_stream")

    async def run_stream(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[Any] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[Union[Message, Any, AgentResponse], None]:
        del cancellation_token, verbose, stream_tokens
        if task is not None:
            self.received_tasks.append(task)

        ctx = context or AgentContext()
        if isinstance(task, UserMessage):
            yield task
            if not ctx.messages:
                ctx.messages.append(task)
        elif isinstance(task, list):
            for message in task:
                yield message
                ctx.messages.append(message)

        response_text = self._response_texts[self._index % len(self._response_texts)]
        self._index += 1
        message = AssistantMessage(content=response_text, source=self.name)
        yield message
        ctx.messages.append(message)
        yield AgentResponse(
            source=self.name,
            context=ctx,
            usage=Usage(tokens_input=5, tokens_output=5, llm_calls=1),
            finish_reason="stop",
        )


def make_result(
    structured_output: Any,
    tokens_input: int = 3,
    tokens_output: int = 1,
) -> ChatCompletionResult:
    return ChatCompletionResult(
        message=AssistantMessage(content="structured", source="model"),
        usage=Usage(
            tokens_input=tokens_input,
            tokens_output=tokens_output,
            llm_calls=1,
        ),
        model="plan-fake",
        structured_output=structured_output,
    )


@pytest.mark.asyncio
class TestPlanBasedOrchestrator:
    async def test_plan_creation_executes_steps_in_order(self) -> None:
        researcher = MockAgent("researcher", ["Found key facts about solar energy."])
        writer = MockAgent("writer", ["Created a clear draft about solar benefits."])
        reviewer = MockAgent("reviewer", ["Reviewed the draft and confirmed accuracy."])
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Gather core facts",
                                agent_name="researcher",
                                reasoning="Best for fact gathering",
                            ),
                            PlanStep(
                                task="Write the note",
                                agent_name="writer",
                                reasoning="Best for polished prose",
                            ),
                            PlanStep(
                                task="Review the note",
                                agent_name="reviewer",
                                reasoning="Best for final verification",
                            ),
                        ]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.95,
                        suggested_improvements=[],
                    )
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher, writer, reviewer],
            termination=MaxMessageTermination(20),
            model_client=client,
        )

        response = await orchestrator.run(
            UserMessage(content="Create a solar note", source="user")
        )

        assert isinstance(response, OrchestrationResponse)
        assert [message.source for message in response.messages] == [
            "user",
            "researcher",
            "writer",
            "reviewer",
        ]
        assert response.stop_message.content == "Execution plan completed"
        assert response.pattern_metadata["steps_completed"] == 3
        assert response.pattern_metadata["steps_failed"] == 0
        assert response.pattern_metadata["total_steps"] == 3

    async def test_default_recent_messages_history_is_preserved(self) -> None:
        researcher = MockAgent("researcher", ["Found facts."])
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Research",
                                agent_name="researcher",
                                reasoning="Best fit",
                            )
                        ]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    )
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
        )
        orchestrator.shared_messages = [
            UserMessage(content="1", source="user"),
            AssistantMessage(content="2", source="researcher"),
            AssistantMessage(content="3", source="researcher"),
            AssistantMessage(content="4", source="researcher"),
            AssistantMessage(content="5", source="researcher"),
            AssistantMessage(content="6", source="researcher"),
        ]
        orchestrator.execution_plan = ExecutionPlan(
            steps=[
                PlanStep(
                    task="Research",
                    agent_name="researcher",
                    reasoning="Best fit",
                )
            ]
        )

        context = await orchestrator.prepare_context_for_agent(researcher)

        assert isinstance(context, list)
        assert [message.content for message in context[:-1]] == ["2", "3", "4", "5", "6"]
        assert isinstance(context[-1], UserMessage)
        assert "PLAN STEP 1" in context[-1].content

    async def test_text_based_history_mode_appends_step_instruction(self) -> None:
        researcher = MockAgent("researcher", ["Found facts."])
        client = FakePlanClient(responses=[])
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
            history_policy=HistoryContextPolicy(
                default=AgentHistoryPolicy(window="full", format="text")
            ),
        )
        orchestrator.shared_messages = [UserMessage(content="task", source="user")]
        orchestrator.execution_plan = ExecutionPlan(
            steps=[
                PlanStep(
                    task="Research the topic",
                    agent_name="researcher",
                    reasoning="Best fit",
                )
            ]
        )

        context = await orchestrator.prepare_context_for_agent(researcher)

        assert isinstance(context, UserMessage)
        assert "It is now your turn." in context.content
        assert "PLAN STEP 1: Research the topic" in context.content

    async def test_planning_failure_uses_fallback_plan(self) -> None:
        researcher = MockAgent("researcher", ["Completed the task with fallback planning."])
        writer = MockAgent("writer", ["unused"])
        client = FakePlanClient(
            responses=[
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.8,
                        suggested_improvements=[],
                    )
                )
            ],
            fail_on_calls={1},
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher, writer],
            termination=MaxMessageTermination(10),
            model_client=client,
        )

        response = await orchestrator.run("Finish the work")

        assert response.messages[1].source == "researcher"
        assert response.pattern_metadata["execution_plan"]["steps"][0]["agent_name"] == "researcher"

    async def test_agent_name_resolution_is_case_insensitive(self) -> None:
        researcher = MockAgent("researcher", ["Gathered useful facts."])
        writer = MockAgent("writer", ["unused"])
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Gather facts",
                                agent_name="RESEARCHER",
                                reasoning="Case-insensitive match",
                            )
                        ]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    )
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher, writer],
            termination=MaxMessageTermination(10),
            model_client=client,
        )

        response = await orchestrator.run("Gather facts")

        assert response.messages[1].source == "researcher"

    async def test_failed_step_retries_with_feedback_and_then_succeeds(self) -> None:
        researcher = MockAgent(
            "researcher",
            ["Initial draft without enough detail.", "Completed a stronger second attempt."],
        )
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Produce a concise research summary",
                                agent_name="researcher",
                                reasoning="Research specialist",
                            )
                        ]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=False,
                        failure_reason="Missing key facts",
                        confidence_score=0.8,
                        suggested_improvements=[
                            "Add two concrete facts",
                            "Clarify the main takeaway",
                        ],
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.95,
                        suggested_improvements=[],
                    )
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
            max_step_retries=2,
        )

        response = await orchestrator.run("Summarize the topic")

        assert [message.source for message in response.messages] == ["user", "researcher", "researcher"]
        assert response.pattern_metadata["total_retries"] == 1
        assert response.pattern_metadata["steps_completed"] == 1

        second_task = researcher.received_tasks[1]
        assert isinstance(second_task, list)
        retry_message = second_task[-1]
        assert isinstance(retry_message, UserMessage)
        assert "Missing key facts" in retry_message.content
        assert "Add two concrete facts" in retry_message.content

    async def test_retry_exhaustion_moves_to_next_step(self) -> None:
        researcher = MockAgent("researcher", ["Attempt one.", "Attempt two."])
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Do the step",
                                agent_name="researcher",
                                reasoning="Only agent",
                            )
                        ]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=False,
                        failure_reason="Still incomplete",
                        confidence_score=0.7,
                        suggested_improvements=["Try again"],
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=False,
                        failure_reason="Still incomplete",
                        confidence_score=0.7,
                        suggested_improvements=["Try again"],
                    )
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
            max_step_retries=1,
        )

        response = await orchestrator.run("Do the task")

        assert response.pattern_metadata["steps_failed"] == 1
        assert response.pattern_metadata["steps_completed"] == 0
        assert response.pattern_metadata["total_retries"] == 1

    async def test_evaluation_fallback_heuristic_keeps_execution_moving(self) -> None:
        researcher = MockAgent("researcher", ["Completed the analysis and created a useful summary."])
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Analyze and summarize",
                                agent_name="researcher",
                                reasoning="Only agent",
                            )
                        ]
                    )
                )
            ],
            fail_on_calls={2},
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
        )

        response = await orchestrator.run("Analyze the subject")

        assert response.pattern_metadata["steps_completed"] == 1
        assert response.stop_message.content == "Execution plan completed"

    async def test_usage_includes_planning_evaluation_and_agent_calls(self) -> None:
        researcher = MockAgent("researcher", ["Completed the research task successfully."])
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Research facts",
                                agent_name="researcher",
                                reasoning="Best fit",
                            )
                        ]
                    ),
                    tokens_input=3,
                    tokens_output=1,
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    ),
                    tokens_input=2,
                    tokens_output=1,
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
        )

        response = await orchestrator.run("Research facts")

        assert response.usage.tokens_input == 10
        assert response.usage.tokens_output == 7
        assert response.usage.llm_calls == 3

    async def test_run_stream_emits_events_and_final_response(self) -> None:
        researcher = MockAgent("researcher", ["Completed the step successfully."])
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[
                            PlanStep(
                                task="Do the work",
                                agent_name="researcher",
                                reasoning="Only agent",
                            )
                        ]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    )
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
        )

        events: list[OrchestrationEvent] = []
        final_response: Optional[OrchestrationResponse] = None
        async for item in orchestrator.run_stream("Do the work", verbose=True):
            if isinstance(item, OrchestrationEvent):
                events.append(item)
            elif isinstance(item, OrchestrationResponse):
                final_response = item

        assert final_response is not None
        assert any(type(event).__name__ == "AgentSelectionEvent" for event in events)
        assert any(type(event).__name__ == "OrchestrationCompleteEvent" for event in events)

    async def test_reset_clears_plan_specific_state_between_runs(self) -> None:
        researcher = MockAgent(
            "researcher",
            ["Completed the first run.", "Completed the second run."],
        )
        client = FakePlanClient(
            responses=[
                make_result(
                    ExecutionPlan(
                        steps=[PlanStep(task="First", agent_name="researcher", reasoning="Only agent")]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    )
                ),
                make_result(
                    ExecutionPlan(
                        steps=[PlanStep(task="Second", agent_name="researcher", reasoning="Only agent")]
                    )
                ),
                make_result(
                    StepProgressEvaluation(
                        step_completed=True,
                        failure_reason="None",
                        confidence_score=0.9,
                        suggested_improvements=[],
                    )
                ),
            ]
        )
        orchestrator = PlanBasedOrchestrator(
            agents=[researcher],
            termination=MaxMessageTermination(10),
            model_client=client,
        )

        first = await orchestrator.run("first task")
        second = await orchestrator.run("second task")

        assert first.pattern_metadata["execution_plan"]["steps"][0]["task"] == "First"
        assert second.pattern_metadata["execution_plan"]["steps"][0]["task"] == "Second"
        assert orchestrator.current_step_index == 1

    async def test_review_gate_disabled_does_not_require_named_agents(self) -> None:
        worker = MockAgent("worker", ["done"])
        client = FakePlanClient(responses=[])

        orchestrator = PlanBasedOrchestrator(
            agents=[worker],
            termination=MaxMessageTermination(5),
            model_client=client,
            review_gate_policy=ReviewGatePolicy(enabled=False),
        )

        assert orchestrator.review_gate_policy.enabled is False

    async def test_review_gate_validation_requires_reviewer_and_revision_agents(self) -> None:
        worker = MockAgent("worker", ["done"])
        client = FakePlanClient(responses=[])

        with pytest.raises(
            ValueError,
            match="ReviewGatePolicy.reviewer_agent_name must reference an agent in the pool",
        ):
            PlanBasedOrchestrator(
                agents=[worker],
                termination=MaxMessageTermination(5),
                model_client=client,
                review_gate_policy=ReviewGatePolicy(enabled=True),
            )

        reviewer = MockAgent("reviewer", ["done"])
        with pytest.raises(
            ValueError,
            match="ReviewGatePolicy.revision_agent_name must reference an agent in the pool",
        ):
            PlanBasedOrchestrator(
                agents=[reviewer],
                termination=MaxMessageTermination(5),
                model_client=client,
                review_gate_policy=ReviewGatePolicy(enabled=True),
            )

    async def test_review_gate_requires_explicit_approval_when_enabled(self) -> None:
        reviewer = MockAgent("reviewer", ["Needs changes"])
        writer = MockAgent("writer", ["unused"])
        client = FakePlanClient(responses=[])
        orchestrator = PlanBasedOrchestrator(
            agents=[writer, reviewer],
            termination=MaxMessageTermination(5),
            model_client=client,
            review_gate_policy=ReviewGatePolicy(enabled=True),
        )
        step = PlanStep(
            task="Review the draft",
            agent_name="reviewer",
            reasoning="Quality gate",
        )
        result = AgentResponse(
            source="reviewer",
            context=AgentContext(
                messages=[
                    AssistantMessage(
                        content="The draft still needs stronger evidence.",
                        source="reviewer",
                    )
                ]
            ),
            finish_reason="stop",
        )

        evaluation = await orchestrator.evaluate_step_progress(step, result)

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

    async def test_review_gate_inserts_revision_loop_and_stops_at_max_rounds(self) -> None:
        reviewer = MockAgent("reviewer", ["Needs changes"])
        writer = MockAgent("writer", ["Revision"])
        client = FakePlanClient(responses=[])
        orchestrator = PlanBasedOrchestrator(
            agents=[writer, reviewer],
            termination=MaxMessageTermination(5),
            model_client=client,
            review_gate_policy=ReviewGatePolicy(enabled=True, max_revision_rounds=1),
        )
        orchestrator.execution_plan = ExecutionPlan(
            steps=[
                PlanStep(
                    task="Review the draft",
                    agent_name="reviewer",
                    reasoning="Quality gate",
                )
            ]
        )
        orchestrator.current_step_index = 0
        orchestrator._last_context_size = 0
        orchestrator._last_context_was_text = False
        result = AgentResponse(
            source="reviewer",
            context=AgentContext(
                messages=[
                    AssistantMessage(
                        content="Needs a clearer evidence-backed revision before approval.",
                        source="reviewer",
                    )
                ]
            ),
            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"
        assert orchestrator._review_revision_rounds == 1

        second_result = AgentResponse(
            source="reviewer",
            context=AgentContext(
                messages=[
                    AssistantMessage(content="Still not approved.", source="reviewer")
                ]
            ),
            finish_reason="stop",
        )
        orchestrator.current_step_index = 2

        await orchestrator.update_shared_state(second_result)

        assert orchestrator._review_revision_rounds == 1
        assert len(orchestrator.execution_plan.steps) == 3

    async def test_final_result_prefers_configured_agent_output(self) -> None:
        writer = MockAgent("writer", ["unused"])
        reviewer = MockAgent("reviewer", ["unused"])
        client = FakePlanClient(responses=[])
        orchestrator = PlanBasedOrchestrator(
            agents=[writer, reviewer],
            termination=MaxMessageTermination(5),
            model_client=client,
            review_gate_policy=ReviewGatePolicy(enabled=True),
            final_result_policy=FinalResultPolicy(prefer_agent="writer"),
        )
        orchestrator.shared_messages = [
            AssistantMessage(content="Initial research notes", source="reviewer"),
            AssistantMessage(content="Final approved article draft", source="writer"),
            AssistantMessage(content="APPROVED", source="reviewer"),
        ]

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