"""T44: deterministic state aggregation probes; no LLM or network calls."""
import asyncio
from copy import deepcopy
from itertools import permutations
import operator
from typing import Annotated, TypedDict

import pytest
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
from langgraph.errors import InvalidUpdateError
from google.adk.agents import BaseAgent, ParallelAgent, SequentialAgent, LlmAgent
from google.adk.models import BaseLlm, LlmResponse
from google.adk.events import Event, EventActions
from google.adk.runners import InMemoryRunner
from google.genai import types
from google.adk import Workflow, Runner
from google.adk.workflow import JoinNode, START as ADK_START
from google.adk.sessions import InMemorySessionService

IDS = ("a", "b", "c")
ORDERS = list(permutations(IDS))


def merge_by_id(left, right):
    """Equal retransmission is accepted; conflicting ownership is rejected."""
    merged = deepcopy(left)
    for key, value in right.items():
        if key in merged and merged[key] != value:
            raise ValueError(f"conflicting result: {key}")
        merged[key] = deepcopy(value)
    return merged


def gather(expected, records, run_id):
    if len(expected) != len(set(expected)):
        raise ValueError("duplicate task ID")
    if set(records) != set(expected):
        raise ValueError("missing or unexpected task ID")
    ordered = []
    for task_id in expected:
        item = records[task_id]
        if (item.get("task_id") != task_id
                or item.get("run_id") != run_id
                or item.get("status") != "ok"
                or not isinstance(item.get("text"), str)
                or not item["text"].strip()):
            raise ValueError(f"invalid or stale result: {task_id}")
        ordered.append(item)
    return ordered


def record(task_id, run_id="r1"):
    return dict(task_id=task_id, run_id=run_id, status="ok", text=task_id.upper())


class Gate:
    """Release writers in an exact order after every worker has started."""
    def __init__(self, order):
        self.order = order
        self.ready = {i: asyncio.Event() for i in order}
        self.release = {i: asyncio.Event() for i in order}
        self.done = {i: asyncio.Event() for i in order}
        self.log = []

    async def enter(self, task_id):
        self.ready[task_id].set()
        await self.release[task_id].wait()

    def finish(self, task_id):
        self.log.append(task_id)
        self.done[task_id].set()

    async def control(self):
        await asyncio.gather(*(e.wait() for e in self.ready.values()))
        for task_id in self.order:
            self.release[task_id].set()
            await self.done[task_id].wait()


class ListState(TypedDict):
    results: Annotated[list, operator.add]
    final: list


class MapState(TypedDict):
    run_id: str
    expected: list[str]
    by_id: Annotated[dict, merge_by_id]
    final: list


@pytest.mark.parametrize("same", [False, True])
@pytest.mark.parametrize("limit", [1, 3])
def test_langgraph_collision(same, limit):
    class S(TypedDict):
        result: str
    b = StateGraph(S)
    b.add_node("a", lambda s: {"result": "A"})
    b.add_node("b", lambda s: {"result": "A" if same else "B"})
    for name in ("a", "b"):
        b.add_edge(START, name)
        b.add_edge(name, END)
    with pytest.raises(InvalidUpdateError):
        b.compile().invoke({"result": ""}, {"max_concurrency": limit})


def test_langgraph_sequential_overwrite():
    class S(TypedDict):
        result: str
    b = StateGraph(S)
    b.add_node("a", lambda s: {"result": "A"})
    b.add_node("b", lambda s: {"result": "B"})
    b.add_edge(START, "a"); b.add_edge("a", "b"); b.add_edge("b", END)
    assert b.compile().invoke({"result": ""})["result"] == "B"


@pytest.mark.parametrize("order", ORDERS)
def test_langgraph_reducer_and_snapshot(order):
    async def run():
        gate, seen = Gate(order), {}
        b = StateGraph(ListState)
        def worker(task_id):
            async def work(s):
                seen[task_id] = list(s["results"])
                await gate.enter(task_id)
                gate.finish(task_id)
                return {"results": [record(task_id)]}
            return work
        for i in IDS:
            b.add_node(i, worker(i)); b.add_edge(START, i)
        b.add_node("join", lambda s: {"final": sorted(s["results"], key=lambda x: x["task_id"])})
        b.add_edge(list(IDS), "join"); b.add_edge("join", END)
        _, out = await asyncio.wait_for(asyncio.gather(gate.control(), b.compile().ainvoke({"results": []})), 5)
        assert gate.log == list(order)
        assert seen == {i: [] for i in IDS}
        assert [r["task_id"] for r in out["final"]] == list(IDS)
        assert len(out["results"]) == 3
    asyncio.run(run())


def test_reducer_full_state_is_not_delta():
    b = StateGraph(ListState)
    for name in ("a", "b"):
        b.add_node(name, lambda s, n=name: {"results": s["results"] + [n]})
        b.add_edge(START, name); b.add_edge(name, END)
    out = b.compile().invoke({"results": ["seed"]})
    assert out["results"].count("seed") == 3


