"""Graph orchestration layered on top of Pipeline primitives."""
from __future__ import annotations

import asyncio
import time
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from typing import Any, AsyncIterator, Awaitable, Callable, TypedDict, cast
from uuid import uuid4

from logos_adk.checkpoint_json import JSONObject, ensure_json_object
from logos_adk.graph_checkpoint_store import GraphCheckpointRecord, GraphCheckpointStore
from logos_adk.reliability import RetryPolicy
from logos_adk.runnable import Runnable
from logos_adk.types import Response, StreamChunk, StreamDone, TextChunk, Usage
from logos_adk.workflow.context import PipelineContext, StepResult
from logos_adk.workflow.step import Step
from logos_adk.workflow.steps import AgentStep, FunctionStep


GraphTarget = Step | Runnable | Callable[[PipelineContext], Awaitable[Any]]
GraphCondition = Callable[[PipelineContext], bool]


@dataclass
class Node:
    """A named executable node inside a graph."""

    name: str
    step: Step
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class Edge:
    """A directed graph edge with optional conditional routing."""

    source: str
    target: str
    condition: GraphCondition | None = None
    label: str | None = None


@dataclass
class GraphNodeEvent:
    """Node-level execution event suitable for a UI event stream."""

    event_id: str
    event: str
    event_type: str
    graph_name: str
    graph_run_id: str
    pipeline_id: str
    trace_id: str | None = None
    session_id: str | None = None
    correlation_id: str | None = None
    sequence: int = 0
    timestamp: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
    parent_event_id: str | None = None
    node_name: str | None = None
    route_source: str | None = None
    route_target: str | None = None
    edge_label: str | None = None
    status: str | None = None
    duration_ms: float | None = None
    output: Any = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class GraphExecutionResult:
    """Serializable result of executing a graph."""

    graph_name: str
    status: str
    last_output: Any
    outputs: dict[str, StepResult]
    events: list[GraphNodeEvent]
    ui_event_payloads: list[dict[str, Any]]
    routing_table: dict[str, list[dict[str, Any]]]
    duration_ms: float


@dataclass
class _PausedGraphRun:
    correlation_id: str
    pipeline_id: str
    paused_node_name: str
    resume_from: str
    context_data: dict[str, Any]
    metadata_data: dict[str, Any]
    blackboard_data: dict[str, Any]
    blackboard_versions_data: dict[str, int]


class _PausedGraphRunCheckpointPayload(TypedDict):
    correlation_id: str
    pipeline_id: str
    paused_node_name: str
    resume_from: str
    context_data: JSONObject
    metadata_data: JSONObject
    blackboard_data: JSONObject
    blackboard_versions_data: JSONObject


