"""Tests for graph runtime built on top of Pipeline."""
from __future__ import annotations

import asyncio

import pytest

from logos_adk.event_bus import SQLiteEventBus
from logos_adk.swarm import Swarm
from logos_adk.workflow import Graph, SQLiteGraphCheckpointStore
from logos_adk.workflow.steps import ApprovalStep, PublishEventStep, WaitForEventStep


@pytest.mark.asyncio
async def test_graph_runs_conditional_routes_and_returns_graph_result():
    async def draft(ctx):
        ctx["needs_review"] = True
        return "draft"

    async def review(ctx):
        return f"{ctx['_last']} reviewed"

    async def publish(ctx):
        return f"{ctx['_last']} published"

    graph = (
        Graph("release_flow")
        .node(draft, name="draft")
        .node(review, name="review")
        .node(publish, name="publish")
        .conditional_routes("draft", {"review": lambda ctx: ctx["needs_review"]}, default="publish")
        .edge("review", "publish")
        .entrypoint("draft")
        .terminal("publish")
    )

    response = await graph.run("ship it")

    assert response.content == "draft reviewed published"
    raw = response.raw["graph_result"]
    assert raw["status"] == "completed"
    assert raw["routing_table"]["draft"][0]["target"] == "review"
    assert raw["outputs"]["publish"]["output"] == "draft reviewed published"
    assert any(event["event"] == "edge_selected" for event in raw["node_events"])


@pytest.mark.asyncio
async def test_graph_compile_to_executor_wraps_into_pipeline_and_preserves_graph_results():
    async def start(ctx):
        return "alpha"

    async def finish(ctx):
        return f"{ctx['_last']}-omega"

    graph = (
        Graph("alpha_graph")
        .node(start, name="start")
        .node(finish, name="finish")
        .edge("start", "finish")
        .terminal("finish")
    )
    executor = graph.compile_to_executor(pipeline_name="alpha_executor")

    response = await executor.run("ignored")

    assert response.content == "alpha-omega"
    raw = response.raw["pipeline_result"]
    assert raw["outputs"]["alpha_graph"]["output"] == "alpha-omega"
    assert raw["graph_results"]["alpha_graph"]["outputs"]["finish"]["output"] == "alpha-omega"


@pytest.mark.asyncio
async def test_graph_stream_events_emits_node_level_events():
    async def first(ctx):
        return "one"

    async def second(ctx):
        return f"{ctx['_last']}-two"

    graph = (
        Graph("event_graph")
        .node(first, name="first")
        .node(second, name="second")
        .edge("first", "second")
        .terminal("second")
    )

    events = [event async for event in graph.stream_events("go")]

    assert events[0].event == "graph_started"
    assert any(event.event == "node_started" and event.node_name == "first" for event in events)
    assert any(event.event == "node_completed" and event.node_name == "second" for event in events)
    assert events[-1].event == "graph_completed"


@pytest.mark.asyncio
async def test_graph_stream_ui_events_emits_correlation_ids_and_separate_payload():
    async def first(ctx):
        return "one"

    async def second(ctx):
        return f"{ctx['_last']}-two"

    graph = (
        Graph("ui_event_graph")
        .node(first, name="first")
        .node(second, name="second")
        .edge("first", "second")
        .terminal("second")
    )

    payloads = [
        item
        async for item in graph.stream_ui_events(
            "go",
            session_id="sess-1",
            correlation_id="corr-123",
            _trace_id="trace-abc",
        )
    ]

    assert payloads[0]["event"] == "graph_started"
    assert payloads[0]["correlation_id"] == "corr-123"
    assert payloads[0]["trace_id"] == "trace-abc"
    assert payloads[0]["session_id"] == "sess-1"
    assert payloads[0]["event_id"]
    assert payloads[0]["graph_run_id"]
    assert payloads[0]["sequence"] == 1
    assert any(item["event"] == "edge_selected" and item["route_source"] == "first" for item in payloads)
    assert payloads[-1]["event"] == "graph_completed"


@pytest.mark.asyncio
async def test_graph_retries_failed_node_using_metadata_policy():
    attempts = {"count": 0}

    async def flaky(ctx):
        attempts["count"] += 1
        if attempts["count"] < 2:
            raise RuntimeError("temporary upstream failure")
        return "ok"

    graph = (
        Graph("retry_graph")
        .node(
            flaky,
            name="flaky",
            metadata={"retry": {"max_attempts": 2, "base_delay": 0.0, "backoff": "fixed"}},
        )
        .terminal("flaky")
    )

    response = await graph.run("go")

    assert attempts["count"] == 2
    assert response.content == "ok"
    assert response.raw["graph_result"]["outputs"]["flaky"]["metadata"]["attempt"] == 2