@pytest.mark.parametrize("order", ORDERS)
def test_id_reducer_laws(order):
    updates = [{i: record(i)} for i in order]
    out = {}
    for update in updates:
        before = deepcopy(out)
        merged = merge_by_id(out, update)
        assert out == before
        out = merged
    assert set(out) == set(IDS)
    assert merge_by_id(out, out) == out
    a, b, c = updates
    assert merge_by_id(merge_by_id(a, b), c) == merge_by_id(a, merge_by_id(b, c))
    assert merge_by_id(a, b) == merge_by_id(b, a)
    assert [x["task_id"] for x in gather(IDS, out, "r1")] == list(IDS)
    with pytest.raises(ValueError, match="conflicting"):
        merge_by_id(out, {"a": record("a", "different")})


@pytest.mark.parametrize("n", [0, 1, 3])
def test_langgraph_send(n):
    b = StateGraph(MapState)
    def route(s):
        return [Send("work", {"task_id": i, "run_id": s["run_id"]}) for i in s["expected"]] or "join"
    def work(s):
        return {"by_id": {s["task_id"]: record(s["task_id"], s["run_id"])}}
    def join(s):
        return {"final": gather(s["expected"], s["by_id"], s["run_id"])}
    b.add_node("work", work); b.add_node("join", join, defer=True)
    b.add_conditional_edges(START, route, ["work", "join"])
    b.add_edge("work", "join"); b.add_edge("join", END)
    out = b.compile().invoke({"run_id": "r1", "expected": list(IDS[:n]), "by_id": {}})
    assert len(out["final"]) == n


@pytest.mark.parametrize("style,count", [("separate", 2), ("barrier", 1), ("defer", 1)])
def test_langgraph_unequal_depth(style, count):
    b, snapshots = StateGraph(ListState), []
    for name in ("a", "b", "b2"):
        b.add_node(name, lambda s, n=name: {"results": [n]})
    def join(s):
        snapshots.append(list(s["results"]))
        return {}
    b.add_node("join", join, defer=style == "defer")
    b.add_edge(START, "a"); b.add_edge(START, "b"); b.add_edge("b", "b2")
    if style == "barrier": b.add_edge(["a", "b2"], "join")
    else:
        b.add_edge("a", "join"); b.add_edge("b2", "join")
    b.add_edge("join", END)
    b.compile().invoke({"results": []})
    assert len(snapshots) == count
    if count == 2: assert "b2" not in snapshots[0]
    assert set(snapshots[-1]) == {"a", "b", "b2"}


@pytest.mark.parametrize("defer,count", [(False, 0), (True, 1)])
def test_langgraph_unselected_branch(defer, count):
    b, calls = StateGraph(ListState), []
    b.add_node("a", lambda s: {"results": ["a"]})
    b.add_node("b", lambda s: {"results": ["b"]})
    b.add_node("join", lambda s: calls.append(1) or {}, defer=defer)
    b.add_conditional_edges(START, lambda s: "a", ["a", "b"])
    if defer:
        b.add_edge("a", "join"); b.add_edge("b", "join")
    else: b.add_edge(["a", "b"], "join")
    b.add_edge("join", END)
    b.compile().invoke({"results": []})
    assert len(calls) == count


class Writer(BaseAgent):
    task_id: str
    key: str
    gate: object
    mode: str = "ok"
    async def _run_async_impl(self, ctx):
        await self.gate.enter(self.task_id)
        if self.mode == "raise":
            self.gate.finish(self.task_id)
            raise RuntimeError("worker failed")
        delta = {} if self.mode == "missing" else {self.key: record(self.task_id)}
        yield Event(author=self.name, branch=ctx.branch,
                    actions=EventActions(state_delta=delta, escalate=self.mode == "escalate"))
        self.gate.finish(self.task_id)


class Gatherer(BaseAgent):
    keys: dict
    strict: bool = True
    async def _run_async_impl(self, ctx):
        records = {i: ctx.session.state[k] for i, k in self.keys.items() if k in ctx.session.state}
        final = gather(IDS, records, "r1") if self.strict else list(records.values())
        yield Event(author=self.name, actions=EventActions(state_delta={"final": final}))


async def adk_run(order, *, shared=False, mode="ok", strict=True, initial=None):
    gate = Gate(order)
    keys = {i: "result" if shared else f"result_{i}" for i in IDS}
    workers = [Writer(name=f"worker_{i}", task_id=i, key=keys[i], gate=gate,
                      mode=mode if i == "b" else "ok") for i in IDS]
    parallel = ParallelAgent(name="parallel", sub_agents=workers)
    join = Gatherer(name="join", keys=keys, strict=strict)
    root = SequentialAgent(name="pipeline", sub_agents=[parallel, join])
    runner = InMemoryRunner(agent=root, app_name="t44")
    session = await runner.session_service.create_session(app_name="t44", user_id="u", state=initial or {})
    async def consume():
        return [e async for e in runner.run_async(user_id="u", session_id=session.id,
                new_message=types.Content(role="user", parts=[types.Part(text="start")]))]
    controller = asyncio.create_task(gate.control())
    try:
        events = await asyncio.wait_for(consume(), 5)
        state = (await runner.session_service.get_session(app_name="t44", user_id="u", session_id=session.id)).state
        return state, events, gate.log
    finally:
        controller.cancel()
        await asyncio.gather(controller, return_exceptions=True)
        await runner.close()


