# Copyright (c) Microsoft. All rights reserved.

from typing import Any

import pytest
from typing_extensions import Never

from agent_framework import (
    Executor,
    InProcRunnerContext,
    Workflow,
    WorkflowBuilder,
    WorkflowContext,
    WorkflowEvent,
    WorkflowEventSource,
    WorkflowRunResult,
    WorkflowRunState,
    handler,
)
from agent_framework._workflows._state import State


class FailingExecutor(Executor):
    """Executor that raises at runtime to test failure signaling."""

    @handler
    async def fail(self, msg: int, ctx: WorkflowContext) -> None:  # pragma: no cover - invoked via workflow
        raise RuntimeError("boom")


async def test_executor_failed_and_workflow_failed_events_streaming():
    failing = FailingExecutor(id="f")
    wf: Workflow = WorkflowBuilder(start_executor=failing).build()

    events: list[object] = []
    with pytest.raises(RuntimeError, match="boom"):
        async for ev in wf.run(0, stream=True):
            events.append(ev)

    # executor_failed event (type='executor_failed') should be emitted before workflow failed event
    executor_failed_events: list[WorkflowEvent[Any]] = [
        e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
    ]
    assert executor_failed_events, "executor_failed event should be emitted when start executor fails"
    assert executor_failed_events[0].executor_id == "f"
    assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK

    # Workflow-level failure and FAILED status should be surfaced
    failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
    assert failed_events
    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)
    status: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
    assert status and status[-1].state == WorkflowRunState.FAILED
    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status)

    # Verify executor_failed event comes before workflow failed event
    executor_failed_idx = events.index(executor_failed_events[0])
    workflow_failed_idx = events.index(failed_events[0])
    assert executor_failed_idx < workflow_failed_idx, (
        "executor_failed event should be emitted before workflow failed event"
    )


async def test_executor_failed_event_emitted_on_direct_execute():
    failing = FailingExecutor(id="f")
    ctx = InProcRunnerContext()
    state = State()
    with pytest.raises(RuntimeError, match="boom"):
        await failing.execute(
            0,
            ["START"],
            state,
            ctx,
        )
    drained = await ctx.drain_events()
    failed = [e for e in drained if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
    assert failed
    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed)


class PassthroughExecutor(Executor):
    """Executor that passes message to the next executor."""

    @handler
    async def passthrough(self, msg: int, ctx: WorkflowContext[int]) -> None:
        await ctx.send_message(msg)


async def test_executor_failed_event_from_second_executor_in_chain():
    """Test that executor_failed event is emitted when a non-start executor fails."""
    passthrough = PassthroughExecutor(id="passthrough")
    failing = FailingExecutor(id="failing")
    wf: Workflow = WorkflowBuilder(start_executor=passthrough).add_edge(passthrough, failing).build()

    events: list[object] = []
    with pytest.raises(RuntimeError, match="boom"):
        async for ev in wf.run(0, stream=True):
            events.append(ev)

    # executor_failed event should be emitted for the failing executor
    executor_failed_events: list[WorkflowEvent[Any]] = [
        e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
    ]
    assert executor_failed_events, "executor_failed event should be emitted when second executor fails"
    assert executor_failed_events[0].executor_id == "failing"
    assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK

    # Workflow-level failure should also be surfaced
    failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
    assert failed_events
    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)

    # Verify executor_failed event comes before workflow failed event
    executor_failed_idx = events.index(executor_failed_events[0])
    workflow_failed_idx = events.index(failed_events[0])
    assert executor_failed_idx < workflow_failed_idx, (
        "executor_failed event should be emitted before workflow failed event"
    )


class SimpleExecutor(Executor):
    """Executor that does nothing, for testing."""

    @handler
    async def run(self, msg: str, ctx: WorkflowContext[str]) -> None:  # pragma: no cover
        await ctx.send_message(msg)


class Requester(Executor):
    """Executor that always requests external info to test idle-with-requests state."""

    @handler
    async def ask(self, _: str, ctx: WorkflowContext) -> None:  # pragma: no cover
        await ctx.request_info("Mock request data", str)


async def test_idle_with_pending_requests_status_streaming():
    simple_executor = SimpleExecutor(id="simple")
    requester = Requester(id="req")
    wf = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()

    events = [ev async for ev in wf.run("start", stream=True)]  # Consume stream fully

    # Ensure a request was emitted
    assert any(isinstance(e, WorkflowEvent) and e.type == "request_info" for e in events)
    status_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
    assert len(status_events) >= 3
    assert status_events[-2].state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS
    assert status_events[-1].state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS


class Completer(Executor):
    """Executor that completes immediately with provided data for testing."""

    @handler
    async def run(self, msg: str, ctx: WorkflowContext[Never, str]) -> None:  # pragma: no cover  # type: ignore[valid-type]
        await ctx.yield_output(msg)


async def test_completed_status_streaming():
    c = Completer(id="c")
    wf = WorkflowBuilder(start_executor=c).build()
    events = [ev async for ev in wf.run("ok", stream=True)]  # no raise
    # Last status should be IDLE
    status = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
    assert status and status[-1].state == WorkflowRunState.IDLE
    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status)


async def test_started_and_completed_event_origins():
    c = Completer(id="c-origin")
    wf = WorkflowBuilder(start_executor=c).build()
    events = [ev async for ev in wf.run("payload", stream=True)]

    started = next(e for e in events if isinstance(e, WorkflowEvent) and e.type == "started")
    assert started.origin is WorkflowEventSource.FRAMEWORK

    # Check for IDLE status indicating completion
    idle_status = next(
        (e for e in events if isinstance(e, WorkflowEvent) and e.type == "status" and e.state == WorkflowRunState.IDLE),
        None,
    )
    assert idle_status is not None
    assert idle_status.origin is WorkflowEventSource.FRAMEWORK


async def test_non_streaming_final_state_helpers():
    # Completed case
    c = Completer(id="c")
    wf1 = WorkflowBuilder(start_executor=c).build()
    result1: WorkflowRunResult = await wf1.run("done")
    assert result1.get_final_state() == WorkflowRunState.IDLE

    # Idle-with-pending-request case
    simple_executor = SimpleExecutor(id="simple")
    requester = Requester(id="req")
    wf2 = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
    result2: WorkflowRunResult = await wf2.run("start")
    assert result2.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS


async def test_run_includes_status_events_completed():
    c = Completer(id="c2")
    wf = WorkflowBuilder(start_executor=c).build()
    result: WorkflowRunResult = await wf.run("ok")
    timeline = result.status_timeline()
    assert timeline, "Expected status timeline in non-streaming run() results"
    assert timeline[-1].state == WorkflowRunState.IDLE


async def test_run_includes_status_events_idle_with_requests():
    simple_executor = SimpleExecutor(id="simple")
    requester = Requester(id="req2")
    wf = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
    result: WorkflowRunResult = await wf.run("start")
    timeline = result.status_timeline()
    assert timeline, "Expected status timeline in non-streaming run() results"
    assert len(timeline) >= 3
    assert timeline[-2].state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS
    assert timeline[-1].state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