@pytest.mark.asyncio
async def test_graph_times_out_node_using_metadata_timeout():
    async def slow(ctx):
        await asyncio.sleep(0.05)
        return "late"

    graph = (
        Graph("timeout_graph")
        .node(slow, name="slow", metadata={"timeout": 0.01})
        .terminal("slow")
    )

    response = await graph.run("go")

    raw = response.raw["graph_result"]
    assert raw["status"] == "failed"
    assert raw["outputs"]["slow"]["status"] == "failed"
    assert "exceeded timeout=0.01s" in raw["outputs"]["slow"]["metadata"]["error"]


@pytest.mark.asyncio
async def test_graph_run_serializes_ui_event_payloads_separately():
    async def first(ctx):
        return "one"

    graph = Graph("serialized_ui_events").node(first, name="first").terminal("first")

    response = await graph.run("go", session_id="sess-2", correlation_id="corr-456", _trace_id="trace-456")

    raw = response.raw["graph_result"]
    assert raw["ui_event_payloads"]
    assert raw["ui_event_payloads"][0]["correlation_id"] == "corr-456"
    assert raw["ui_event_payloads"][0]["trace_id"] == "trace-456"
    assert raw["ui_event_payloads"][0]["session_id"] == "sess-2"


@pytest.mark.asyncio
async def test_graph_resume_restores_paused_context_by_correlation_id():
    async def draft(ctx):
        ctx.blackboard["draft"] = "saved-draft"
        return "draft"

    async def publish(ctx):
        return f"{ctx['draft']}::{ctx.blackboard['draft']}::{ctx['review']}"

    graph = (
        Graph("resume_graph")
        .node(draft, name="draft")
        .node(ApprovalStep("review", "Need approval"), name="review")
        .node(publish, name="publish")
        .edge("draft", "review")
        .edge("review", "publish")
        .entrypoint("draft")
        .terminal("publish")
    )

    initial_events = [
        item
        async for item in graph.stream_ui_events(
            "go",
            session_id="sess-resume",
            correlation_id="corr-resume",
        )
    ]
    assert initial_events[-1]["event"] == "graph_completed"
    assert initial_events[-1]["status"] == "paused"

    resumed_events = [
        item
        async for item in graph.resume(
            "approved",
            session_id="sess-resume",
            correlation_id="corr-resume",
        )
    ]
    assert resumed_events[-1]["event"] == "graph_completed"
    assert resumed_events[-1]["status"] == "completed"
    assert resumed_events[-1]["output"] == "draft::saved-draft::approved"


@pytest.mark.asyncio
async def test_graph_resume_survives_restart_with_sqlite_checkpoint_store(tmp_path):
    db_path = tmp_path / "graph_checkpoints.db"
    before_store = SQLiteGraphCheckpointStore(str(db_path))
    before_calls = {"draft": 0, "publish": 0}

    async def before_draft(ctx):
        before_calls["draft"] += 1
        ctx.blackboard["draft"] = "saved-draft"
        return "draft"

    async def before_publish(ctx):
        before_calls["publish"] += 1
        return f"{ctx['draft']}::{ctx.blackboard['draft']}::{ctx['review']}"

    graph_before = (
        Graph("restart_graph", checkpoint_store=before_store)
        .node(before_draft, name="draft")
        .node(ApprovalStep("review", "Need approval"), name="review")
        .node(before_publish, name="publish")
        .edge("draft", "review")
        .edge("review", "publish")
        .entrypoint("draft")
        .terminal("publish")
    )

    paused_events = [
        item
        async for item in graph_before.stream_ui_events(
            "go",
            correlation_id="corr-restart",
        )
    ]
    assert paused_events[-1]["status"] == "paused"

    after_calls = {"draft": 0, "publish": 0}

    async def after_draft(ctx):
        after_calls["draft"] += 1
        ctx.blackboard["draft"] = "should-not-rerun"
        return "draft-rerun"

    async def after_publish(ctx):
        after_calls["publish"] += 1
        return f"{ctx['draft']}::{ctx.blackboard['draft']}::{ctx['review']}"

    graph_after = (
        Graph(
            "restart_graph",
            checkpoint_store=SQLiteGraphCheckpointStore(str(db_path)),
        )
        .node(after_draft, name="draft")
        .node(ApprovalStep("review", "Need approval"), name="review")
        .node(after_publish, name="publish")
        .edge("draft", "review")
        .edge("review", "publish")
        .entrypoint("draft")
        .terminal("publish")
    )

    resumed_events = [
        item
        async for item in graph_after.resume(
            "approved",
            correlation_id="corr-restart",
        )
    ]

    assert resumed_events[-1]["event"] == "graph_completed"
    assert resumed_events[-1]["status"] == "completed"
    assert resumed_events[-1]["output"] == "draft::saved-draft::approved"
    assert before_calls["draft"] == 1
    assert before_calls["publish"] == 0
    assert after_calls["draft"] == 0
    assert after_calls["publish"] == 1