@pytest.mark.parametrize("order", ORDERS)
def test_adk_distinct_keys(order):
    state, events, log = asyncio.run(adk_run(order))
    assert log == list(order)
    assert [x["task_id"] for x in state["final"]] == list(IDS)
    authors = [e.author for e in events]
    assert authors[-1] == "join"


@pytest.mark.parametrize("order", ORDERS)
def test_adk_same_key_silently_overwrites(order):
    state, events, log = asyncio.run(adk_run(order, shared=True, strict=False))
    assert state["result"]["task_id"] == order[-1]
    assert len({x["task_id"] for x in state["final"]}) == 1
    assert len([e for e in events if e.author.startswith("worker_")]) == 3


def test_adk_missing_result():
    with pytest.raises(ValueError, match="missing"):
        asyncio.run(adk_run(IDS, mode="missing"))


def test_adk_old_result_is_rejected():
    with pytest.raises(ValueError, match="stale"):
        asyncio.run(adk_run(IDS, mode="missing", initial={"result_b": record("b", "old")}))


def test_adk_worker_exception():
    with pytest.raises(RuntimeError, match="worker failed"):
        asyncio.run(adk_run(("a", "b", "c"), mode="raise"))


def test_adk_escalation_does_not_prove_all_done():
    with pytest.raises(ValueError, match="missing"):
        asyncio.run(adk_run(("b", "a", "c"), mode="escalate"))


class FakeModel(BaseLlm):
    task_id: str
    gate: object
    async def generate_content_async(self, llm_request, stream=False):
        await self.gate.enter(self.task_id)
        self.gate.finish(self.task_id)
        yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text=self.task_id.upper())]))


@pytest.mark.parametrize("order", ORDERS)
@pytest.mark.parametrize("shared", [False, True])
def test_adk_llm_output_key(order, shared):
    async def run():
        gate = Gate(order)
        workers = [LlmAgent(name=f"worker_{i}", model=FakeModel(model="fake", task_id=i, gate=gate),
                           output_key="result" if shared else f"result_{i}") for i in IDS]
        runner = InMemoryRunner(agent=ParallelAgent(name="parallel", sub_agents=workers), app_name="t44")
        session = await runner.session_service.create_session(app_name="t44", user_id="u")
        async def consume():
            return [e async for e in runner.run_async(user_id="u", session_id=session.id,
                    new_message=types.Content(role="user", parts=[types.Part(text="start")]))]
        try:
            await asyncio.wait_for(asyncio.gather(gate.control(), consume()), 5)
            state = (await runner.session_service.get_session(app_name="t44", user_id="u", session_id=session.id)).state
            if shared: assert state["result"] == order[-1].upper()
            else: assert {i: state[f"result_{i}"] for i in IDS} == {i: i.upper() for i in IDS}
        finally: await runner.close()
    asyncio.run(run())


@pytest.mark.parametrize("order", ORDERS)
def test_adk_workflow_join_by_node_name(order):
    async def run():
        gate = Gate(order)
        async def a(node_input: str):
            await gate.enter("a"); gate.finish("a"); return record("a")
        async def b(node_input: str):
            await gate.enter("b"); gate.finish("b"); return record("b")
        async def c(node_input: str):
            await gate.enter("c"); gate.finish("c"); return record("c")
        def collect(node_input: dict):
            return gather(IDS, node_input, "r1")
        workflow = Workflow(name="fan", edges=[(ADK_START, (a, b, c), JoinNode(name="join"), collect)])
        service = InMemorySessionService()
        runner = Runner(node=workflow, app_name="t44", session_service=service)
        session = await service.create_session(app_name="t44", user_id="u")
        async def consume():
            return [e async for e in runner.run_async(user_id="u", session_id=session.id,
                    new_message=types.Content(role="user", parts=[types.Part(text="topic")]))]
        try:
            _, events = await asyncio.wait_for(asyncio.gather(gate.control(), consume()), 5)
            assert gate.log == list(order)
            assert [r["task_id"] for r in events[-1].output] == list(IDS)
        finally: await runner.close()
    asyncio.run(run())


@pytest.mark.parametrize("bad", ["missing", "extra", "wrong_id", "empty", "failed", "duplicate_expected"])
def test_gather_rejects_bad_manifest(bad):
    records = {i: record(i) for i in IDS}
    expected = list(IDS)
    if bad == "missing": del records["b"]
    elif bad == "extra": records["d"] = record("d")
    elif bad == "wrong_id": records["b"]["task_id"] = "a"
    elif bad == "empty": records["b"]["text"] = " "
    elif bad == "failed": records["b"]["status"] = "error"
    else: expected.append("a")
    with pytest.raises(ValueError): gather(expected, records, "r1")