import asyncio

import pytest
from typing_extensions import TypedDict

from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context
from saf_python_sdk.types import Command, Send

pytestmark = pytest.mark.anyio


class StreamState(TypedDict):
    done: bool


async def test_custom_stream_receive_and_close() -> None:
    graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)

    async def start_node(ctx: Context, state: StreamState) -> Command:
        ctx.send_custom_stream_event({"step": "start", "value": 1})
        await asyncio.sleep(0.08)
        ctx.send_custom_stream_event({"step": "start", "value": 2})
        return Command(update=state, goto=Send("finish_node", None))

    async def finish_node(state: StreamState) -> dict[str, bool]:
        return {"done": True}

    graph.add_entry_node(start_node)
    graph.add_finish_node(finish_node)

    handler = await graph.compile().astart({"done": False}, stream_mode="custom")

    event = await handler.receive_stream()
    assert isinstance(event, dict)
    assert event["step"] == "start"
    assert event["value"] == 1

    handler.close_stream()
    assert await handler.receive_stream() is None

    result = await handler.aresult()
    assert result["done"] is True


async def test_only_custom_stream_mode_supported() -> None:
    graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)

    async def start_node(ctx: Context, state: StreamState) -> Command:
        ctx.send_custom_stream_event({"hello": "world"})
        return Command(update=state)

    graph.add_entry_node(start_node)

    handler = await graph.compile().astart({"done": False}, stream_mode="values")
    with pytest.raises(Exception, match="only `custom` is supported"):
        await handler.aresult()