@pytest.mark.asyncio
async def test_graph_can_hybridize_with_swarm_via_persistent_event_bus(tmp_path):
    graph_db_path = tmp_path / "graph_checkpoints.db"
    event_bus_path = tmp_path / "event_bus.db"
    before_calls = {"draft": 0, "finalize": 0}

    async def before_draft(ctx):
        before_calls["draft"] += 1
        return "summer sale"

    async def before_finalize(ctx):
        before_calls["finalize"] += 1
        return f"{ctx['draft']}::{ctx['await_seo']['optimized']}"

    graph_before = (
        Graph(
            "hybrid_graph",
            checkpoint_store=SQLiteGraphCheckpointStore(str(graph_db_path)),
        )
        .node(before_draft, name="draft")
        .node(
            PublishEventStep(
                "publish_seo",
                SQLiteEventBus(str(event_bus_path)),
                topic="need_seo",
                payload=lambda ctx: {"prompt": ctx["draft"]},
            ),
            name="publish_seo",
        )
        .node(
            WaitForEventStep(
                "await_seo",
                SQLiteEventBus(str(event_bus_path)),
                topic="seo_done",
                message="Waiting for SEO optimization",
            ),
            name="await_seo",
        )
        .node(before_finalize, name="finalize")
        .edge("draft", "publish_seo")
        .edge("publish_seo", "await_seo")
        .edge("await_seo", "finalize")
        .entrypoint("draft")
        .terminal("finalize")
    )

    initial_events = [
        item
        async for item in graph_before.stream_ui_events(
            "launch",
            correlation_id="corr-hybrid",
        )
    ]
    assert initial_events[-1]["status"] == "paused"

    swarm = Swarm("seo-swarm", event_bus=SQLiteEventBus(str(event_bus_path)))

    async def optimize(event):
        return {"optimized": f"seo::{event.payload['prompt']}"}

    swarm.subscribe("need_seo", optimize, publish_to="seo_done")
    processed = await swarm.drain()
    assert processed == 1

    after_calls = {"draft": 0, "finalize": 0}

    async def after_draft(ctx):
        after_calls["draft"] += 1
        return "should-not-rerun"

    async def after_finalize(ctx):
        after_calls["finalize"] += 1
        return f"{ctx['draft']}::{ctx['await_seo']['optimized']}"

    graph_after = (
        Graph(
            "hybrid_graph",
            checkpoint_store=SQLiteGraphCheckpointStore(str(graph_db_path)),
        )
        .node(after_draft, name="draft")
        .node(
            PublishEventStep(
                "publish_seo",
                SQLiteEventBus(str(event_bus_path)),
                topic="need_seo",
                payload=lambda ctx: {"prompt": ctx["draft"]},
            ),
            name="publish_seo",
        )
        .node(
            WaitForEventStep(
                "await_seo",
                SQLiteEventBus(str(event_bus_path)),
                topic="seo_done",
                message="Waiting for SEO optimization",
            ),
            name="await_seo",
        )
        .node(after_finalize, name="finalize")
        .edge("draft", "publish_seo")
        .edge("publish_seo", "await_seo")
        .edge("await_seo", "finalize")
        .entrypoint("draft")
        .terminal("finalize")
    )

    resumed_events = [
        item
        async for item in graph_after.resume(
            "continue",
            correlation_id="corr-hybrid",
        )
    ]

    assert resumed_events[-1]["event"] == "graph_completed"
    assert resumed_events[-1]["status"] == "completed"
    assert resumed_events[-1]["output"] == "summer sale::seo::summer sale"
    assert before_calls["draft"] == 1
    assert before_calls["finalize"] == 0
    assert after_calls["draft"] == 0
    assert after_calls["finalize"] == 1
