"""Tests for WebUI execution and session behavior."""

from __future__ import annotations

import pytest
from pydantic import BaseModel

from agentbyte.messages import UserMessage
from agentbyte.webui.execution import ExecutionEngine
from agentbyte.webui.sessions import SessionBindingError, SessionManager
from agentbyte.workflow import FunctionStep, StepMetadata, Workflow, WorkflowConfig

from .helpers import DummyAgent, DummyOrchestrator


@pytest.mark.asyncio
async def test_execute_agent_stream_emits_final_response_and_persists_session() -> None:
    manager = SessionManager()
    engine = ExecutionEngine(manager)
    agent = DummyAgent()

    events = []
    async for item in engine.execute_agent_stream(
        "dummy-agent",
        agent,
        messages=[UserMessage(content="hello", source="user")],
        session_id="session-a",
        user_id="user-a",
        stream_tokens=False,
    ):
        events.append(item)

    assert any("echo: hello" in item for item in events)
    stored = await manager.get("session-a")
    assert stored is not None
    assert stored.user_id == "user-a"
    assert stored.messages[-1].content == "echo: hello"


@pytest.mark.asyncio
async def test_execute_agent_stream_supports_approval_resume() -> None:
    manager = SessionManager()
    engine = ExecutionEngine(manager)
    agent = DummyAgent(require_approval=True)

    first_run = []
    async for item in engine.execute_agent_stream(
        "dummy-agent",
        agent,
        messages=[UserMessage(content="needs approval", source="user")],
        session_id="approval-session",
        stream_tokens=False,
    ):
        first_run.append(item)

    assert any("approval_needed" in item for item in first_run)

    context = await manager.get("approval-session")
    assert context is not None
    assert context.waiting_for_approval is True

    approval = context.pending_approval_requests[0].create_response(approved=True)
    second_run = []
    async for item in engine.execute_agent_stream(
        "dummy-agent",
        agent,
        messages=[],
        session_id="approval-session",
        stream_tokens=False,
        approval_responses=[approval],
    ):
        second_run.append(item)

    assert any("echo: needs approval" in item for item in second_run)
    resumed = await manager.get("approval-session")
    assert resumed is not None
    assert resumed.waiting_for_approval is False


@pytest.mark.asyncio
async def test_execute_orchestrator_stream_emits_orchestration_response() -> None:
    manager = SessionManager()
    engine = ExecutionEngine(manager)
    orchestrator = DummyOrchestrator(DummyAgent())

    events = []
    async for item in engine.execute_orchestrator_stream(
        "dummy-orchestrator",
        orchestrator,
        messages=[UserMessage(content="start orchestration", source="user")],
        session_id="orch-session",
        user_id="user-o",
    ):
        events.append(item)

    assert any("orchestration_complete" in item or "final_result" in item for item in events)
    context = await manager.get("orch-session")
    assert context is not None
    assert context.user_id == "user-o"
    assert context.metadata["final_result"]
    assert context.metadata["stop_reason"] == "Pattern complete"


@pytest.mark.asyncio
async def test_execute_workflow_stream_emits_workflow_events() -> None:
    manager = SessionManager()
    engine = ExecutionEngine(manager)

    workflow = Workflow(config=WorkflowConfig(name="demo-workflow", description="Demo"))

    class WorkflowInput(BaseModel):
        task: str

    class WorkflowOutput(BaseModel):
        result: str

    async def step_fn(input_data, context):
        return {"result": "done"}

    step = FunctionStep(
        step_id="run",
        metadata=StepMetadata(name="Run"),
        input_type=WorkflowInput,
        output_type=WorkflowOutput,
        func=step_fn,
    )
    workflow.add_step(step)
    workflow.set_start_step("run")
    workflow.add_end_step("run")

    events = []
    async for item in engine.execute_workflow_stream(
        "demo-workflow",
        workflow,
        input_data={"task": "demo"},
        session_id="wf-session",
        user_id="user-w",
    ):
        events.append(item)

    assert any("workflow_started" in item for item in events)
    assert any("workflow_completed" in item for item in events)
    assert any('"final_output":{"result":"done"}' in item for item in events)
    context = await manager.get("wf-session")
    assert context is not None
    assert context.user_id == "user-w"
    assert context.metadata["last_input"]["task"] == "demo"


@pytest.mark.asyncio
async def test_execute_agent_rejects_conflicting_user_id_on_existing_session() -> None:
    manager = SessionManager()
    engine = ExecutionEngine(manager)
    agent = DummyAgent()

    await engine.execute_agent(
        "dummy-agent",
        agent,
        messages=[UserMessage(content="hello", source="user")],
        session_id="session-conflict",
        user_id="alice",
    )

    with pytest.raises(SessionBindingError, match="owned by user alice"):
        await engine.execute_agent(
            "dummy-agent",
            agent,
            messages=[UserMessage(content="again", source="user")],
            session_id="session-conflict",
            user_id="bob",
        )
