import pytest

from aegis.agents import EvaluatorAgent, PlannerAgent
from aegis.core import (
    Agent,
    AgentContext,
    AgentRegistry,
    AgentResult,
    Conductor,
    MemoryStore,
    Task,
    WorkflowEngine,
)


class WorkerAgent(Agent):
    name = "worker"

    async def run(self, context: AgentContext) -> AgentResult:
        return AgentResult(
            agent_name=self.name,
            success=True,
            output={
                "task": context.task.name,
                "completed": True,
            },
        )


class FailingWorkerAgent(Agent):
    name = "failing-worker"

    async def run(self, context: AgentContext) -> AgentResult:
        return AgentResult(
            agent_name=self.name,
            success=False,
            errors=["Worker execution failed."],
        )


class FailingPlannerAgent(Agent):
    name = "planner"

    async def run(self, context: AgentContext) -> AgentResult:
        return AgentResult(
            agent_name=self.name,
            success=False,
            errors=["Planning failed."],
        )


def make_context() -> AgentContext:
    return AgentContext(
        task=Task(
            name="Execute workflow",
            description="Run and evaluate a target agent.",
        )
    )


@pytest.mark.asyncio
async def test_workflow_runs_plan_execution_and_evaluation() -> None:
    registry = AgentRegistry(
        [
            PlannerAgent(),
            WorkerAgent(),
            EvaluatorAgent(),
        ]
    )
    memory = MemoryStore()
    engine = WorkflowEngine(
        conductor=Conductor(registry),
        memory=memory,
    )
    context = make_context()

    result = await engine.run(context, "worker")

    assert result.plan.agent_name == "planner"
    assert result.plan.success is True

    assert result.execution.agent_name == "worker"
    assert result.execution.success is True
    assert result.execution.output["completed"] is True

    assert result.evaluation.agent_name == "evaluator"
    assert result.evaluation.success is True
    assert result.evaluation.output["passed"] is True
    assert result.evaluation.output["score"] == 100

    assert len(result.memory_keys) == 3


@pytest.mark.asyncio
async def test_workflow_stores_all_results_in_memory() -> None:
    registry = AgentRegistry(
        [
            PlannerAgent(),
            WorkerAgent(),
            EvaluatorAgent(),
        ]
    )
    memory = MemoryStore()
    engine = WorkflowEngine(
        conductor=Conductor(registry),
        memory=memory,
    )
    context = make_context()

    await engine.run(context, "worker")

    prefix = f"workflow.{context.run_id}"

    plan = memory.get(f"{prefix}.plan")
    execution = memory.get(f"{prefix}.execution")
    evaluation = memory.get(f"{prefix}.evaluation")

    assert plan.metadata["stage"] == "plan"
    assert plan.metadata["agent_name"] == "planner"

    assert execution.metadata["stage"] == "execution"
    assert execution.metadata["agent_name"] == "worker"

    assert evaluation.metadata["stage"] == "evaluation"
    assert evaluation.metadata["agent_name"] == "evaluator"


@pytest.mark.asyncio
async def test_workflow_evaluates_failed_execution() -> None:
    registry = AgentRegistry(
        [
            PlannerAgent(),
            FailingWorkerAgent(),
            EvaluatorAgent(),
        ]
    )
    engine = WorkflowEngine(
        conductor=Conductor(registry),
        memory=MemoryStore(),
    )

    result = await engine.run(
        make_context(),
        "failing-worker",
    )

    assert result.execution.success is False
    assert result.evaluation.success is True
    assert result.evaluation.output["passed"] is False
    assert result.evaluation.output["score"] == 25


@pytest.mark.asyncio
async def test_workflow_stops_when_planner_fails() -> None:
    registry = AgentRegistry(
        [
            FailingPlannerAgent(),
            WorkerAgent(),
            EvaluatorAgent(),
        ]
    )
    memory = MemoryStore()
    engine = WorkflowEngine(
        conductor=Conductor(registry),
        memory=memory,
    )
    context = make_context()

    with pytest.raises(
        RuntimeError,
        match="Planner failed to produce a valid plan",
    ):
        await engine.run(context, "worker")

    prefix = f"workflow.{context.run_id}"

    assert memory.contains(f"{prefix}.plan") is True
    assert memory.contains(f"{prefix}.execution") is False
    assert memory.contains(f"{prefix}.evaluation") is False