from __future__ import annotations

from core.agent.orchestrator import AgentChatOrchestrator


async def _intent_stage(
    *, message: str, context: dict[str, object]
) -> dict[str, object]:
    sequence = context.setdefault("sequence", [])
    if isinstance(sequence, list):
        sequence.append("intent")
    return {
        "content": f"intent:{message}",
        "usage": {"input_tokens": 2, "output_tokens": 1, "cost": "0.001000"},
    }


async def _execution_stage(
    *, message: str, context: dict[str, object]
) -> dict[str, object]:
    sequence = context.setdefault("sequence", [])
    if isinstance(sequence, list):
        sequence.append("execution")
    return {
        "content": f"execution:{message}",
        "usage": {"input_tokens": 3, "output_tokens": 2, "cost": "0.002000"},
    }


async def _organization_stage(
    *, message: str, context: dict[str, object]
) -> dict[str, object]:
    sequence = context.setdefault("sequence", [])
    if isinstance(sequence, list):
        sequence.append("organization")
    return {
        "content": "final answer",
        "usage": {"input_tokens": 4, "output_tokens": 1, "cost": "0.001500"},
    }


def test_orchestrator_runs_three_stages_in_order() -> None:
    orchestrator = AgentChatOrchestrator(
        intent_stage=_intent_stage,
        execution_stage=_execution_stage,
        organization_stage=_organization_stage,
    )

    result = orchestrator.run_sync(run_id="run-1", user_message="hello")

    assert result.context["sequence"] == ["intent", "execution", "organization"]
    assert result.output == "final answer"
    assert result.usage["total_tokens"] == 13
    assert result.events[0]["type"] == "run.started"
    assert result.events[-1]["type"] == "run.completed"


async def _failing_execution_stage(
    *, message: str, context: dict[str, object]
) -> dict[str, object]:
    sequence = context.setdefault("sequence", [])
    if isinstance(sequence, list):
        sequence.append("execution")
    raise RuntimeError("boom")


def test_orchestrator_stops_and_marks_failed_when_middle_stage_raises() -> None:
    orchestrator = AgentChatOrchestrator(
        intent_stage=_intent_stage,
        execution_stage=_failing_execution_stage,
        organization_stage=_organization_stage,
    )

    result = orchestrator.run_sync(run_id="run-2", user_message="hello")

    assert result.context["sequence"] == ["intent", "execution"]
    assert result.events[-1]["type"] == "run.failed"
    assert result.events[-1]["run_id"] == "run-2"
    assert "boom" in (result.events[-1].get("error") or "")
    assert result.failed is True
    assert "boom" in (result.error or "")


def test_orchestrator_emits_stage_event_payload_shape() -> None:
    orchestrator = AgentChatOrchestrator(
        intent_stage=_intent_stage,
        execution_stage=_execution_stage,
        organization_stage=_organization_stage,
    )

    result = orchestrator.run_sync(run_id="run-3", user_message="hello")

    for event in result.events:
        assert "type" in event
        assert event.get("run_id") == "run-3"

    stage_events = [
        event for event in result.events if event["type"] == "stage.completed"
    ]
    assert [event["stage"] for event in stage_events] == [
        "intent",
        "execution",
        "organization",
    ]