class Graph(Step, Runnable):
    """Stateful graph runtime that compiles onto the existing Pipeline model."""

    def __init__(
        self,
        name: str,
        *,
        max_visits: int = 64,
        checkpoint_store: GraphCheckpointStore | None = None,
    ) -> None:
        super().__init__(name)
        self._nodes: dict[str, Node] = {}
        self._edges: list[Edge] = []
        self._entrypoint: str | None = None
        self._terminal_nodes: set[str] = set()
        self._max_visits = max(1, int(max_visits))
        self._auto_counter = 0
        self._source_path: str | None = None
        self._paused_runs: dict[str, _PausedGraphRun] = {}
        self._checkpoint_store = checkpoint_store

    def to_dict(self) -> dict[str, Any]:
        """Convert graph to a YAML-compatible dictionary."""
        nodes_data = []
        for node in self._nodes.values():
            node_item = {"name": node.name}
            if node.metadata:
                node_item["metadata"] = node.metadata

            # Detect node type for YAML serialization
            step = node.step
            from logos_adk.workflow.steps import ApprovalStep, AgentStep, FunctionStep
            if isinstance(step, ApprovalStep):
                node_item["approval"] = {"message": step._message}
                if step._options:
                    node_item["approval"]["options"] = step._options
            elif isinstance(step, AgentStep):
                node_item["agent"] = step._agent.name
            elif isinstance(step, FunctionStep):
                node_item["function"] = step._fn.__name__
            elif isinstance(step, Graph):
                node_item["graph"] = step.to_dict()
            else:
                # Fallback: check if it's a wrapped Agent or Runnable
                agent = getattr(step, "_agent", None) or getattr(step, "agent", None)
                if agent and hasattr(agent, "name"):
                    node_item["agent"] = agent.name
                else:
                    node_item["type"] = step.__class__.__name__
            nodes_data.append(node_item)

        edges_data = []
        for edge in self._edges:
            edge_item = {"from": edge.source, "to": edge.target}
            if edge.label:
                edge_item["label"] = edge.label
            # Note: complex conditions (when) are hard to serialize back to string 
            # if they were passed as Python callables. 
            # For now, we assume simple string-based conditions if loaded from YAML.
            edges_data.append(edge_item)

        return {
            "name": self.name,
            "max_visits": self._max_visits,
            "entrypoint": self._entrypoint,
            "terminal_nodes": list(self._terminal_nodes),
            "nodes": nodes_data,
            "edges": edges_data,
        }

    def node(self, target: GraphTarget, *, name: str | None = None,
             metadata: dict[str, Any] | None = None) -> Graph:
        step = self._wrap_target(target, name=name)
        node = Node(name=step.name, step=step, metadata=metadata or {})
        self._nodes[node.name] = node
        if self._entrypoint is None:
            self._entrypoint = node.name
        return self

    def edge(self, source: str, target: str, *, when: GraphCondition | None = None,
             label: str | None = None) -> Graph:
        self._edges.append(Edge(source=source, target=target, condition=when, label=label))
        return self

    def conditional_routes(
        self,
        source: str,
        routes: dict[str, GraphCondition],
        *,
        default: str | None = None,
    ) -> Graph:
        for target, condition in routes.items():
            self.edge(source, target, when=condition, label=target)
        if default is not None:
            self.edge(source, default, label="default")
        return self

    def entrypoint(self, node_name: str) -> Graph:
        self._entrypoint = node_name
        return self

    def terminal(self, *node_names: str) -> Graph:
        self._terminal_nodes.update(node_names)
        return self

    def get_graph_definition(self) -> dict[str, Any]:
        """Return a serializable definition of the graph structure for UI/Dashboard."""
        nodes_data = []
        for node in self._nodes.values():
            node_item = {
                "name": node.name,
                "type": node.step.__class__.__name__,
                "metadata": node.metadata,
            }
            
            # Add YAML-specific fields for tools
            step = node.step
            from logos_adk.workflow.steps import ApprovalStep, AgentStep, FunctionStep
            if isinstance(step, ApprovalStep):
                node_item["approval"] = {"message": step._message, "options": step._options}
            elif isinstance(step, AgentStep):
                node_item["agent"] = step._agent.name
            elif isinstance(step, FunctionStep):
                node_item["function"] = step._fn.__name__
            
            nodes_data.append(node_item)

        return {
            "name": self.name,
            "entrypoint": self._entrypoint,
            "terminal_nodes": list(self._terminal_nodes),
            "nodes": nodes_data,
            "edges": [
                {
                    "source": edge.source,
                    "target": edge.target,
                    "label": edge.label,
                    "conditional": edge.condition is not None,
                }
                for edge in self._edges
            ],
        }

    def update_from_definition(self, definition: dict[str, Any]) -> None:
        """Update graph structure from a UI definition."""
        from logos_adk.workflow.steps import ApprovalStep, AgentStep
        
        new_entrypoint = definition.get("entrypoint")
        new_terminal_nodes = set(definition.get("terminal_nodes", []))
        
        self._entrypoint = new_entrypoint
        self._terminal_nodes = new_terminal_nodes
        
        # Update or create nodes
        for node_def in definition.get("nodes", []):
            name = node_def.get("name")
            metadata = node_def.get("metadata", {})
            
            if name in self._nodes:
                self._nodes[name].metadata = metadata
            else:
                # Dynamic node creation
                target = None
                if "approval" in node_def:
                    app_info = node_def["approval"]
                    target = ApprovalStep(name, app_info.get("message"), options=app_info.get("options"))
                elif "agent" in node_def:
                    # This is tricky: we need the Agent object. 
                    # For now, we'll try to find it in the existing agent nodes
                    agent_name = node_def["agent"]
                    # Find any existing AgentStep with this agent name
                    for existing_node in self._nodes.values():
                        if isinstance(existing_node.step, AgentStep) and existing_node.step._agent.name == agent_name:
                            target = AgentStep(name, existing_node.step._agent)
                            break
                
                if target:
                    self.node(target, name=name, metadata=metadata)
        
        # Replace edges
        self._edges = []
        for edge_def in definition.get("edges", []):
            self.edge(
                edge_def["source"] if "source" in edge_def else edge_def["from"], 
                edge_def["target"] if "target" in edge_def else edge_def["to"], 
                label=edge_def.get("label")
            )

    def routing_table(self) -> dict[str, list[dict[str, Any]]]:
        table: dict[str, list[dict[str, Any]]] = {name: [] for name in self._nodes}
        for edge in self._edges:
            table.setdefault(edge.source, []).append(
                {
                    "target": edge.target,
                    "label": edge.label,
                    "conditional": edge.condition is not None,
                }
            )
        return table

    def compile_to_executor(
        self,
        *,
        pipeline_name: str | None = None,
        step_name: str | None = None,
        retry: Any = None,
        timeout: float | None = None,
        state_store: Any = None,
        learning_engine: Any = None,
    ) -> Any:
        from logos_adk.workflow.pipeline import Pipeline

        pipeline = Pipeline(
            pipeline_name or f"{self.name}_pipeline",
            retry=retry,
            timeout=timeout,
            state_store=state_store,
            learning_engine=learning_engine,
        )
        pipeline.graph(self, name=step_name or self.name, retry=retry, timeout=timeout)
        return pipeline

    def compile(self, **kwargs: Any) -> Any:
        return self.compile_to_executor(**kwargs)

    async def execute(self, ctx: PipelineContext) -> StepResult:
        result = await self._execute_graph(ctx)
        ctx.metadata.setdefault("adk_graph_results", {})[self.name] = _serialize_graph_result(result)
        return StepResult(
            output=result.last_output,
            status=result.status,
            duration_ms=result.duration_ms,
            metadata={
                "node_outputs": {
                    name: _serialize_step_result(step_result)
                    for name, step_result in result.outputs.items()
                },
                "node_events": [asdict(event) for event in result.events],
                "ui_event_payloads": list(result.ui_event_payloads),
                "routing_table": result.routing_table,
            },
        )

    async def run(self, prompt: str, *, session_id: str | None = None, **kwargs: Any) -> Response:
        ctx = PipelineContext(
            pipeline_id=uuid4().hex,
            metadata={
                "trace_id": kwargs.get("_trace_id", uuid4().hex),
                "correlation_id": kwargs.get("correlation_id") or kwargs.get("_graph_correlation_id"),
                "workspace_id": kwargs.get("_workspace_id"),
                "task_type": kwargs.get("_task_type"),
                "user_segment": kwargs.get("_user_segment"),
                "session_id": session_id,
            },
        )
        ctx["_last"] = prompt
        result = await self._execute_graph(ctx)
        ctx.metadata.setdefault("adk_graph_results", {})[self.name] = _serialize_graph_result(result)
        last_output = "" if result.last_output is None else str(result.last_output)
        return Response(
            content=last_output,
            tool_calls=[],
            usage=Usage(input_tokens=0, output_tokens=0),
            raw={"graph_result": _serialize_graph_result(result)},
        )

    async def stream(self, prompt: str, **kwargs: Any) -> AsyncIterator[StreamChunk]:
        response = await self.run(prompt, session_id=kwargs.pop("session_id", None), **kwargs)
        if response.content:
            yield TextChunk(text=response.content)
        yield StreamDone(response=response)

    async def stream_events(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[GraphNodeEvent]:
        queue: asyncio.Queue[GraphNodeEvent] = asyncio.Queue()

        async def callback(event: GraphNodeEvent) -> None:
            await queue.put(event)

        ctx = PipelineContext(
            pipeline_id=uuid4().hex,
            metadata={
                "trace_id": kwargs.get("_trace_id", uuid4().hex),
                "correlation_id": correlation_id or kwargs.get("_graph_correlation_id"),
                "workspace_id": kwargs.get("_workspace_id"),
                "task_type": kwargs.get("_task_type"),
                "user_segment": kwargs.get("_user_segment"),
                "session_id": session_id,
                "_graph_event_callback": callback,
            },
        )
        ctx["_last"] = prompt
        task = asyncio.create_task(self._execute_graph(ctx))
        while True:
            if task.done() and queue.empty():
                break
            try:
                event = await asyncio.wait_for(queue.get(), timeout=0.05)
            except asyncio.TimeoutError:
                continue
            yield event
        await task

    async def stream_ui_events(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        async for event in self.stream_events(
            prompt,
            session_id=session_id,
            correlation_id=correlation_id,
            **kwargs,
        ):
            if isinstance(event, dict):
                yield event
            else:
                yield _serialize_graph_ui_event(event)

    async def resume(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        """Resume a paused graph execution with new input (approval/rejection)."""
        queue: asyncio.Queue[GraphNodeEvent | dict[str, Any]] = asyncio.Queue()

        async def callback(event: GraphNodeEvent | dict[str, Any]) -> None:
            await queue.put(event)

        target_correlation_id = correlation_id or kwargs.get("_graph_correlation_id")
        if not isinstance(target_correlation_id, str) or not target_correlation_id:
            raise ValueError("Graph resume requires correlation_id")
        paused_run = self._load_paused_run(target_correlation_id)
        if paused_run is None:
            raise ValueError(
                f"No paused graph run found for correlation_id '{target_correlation_id}'"
            )
        paused_node_name = kwargs.get("node_name") or paused_run.paused_node_name
        if not isinstance(paused_node_name, str) or not paused_node_name:
            raise ValueError("Graph resume requires a paused node name")

        ctx = PipelineContext(
            pipeline_id=paused_run.pipeline_id,
            data=dict(paused_run.context_data),
            metadata={
                **dict(paused_run.metadata_data),
                "_blackboard_versions": dict(paused_run.blackboard_versions_data),
            },
            blackboard=dict(paused_run.blackboard_data),
        )
        ctx.metadata.update(
            {
                "trace_id": kwargs.get("_trace_id", ctx.metadata.get("trace_id", uuid4().hex)),
                "correlation_id": target_correlation_id,
                "session_id": session_id if session_id is not None else ctx.metadata.get("session_id"),
                "_graph_event_callback": callback,
                "is_resume": True,
            }
        )
        ctx["_last"] = prompt
        
        # --- Self-Learning Capture ---
        # If the prompt starts with 'rejected:' and has feedback, teach the agent
        if prompt.startswith("rejected:") and paused_node_name:
            feedback = prompt.replace("rejected:", "").strip()
            # Search for the node (could be a sub-step in ParallelStep)
            target_node = self._nodes.get(paused_node_name)
            if not target_node:
                # Search inside parallel steps
                for parent_node in self._nodes.values():
                    from logos_adk.workflow.steps import ParallelStep
                    if isinstance(parent_node.step, ParallelStep):
                        for sub_step in parent_node.step._steps:
                            if sub_step.name == paused_node_name:
                                from logos_adk.workflow.steps import AgentStep
                                if isinstance(sub_step, AgentStep) and hasattr(sub_step._agent, "learn_from_feedback"):
                                    sub_step._agent.learn_from_feedback(feedback)
                                break
            elif hasattr(target_node.step, "_agent") and hasattr(target_node.step._agent, "learn_from_feedback"):
                target_node.step._agent.learn_from_feedback(feedback)
        # ----------------------------

        # If the node is a sub-step, we actually want to resume the PARENT node (the ParallelStep)
        actual_resume_node = paused_run.resume_from
        if paused_node_name in self._nodes:
            actual_resume_node = paused_node_name

        ctx.metadata["_resume_node"] = paused_node_name

        task = asyncio.create_task(self._execute_graph(ctx, resume_from=actual_resume_node))

        
        while True:
            if task.done() and queue.empty():
                break
            try:
                event = await asyncio.wait_for(queue.get(), timeout=0.05)
            except asyncio.TimeoutError:
                continue
            
            if isinstance(event, dict):
                yield event
            else:
                yield _serialize_graph_ui_event(event)
        await task

    async def _execute_graph(self, ctx: PipelineContext, resume_from: str | None = None) -> GraphExecutionResult:
        import logging
        graph_logger = logging.getLogger(f"logos_adk.graph.{self.name}")
        if resume_from:
            graph_logger.info(f"Возобновляю выполнение графа с узла: {resume_from}")
        else:
            graph_logger.info(f"Запускаю выполнение графа '{self.name}'")

        self._validate_graph()
        
        # --- Swarm Pulse Integration ---
        # Try to find global event bus to listen for swarm activity
        from logos_adk.runtime import _default_runtime_holder
        runtime = _default_runtime_holder.get()
        bus = getattr(runtime, "event_bus", None) if runtime else None
        
        if bus:
            async def swarm_pulse_observer(record):
                pulse_event = {
                    "event": "swarm_event",
                    "event_type": "swarm",
                    "topic": record.topic,
                    "publisher": record.publisher or "unknown",
                    "timestamp": record.created_at,
                    # Shorten payload for UI
                    "payload_summary": str(record.payload)[:100] + "..." if len(str(record.payload)) > 100 else str(record.payload)
                }
                ui_event_payloads.append(pulse_event)
                await self._dispatch_event(ctx, pulse_event)
            
            bus.on_publish(swarm_pulse_observer)
        # -------------------------------

        # Inject current graph instance for planning tools
        ctx.metadata["_executing_graph"] = self
        
        started = time.monotonic()
        outputs: dict[str, StepResult] = {}
        events: list[GraphNodeEvent] = []
        ui_event_payloads: list[dict[str, Any]] = []
        visits = 0
        current = resume_from or self._entrypoint
        status = "completed"
        last_output: Any = ctx.get("_last")
        event_context = self._ensure_event_context(ctx)

        graph_started = self._make_event(
            ctx,
            event_context,
            event="graph_started",
            event_type="graph",
        )
        events.append(graph_started)
        ui_event_payloads.append(_serialize_graph_ui_event(graph_started))
        await self._dispatch_event(ctx, graph_started)

        while current is not None:
            visits += 1
            if visits > self._max_visits:
                status = "failed"
                outputs[current] = StepResult(
                    output=None,
                    status="failed",
                    duration_ms=0,
                    metadata={"error": f"Graph exceeded max_visits={self._max_visits}"},
                )
                break

            node = self._nodes[current]
            node_started = self._make_event(
                ctx,
                event_context,
                event="node_started",
                event_type="node",
                node_name=node.name,
                parent_event_id=graph_started.event_id,
                metadata=dict(node.metadata),
            )
            events.append(node_started)
            ui_event_payloads.append(_serialize_graph_ui_event(node_started))
            await self._dispatch_event(ctx, node_started)

            result = await self._execute_node_step(node, ctx)
            outputs[node.name] = result
            last_output = result.output
            node_finished = self._make_event(
                ctx,
                event_context,
                event=f"node_{result.status}",
                event_type="node",
                node_name=node.name,
                parent_event_id=node_started.event_id,
                status=result.status,
                duration_ms=result.duration_ms,
                output=_serialize_output(result.output),
                metadata=dict(result.metadata),
            )
            events.append(node_finished)
            ui_event_payloads.append(_serialize_graph_ui_event(node_finished))
            await self._dispatch_event(ctx, node_finished)

            if result.status in {"failed", "paused"}:
                status = result.status
                if result.status == "paused":
                    self._store_paused_run(
                        ctx,
                        event_context,
                        node_name=node.name,
                        result=result,
                    )
                else:
                    self._clear_paused_run(event_context)
                break

            ctx[node.name] = result.output
            ctx["_last"] = result.output

            if node.name in self._terminal_nodes:
                break

            try:
                next_edge = self._select_next_edge(node.name, ctx)
            except Exception as exc:
                status = "failed"
                outputs[node.name] = StepResult(
                    output=result.output,
                    status="failed",
                    duration_ms=result.duration_ms,
                    metadata={
                        **dict(result.metadata),
                        "routing_error": f"{type(exc).__name__}: {exc}",
                    },
                )
                self._clear_paused_run(event_context)
                break
            if next_edge is None:
                graph_logger.info(f"Граф '{self.name}' завершен (нет доступных переходов из '{node.name}').")
                break

            graph_logger.info(f"Переход: {next_edge.source} --({next_edge.label or ''})--> {next_edge.target}")
            transition = self._make_event(
                ctx,
                event_context,
                event="edge_selected",
                event_type="edge",
                node_name=node.name,
                route_source=node.name,
                parent_event_id=node_finished.event_id,
                route_target=next_edge.target,
                edge_label=next_edge.label,
            )
            events.append(transition)
            ui_event_payloads.append(_serialize_graph_ui_event(transition))
            await self._dispatch_event(ctx, transition)
            current = next_edge.target

        duration_ms = (time.monotonic() - started) * 1000
        graph_finished = self._make_event(
            ctx,
            event_context,
            event="graph_completed",
            event_type="graph",
            parent_event_id=graph_started.event_id,
            status=status,
            duration_ms=duration_ms,
            output=_serialize_output(last_output),
        )
        events.append(graph_finished)
        ui_event_payloads.append(_serialize_graph_ui_event(graph_finished))
        await self._dispatch_event(ctx, graph_finished)
        if status != "paused":
            self._clear_paused_run(event_context)
        return GraphExecutionResult(
            graph_name=self.name,
            status=status,
            last_output=last_output,
            outputs=outputs,
            events=events,
            ui_event_payloads=ui_event_payloads,
            routing_table=self.routing_table(),
            duration_ms=duration_ms,
        )

    async def _dispatch_event(self, ctx: PipelineContext, event: GraphNodeEvent | dict[str, Any]) -> None:
        callback = ctx.metadata.get("_graph_event_callback")
        if callback is None:
            return
        
        # If it's a dict, it's a raw event from a sub-step, needs serializing
        # unless the callback already handles dicts. 
        # In our case, the callback from stream_ui_events handles GraphNodeEvent objects.
        # But we can make it polymorphic.
        
        result = callback(event)
        if asyncio.iscoroutine(result):
            await result

    def _resolve_node_timeout(self, node: Node) -> float | None:
        raw_timeout = node.metadata.get("timeout")
        if raw_timeout is None:
            return None
        try:
            timeout = float(raw_timeout)
        except (TypeError, ValueError):
            return None
        return timeout if timeout > 0 else None

    def _resolve_node_retry_policy(self, node: Node) -> RetryPolicy | None:
        retry_data = node.metadata.get("retry")
        if retry_data is None:
            return None
        if isinstance(retry_data, RetryPolicy):
            return retry_data
        if not isinstance(retry_data, dict):
            return None
        return RetryPolicy(
            max_attempts=int(retry_data.get("max_attempts", 3) or 3),
            base_delay=float(retry_data.get("base_delay", 0.25) or 0.25),
            backoff=str(retry_data.get("backoff", "exponential") or "exponential"),
        )

    def _retry_delay_for_attempt(self, retry: RetryPolicy, attempt: int) -> float:
        if retry.backoff == "fixed":
            return retry.base_delay
        return retry.base_delay * (2 ** (attempt - 1))

    async def _execute_node_step(self, node: Node, ctx: PipelineContext) -> StepResult:
        timeout = self._resolve_node_timeout(node)
        retry = self._resolve_node_retry_policy(node)
        max_attempts = retry.max_attempts if retry is not None else 1
        last_result: StepResult | None = None

        for attempt in range(1, max_attempts + 1):
            started = time.monotonic()
            try:
                execution = node.step.execute(ctx)
                if timeout is not None:
                    result = await asyncio.wait_for(execution, timeout=timeout)
                else:
                    result = await execution
            except Exception as exc:
                duration_ms = (time.monotonic() - started) * 1000
                if isinstance(exc, asyncio.TimeoutError):
                    error = f"TimeoutError: graph node '{node.name}' exceeded timeout={timeout}s"
                else:
                    error = f"{type(exc).__name__}: {exc}"
                result = StepResult(
                    output=None,
                    status="failed",
                    duration_ms=duration_ms,
                    metadata={"error": error},
                )

            metadata = dict(result.metadata)
            metadata["attempt"] = attempt
            metadata["max_attempts"] = max_attempts
            result = StepResult(
                output=result.output,
                status=result.status,
                duration_ms=result.duration_ms,
                metadata=metadata,
            )
            last_result = result
            if result.status != "failed" or attempt >= max_attempts:
                return result
            if retry is not None:
                await asyncio.sleep(self._retry_delay_for_attempt(retry, attempt))

        return last_result or StepResult(
            output=None,
            status="failed",
            duration_ms=0,
            metadata={"error": f"Graph node '{node.name}' failed without a result"},
        )

    def _select_next_edge(self, source: str, ctx: PipelineContext) -> Edge | None:
        fallback: Edge | None = None
        for edge in self._edges:
            if edge.source != source:
                continue
            if edge.condition is None:
                if fallback is None:
                    fallback = edge
                continue
            if edge.condition(ctx):
                return edge
        return fallback

    def _validate_graph(self) -> None:
        if not self._nodes:
            raise ValueError("Graph must contain at least one node")
        if self._entrypoint is None:
            raise ValueError("Graph requires an entrypoint")
        if self._entrypoint not in self._nodes:
            raise ValueError(f"Graph entrypoint '{self._entrypoint}' is not a defined node")
        unconditional_edges: dict[str, int] = {}
        for edge in self._edges:
            if edge.source not in self._nodes:
                raise ValueError(f"Graph edge source '{edge.source}' is not a defined node")
            if edge.target not in self._nodes:
                raise ValueError(f"Graph edge target '{edge.target}' is not a defined node")
            if edge.condition is None:
                unconditional_edges[edge.source] = unconditional_edges.get(edge.source, 0) + 1
                if unconditional_edges[edge.source] > 1:
                    raise ValueError(
                        f"Graph node '{edge.source}' defines multiple unconditional fallback edges"
                    )

    def _wrap_target(self, target: GraphTarget, *, name: str | None = None) -> Step:
        if isinstance(target, Step):
            if name and target.name != name:
                target._name = name
            return target
        if isinstance(target, Runnable):
            return AgentStep(name or self._auto_name("agent"), target)
        if callable(target):
            return FunctionStep(name or getattr(target, "__name__", self._auto_name("node")), target)
        raise TypeError(f"Cannot wrap {type(target)} as Graph node")

    def _auto_name(self, prefix: str) -> str:
        name = f"{prefix}_{self._auto_counter}"
        self._auto_counter += 1
        return name

    def _ensure_event_context(self, ctx: PipelineContext) -> dict[str, Any]:
        key = f"adk_graph_event_context:{self.name}"
        existing = ctx.metadata.get(key)
        if isinstance(existing, dict):
            return existing
        trace_id = ctx.metadata.get("trace_id")
        session_id = ctx.metadata.get("session_id")
        graph_run_id = uuid4().hex
        correlation_id = (
            ctx.metadata.get("correlation_id")
            or ctx.metadata.get("_graph_correlation_id")
            or trace_id
            or graph_run_id
        )
        event_context = {
            "graph_run_id": graph_run_id,
            "pipeline_id": ctx.pipeline_id,
            "trace_id": trace_id,
            "session_id": session_id,
            "correlation_id": correlation_id,
            "sequence": 0,
        }
        ctx.metadata[key] = event_context
        return event_context

    def _load_paused_run(self, correlation_id: str) -> _PausedGraphRun | None:
        paused_run = self._paused_runs.get(correlation_id)
        if paused_run is not None:
            return paused_run
        if self._checkpoint_store is None:
            return None
        record = self._checkpoint_store.load_checkpoint(self.name, correlation_id)
        if record is None:
            return None
        paused_run = self._deserialize_paused_run(record.payload)
        self._paused_runs[correlation_id] = paused_run
        return paused_run

    def _serialize_paused_run(
        self,
        paused_run: _PausedGraphRun,
    ) -> _PausedGraphRunCheckpointPayload:
        payload: _PausedGraphRunCheckpointPayload = {
            "correlation_id": paused_run.correlation_id,
            "pipeline_id": paused_run.pipeline_id,
            "paused_node_name": paused_run.paused_node_name,
            "resume_from": paused_run.resume_from,
            "context_data": dict(paused_run.context_data),
            "metadata_data": dict(paused_run.metadata_data),
            "blackboard_data": dict(paused_run.blackboard_data),
            "blackboard_versions_data": dict(paused_run.blackboard_versions_data),
        }
        return cast(
            _PausedGraphRunCheckpointPayload,
            ensure_json_object(payload, path="graph_checkpoint"),
        )

    def _deserialize_paused_run(
        self,
        payload: _PausedGraphRunCheckpointPayload | JSONObject,
    ) -> _PausedGraphRun:
        checkpoint_payload = cast(
            _PausedGraphRunCheckpointPayload,
            ensure_json_object(payload, path="graph_checkpoint"),
        )
        return _PausedGraphRun(
            correlation_id=checkpoint_payload["correlation_id"],
            pipeline_id=checkpoint_payload["pipeline_id"],
            paused_node_name=checkpoint_payload["paused_node_name"],
            resume_from=checkpoint_payload["resume_from"],
            context_data=dict(checkpoint_payload.get("context_data", {})),
            metadata_data=dict(checkpoint_payload.get("metadata_data", {})),
            blackboard_data=dict(checkpoint_payload.get("blackboard_data", {})),
            blackboard_versions_data={
                str(key): int(value)
                for key, value in dict(checkpoint_payload.get("blackboard_versions_data", {})).items()
            },
        )

    def _store_paused_run(
        self,
        ctx: PipelineContext,
        event_context: dict[str, Any],
        *,
        node_name: str,
        result: StepResult,
    ) -> None:
        correlation_id = event_context.get("correlation_id")
        if not isinstance(correlation_id, str) or not correlation_id:
            return
        paused_node_name = node_name
        paused_steps = result.metadata.get("paused_steps")
        if isinstance(paused_steps, list) and paused_steps:
            last_paused = paused_steps[-1]
            if isinstance(last_paused, dict) and isinstance(last_paused.get("node_name"), str):
                paused_node_name = last_paused["node_name"]
        metadata = dict(ctx.metadata)
        metadata.pop("_graph_event_callback", None)
        metadata.pop("_executing_graph", None)
        paused_run = _PausedGraphRun(
            correlation_id=correlation_id,
            pipeline_id=ctx.pipeline_id,
            paused_node_name=paused_node_name,
            resume_from=node_name,
            context_data=dict(ctx.data),
            metadata_data=metadata,
            blackboard_data=dict(ctx.blackboard),
            blackboard_versions_data=ctx.blackboard.versions_snapshot(),
        )
        self._paused_runs[correlation_id] = paused_run
        if self._checkpoint_store is not None:
            now = datetime.now(UTC).isoformat()
            self._checkpoint_store.save_checkpoint(
                GraphCheckpointRecord(
                    correlation_id=correlation_id,
                    graph_name=self.name,
                    status="paused",
                    created_at=now,
                    updated_at=now,
                    payload=self._serialize_paused_run(paused_run),
                )
            )

    def _clear_paused_run(self, event_context: dict[str, Any]) -> None:
        correlation_id = event_context.get("correlation_id")
        if isinstance(correlation_id, str) and correlation_id:
            self._clear_paused_run_by_correlation_id(correlation_id)

    def _clear_paused_run_by_correlation_id(self, correlation_id: str) -> None:
        self._paused_runs.pop(correlation_id, None)
        if self._checkpoint_store is not None:
            self._checkpoint_store.delete_checkpoint(self.name, correlation_id)

    def _make_event(
        self,
        ctx: PipelineContext,
        event_context: dict[str, Any],
        *,
        event: str,
        event_type: str,
        parent_event_id: str | None = None,
        node_name: str | None = None,
        route_source: str | None = None,
        route_target: str | None = None,
        edge_label: str | None = None,
        status: str | None = None,
        duration_ms: float | None = None,
        output: Any = None,
        metadata: dict[str, Any] | None = None,
    ) -> GraphNodeEvent:
        event_context["sequence"] = int(event_context.get("sequence", 0)) + 1
        meta = dict(metadata or {})
        meta["_blackboard"] = dict(ctx.blackboard)
        
        return GraphNodeEvent(
            event_id=uuid4().hex,
            event=event,
            event_type=event_type,
            graph_name=self.name,
            graph_run_id=str(event_context["graph_run_id"]),
            pipeline_id=ctx.pipeline_id,
            trace_id=event_context.get("trace_id"),
            session_id=event_context.get("session_id"),
            correlation_id=event_context.get("correlation_id"),
            sequence=int(event_context["sequence"]),
            parent_event_id=parent_event_id,
            node_name=node_name,
            route_source=route_source,
            route_target=route_target,
            edge_label=edge_label,
            status=status,
            duration_ms=duration_ms,
            output=output,
            metadata=meta,
        )


def _serialize_output(value: Any) -> Any:
    if isinstance(value, (str, int, float, bool, type(None), dict, list)):
        return value
    return str(value)


def _serialize_step_result(result: StepResult) -> dict[str, Any]:
    return {
        "output": _serialize_output(result.output),
        "status": result.status,
        "duration_ms": result.duration_ms,
        "metadata": result.metadata,
    }


def _serialize_graph_result(result: GraphExecutionResult) -> dict[str, Any]:
    return {
        "graph_name": result.graph_name,
        "status": result.status,
        "last_output": _serialize_output(result.last_output),
        "outputs": {
            name: _serialize_step_result(step_result)
            for name, step_result in result.outputs.items()
        },
        "node_events": [asdict(event) for event in result.events],
        "ui_event_payloads": list(result.ui_event_payloads),
        "routing_table": result.routing_table,
        "duration_ms": result.duration_ms,
    }


def _serialize_graph_ui_event(event: GraphNodeEvent) -> dict[str, Any]:
    return {
        "event_id": event.event_id,
        "parent_event_id": event.parent_event_id,
        "event": event.event,
        "event_type": event.event_type,
        "graph_name": event.graph_name,
        "graph_run_id": event.graph_run_id,
        "pipeline_id": event.pipeline_id,
        "trace_id": event.trace_id,
        "session_id": event.session_id,
        "correlation_id": event.correlation_id,
        "sequence": event.sequence,
        "timestamp": event.timestamp,
        "node_name": event.node_name,
        "route_source": event.route_source,
        "route_target": event.route_target,
        "edge_label": event.edge_label,
        "status": event.status,
        "duration_ms": event.duration_ms,
        "output": _serialize_output(event.output),
        "metadata": dict(event.metadata),
        "blackboard": event.metadata.get("_blackboard", {})
    }




__all__ = [
    "Edge",
    "Graph",
    "GraphExecutionResult",
    "GraphNodeEvent",
    "Node",
]
