"""LangGraph implementation integrated with RxPy (actor-first)."""

from __future__ import annotations

import asyncio
import collections
import concurrent.futures
import contextlib
import logging
from collections import defaultdict
from collections.abc import Callable
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field
from rx.core import Observer  # type: ignore[attr-defined]
from rx.scheduler.eventloop import AsyncIOScheduler  # type: ignore[attr-defined]
from rx.subject import Subject  # type: ignore[attr-defined]

from cleveragents.agents.base import Agent
from cleveragents.langgraph.nodes import Edge, Node, NodeConfig, NodeType
from cleveragents.langgraph.state import GraphState, StateManager
from cleveragents.reactive.stream_router import (
    ReactiveStreamRouter,
    StreamConfig,
    StreamMessage,
    StreamType,
)

_OnNextCallable = Callable[[Any], None]
_OnErrorCallable = Callable[[Exception], None]

MAX_EXECUTION_HISTORY: int = 1000

# Default timeout (seconds) for blocking on executor futures.
_DEFAULT_EXECUTOR_TIMEOUT: float = 300.0


class GraphConfig(BaseModel):  # pylint: disable=too-many-instance-attributes
    name: str
    nodes: dict[str, NodeConfig] = Field(default_factory=dict)
    edges: list[Edge] = Field(default_factory=list)
    entry_point: str = "start"
    state_class: type | None = None
    checkpointing: bool = False
    checkpoint_dir: Path | None = None
    enable_time_travel: bool = False
    parallel_execution: bool = True
    metadata: dict[str, Any] = Field(default_factory=dict)


class LangGraph:  # pylint: disable=too-many-instance-attributes
    def __init__(
        self,
        config: GraphConfig,
        agents: dict[str, Agent] | None = None,
        stream_router: ReactiveStreamRouter | None = None,
        scheduler: AsyncIOScheduler | None = None,
    ):
        self.config = config
        self.name = config.name
        self.agents = agents or {}
        self.logger = logging.getLogger(__name__)

        if not scheduler:
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
            scheduler = AsyncIOScheduler(loop)

        self.scheduler = scheduler
        self.stream_router = stream_router or ReactiveStreamRouter(self.scheduler)
        self.nodes: dict[str, Node] = {}
        self._initialize_nodes()

        state_class = config.state_class or GraphState
        self.state_manager = StateManager(
            initial_state=state_class(),
            checkpoint_dir=config.checkpoint_dir if config.checkpointing else None,
            enable_time_travel=config.enable_time_travel,
        )

        self._node_executors: dict[str, Callable[[StreamMessage], StreamMessage]] = {}
        # Default: min(32, cpu_count+4) — suitable for I/O-bound node executors
        self._executor_pool: concurrent.futures.ThreadPoolExecutor = (
            concurrent.futures.ThreadPoolExecutor()
        )
        self.execution_history: collections.deque[str] = collections.deque(
            maxlen=MAX_EXECUTION_HISTORY
        )
        # Disposables returned by observable.subscribe(); disposed in stop().
        self._subscriptions: list[Any] = []
        self.is_running = False
        self._create_graph_streams()
        self._analyze_graph()

    def __del__(self) -> None:
        with contextlib.suppress(Exception):
            self._executor_pool.shutdown(wait=False)

    async def execute(self, input_data: GraphState | dict[str, Any]) -> GraphState:
        if not self.is_running:
            raise RuntimeError(
                f"Cannot execute graph {self.name!r}: graph is not running"
            )
        state = (
            input_data
            if isinstance(input_data, GraphState)
            else GraphState.from_dict(input_data)
        )
        # Use replace_state() so the is_closed guard is enforced and state stream
        # subscribers are notified through the proper StateManager API, preventing
        # silent state corruption after StateManager.close() has been called.
        self.state_manager.replace_state(state)
        start_stream = f"__{self.name}_node_start__"
        if start_stream in self.stream_router.streams:
            self.stream_router.send_message(start_stream, state.model_copy(deep=True))
        else:
            raise ValueError("Start stream not initialized")
        return self.state_manager.get_state()

    def get_execution_history(self) -> list[str]:
        return list(self.execution_history)

    def _topological_levels(self) -> dict[int, set[str]]:
        # Simple level-order grouping using in-degree; excludes start/end guards
        in_degree = {node: 0 for node in self.nodes if node not in {"start", "end"}}
        for edge in self.config.edges:
            if edge.target in in_degree:
                in_degree[edge.target] = in_degree.get(edge.target, 0) + 1
        levels: dict[int, set[str]] = {}
        current_level = 0
        current_nodes = {n for n, deg in in_degree.items() if deg == 0}
        while current_nodes:
            levels[current_level] = set(current_nodes)
            next_nodes: set[str] = set()
            for node in current_nodes:
                for neighbor in self.adjacency_list.get(node, []):
                    if neighbor in in_degree:
                        in_degree[neighbor] -= 1
                        if in_degree[neighbor] == 0:
                            next_nodes.add(neighbor)
            current_nodes = next_nodes
            current_level += 1
        return levels

    def _initialize_nodes(self) -> None:
        if "start" not in self.config.nodes:
            self.config.nodes["start"] = NodeConfig(name="start", type=NodeType.START)
        if "end" not in self.config.nodes:
            self.config.nodes["end"] = NodeConfig(name="end", type=NodeType.END)
        for node_name, node_config in self.config.nodes.items():
            self.nodes[node_name] = Node(node_config, self.agents)

    def _create_graph_streams(self) -> None:
        for node_name in self.nodes:
            stream_name = f"__{self.name}_node_{node_name}__"
            self._register_node_executor(node_name)
            config = StreamConfig(name=stream_name, type=StreamType.COLD)
            self.stream_router.create_stream(config)
        self._create_control_streams()
        self._setup_node_stream_subscriptions()

    def _create_control_streams(self) -> None:
        control_stream_name = f"__{self.name}_control__"
        control_stream = Subject()
        self.stream_router.streams[control_stream_name] = control_stream
        self.stream_router.observables[control_stream_name] = control_stream

        state_stream_name = f"__{self.name}_state__"
        self.stream_router.streams[state_stream_name] = self.state_manager.state_stream
        self.stream_router.observables[state_stream_name] = (
            self.state_manager.state_stream
        )

    def _make_on_next_handler(self, name: str) -> _OnNextCallable:
        """Create an ``on_next`` callback for the node named *name*.

        The returned closure looks up the executor from
        ``self._node_executors`` at invocation time, wraps non-StreamMessage
        inputs, and catches executor exceptions so they do not propagate
        through the RxPy Subject call chain.
        """

        def on_next(msg: Any) -> None:
            executor = self._node_executors.get(name)
            if executor is not None:
                if not isinstance(msg, StreamMessage):
                    msg = StreamMessage(content=msg, metadata={})
                try:
                    # TODO(#10799): route executor return value to successor
                    # node streams based on self.adjacency_list[name] to
                    # implement downstream propagation.
                    executor(msg)
                except TimeoutError:
                    self.logger.warning("Executor timed out for node %s", name)
                except RuntimeError as exc:
                    if "stopping" in str(exc) or "not running" in str(exc):
                        self.logger.warning(
                            "Executor interrupted for node %s: %s", name, exc
                        )
                    else:
                        self.logger.exception("Executor failed for node %s", name)
                except Exception:
                    self.logger.exception("Executor failed for node %s", name)
            else:
                self.logger.warning("No executor found for node %s", name)

        return on_next

    def _make_on_error_handler(self, stream_name: str) -> _OnErrorCallable:
        """Create an ``on_error`` callback for *stream_name*."""

        def on_error(error: Exception) -> None:
            self.logger.error("Error in node stream %s: %s", stream_name, error)

        return on_error

    def _setup_node_stream_subscriptions(self) -> None:
        def on_completed() -> None:
            pass

        for node_name in self.nodes:
            stream_name = f"__{self.name}_node_{node_name}__"
            if stream_name not in self.stream_router.observables:
                self.logger.debug(
                    "No observable for stream %s; skipping subscription",
                    stream_name,
                )
                continue
            observable = self.stream_router.observables[stream_name]
            observer = Observer(
                on_next=self._make_on_next_handler(node_name),
                on_error=self._make_on_error_handler(stream_name),
                on_completed=on_completed,
            )
            disposable = observable.subscribe(observer)
            self._subscriptions.append(disposable)

    def _register_node_executor(self, node_name: str) -> None:
        node = self.nodes[node_name]

        async def async_executor(msg: StreamMessage) -> StreamMessage:
            state = self.state_manager.get_state()
            updates = await node.execute(state)
            self.state_manager.update_state(updates, node_id=node_name)
            # Thread-safe under CPython GIL; deque.append is atomic.
            self.execution_history.append(node_name)
            return StreamMessage(
                content=updates,
                metadata={
                    "node": node_name,
                    "execution_count": getattr(node, "execution_count", 0),
                    "graph": self.name,
                },
            )

        def sync_executor(msg: StreamMessage) -> StreamMessage:
            if not self.is_running:
                raise RuntimeError(
                    f"Cannot execute node {node_name!r}: graph is not running"
                )
            # Access the scheduler's internal event loop.  This relies on
            # RxPY's ``AsyncIOScheduler`` exposing ``_loop`` (private API;
            # verified: RxPY 4.x).  If the attribute is absent the code
            # falls back to the thread-pool path gracefully.
            scheduler_loop: asyncio.AbstractEventLoop | None = getattr(
                self.scheduler, "_loop", None
            )

            # Determine whether we are already running inside the
            # scheduler's event loop thread.  Blocking with
            # ``future.result()`` from the same thread that owns the loop
            # would deadlock, so we only use ``run_coroutine_threadsafe``
            # when called from a *different* thread.
            try:
                running_loop = asyncio.get_running_loop()
            except RuntimeError:
                running_loop = None

            timeout: float = (
                node.config.timeout
                if node.config.timeout is not None
                else _DEFAULT_EXECUTOR_TIMEOUT
            )

            if (
                scheduler_loop is not None
                and scheduler_loop.is_running()
                and running_loop is not scheduler_loop
            ):
                future = asyncio.run_coroutine_threadsafe(
                    async_executor(msg), scheduler_loop
                )
                try:
                    return future.result(timeout=timeout)
                except concurrent.futures.TimeoutError as exc:
                    # Cancellation is best-effort: a timed-out task may still
                    # complete and update state if the coroutine has already
                    # started and does not cooperatively check for cancellation.
                    future.cancel()
                    raise TimeoutError(
                        f"Node {node_name!r} executor timed out after {timeout}s"
                    ) from exc
                except concurrent.futures.CancelledError as exc:
                    raise RuntimeError(
                        f"Node {node_name!r} execution was cancelled (graph stopping)"
                    ) from exc

            def run_async() -> StreamMessage:
                return asyncio.run(async_executor(msg))

            try:
                future_tp = self._executor_pool.submit(run_async)
            except RuntimeError as exc:
                raise RuntimeError(
                    f"Cannot execute node {node_name!r}: graph is not running"
                ) from exc
            try:
                return future_tp.result(timeout=timeout)
            except concurrent.futures.TimeoutError as exc:
                # Cancellation is best-effort: ``cancel()`` only succeeds if
                # the task has not yet started running.  A timed-out task may
                # still complete and update state.
                if not future_tp.cancel():
                    self.logger.warning(
                        "Timed-out task for node %s could not be cancelled; "
                        "thread slot remains occupied until completion",
                        node_name,
                    )
                raise TimeoutError(
                    f"Node {node_name!r} executor timed out after {timeout}s"
                ) from exc
            except concurrent.futures.CancelledError as exc:
                raise RuntimeError(
                    f"Node {node_name!r} execution was cancelled (graph stopping)"
                ) from exc

        self._node_executors[node_name] = sync_executor

    def _analyze_graph(self) -> None:
        self.adjacency_list: dict[str, list[str]] = defaultdict(list)
        self.reverse_adjacency_list: dict[str, list[str]] = defaultdict(list)
        for edge in self.config.edges:
            source = edge.source if edge.source != "END" else "end"
            target = edge.target if edge.target != "END" else "end"
            self.adjacency_list[source].append(target)
            self.reverse_adjacency_list[target].append(source)
        self.has_cycles = self._detect_cycles()
        self.parallel_groups = self._find_parallel_groups()
        self._validate_graph()

    def _detect_cycles(self) -> bool:
        visited = set()
        rec_stack = set()

        def has_cycle(node: str) -> bool:
            visited.add(node)
            rec_stack.add(node)
            for neighbor in self.adjacency_list[node]:
                if neighbor not in visited:
                    if has_cycle(neighbor):
                        return True
                elif neighbor in rec_stack:
                    return True
            rec_stack.remove(node)
            return False

        return any(node not in visited and has_cycle(node) for node in self.nodes)

    def _find_parallel_groups(self) -> list[set[str]]:
        if not self.config.parallel_execution:
            return []
        in_degree = {node: 0 for node in self.nodes}
        for edge in self.config.edges:
            in_degree[edge.target] = in_degree.get(edge.target, 0) + 1

        groups: list[set[str]] = []
        current_group = {node for node, degree in in_degree.items() if degree == 0}

        while current_group:
            groups.append(current_group)
            next_group: set[str] = set()
            for node in current_group:
                for neighbor in self.adjacency_list.get(node, []):
                    in_degree[neighbor] -= 1
                    if in_degree[neighbor] == 0:
                        next_group.add(neighbor)
            current_group = next_group

        return groups

    def _validate_graph(self) -> None:
        if "start" not in self.nodes or "end" not in self.nodes:
            raise ValueError("Graph must have start and end nodes")

    def start(self, initial_input: Any = None) -> None:
        """Start (or restart) the graph.

        Shuts down the previous executor pool before creating a new one
        to prevent leaking threads and to ensure in-flight tasks from a
        prior cycle cannot corrupt state in the new cycle.

        .. warning::

            ``shutdown(wait=True, cancel_futures=True)`` blocks until all
            **already-running** tasks complete.  ``cancel_futures`` only
            cancels *queued-but-not-started* tasks.  If a running task is
            stuck in blocking I/O without a timeout, this method will hang
            indefinitely.  Callers that need bounded startup latency should
            ensure all submitted callables respect a timeout or cooperative
            cancellation.
        """
        # Prevent new submissions while the pool is being swapped.
        self.is_running = False
        self._executor_pool.shutdown(wait=True, cancel_futures=True)
        self._executor_pool = concurrent.futures.ThreadPoolExecutor()
        self.is_running = True
        start_stream = f"__{self.name}_node_start__"
        if start_stream not in self.stream_router.streams:
            raise ValueError("Start stream not initialized")
        if initial_input is None:
            initial_input = GraphState()
        self.stream_router.send_message(start_stream, initial_input)

    def stop(self) -> None:
        self.is_running = False
        for disposable in self._subscriptions:
            with contextlib.suppress(Exception):
                disposable.dispose()
        self._subscriptions.clear()
        self._executor_pool.shutdown(wait=True, cancel_futures=True)

    def get_state(self) -> GraphState:
        return self.state_manager.get_state()
