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

from __future__ import annotations

import asyncio
import concurrent.futures
import logging
from collections import defaultdict
from pathlib import Path
from typing import Any, cast

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,
)


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.execution_history: list[str] = []
        self.is_running = False
        self._create_graph_streams()
        self._analyze_graph()

    async def execute(self, input_data: GraphState | dict[str, Any]) -> GraphState:
        state = (
            input_data
            if isinstance(input_data, GraphState)
            else GraphState.from_dict(cast(dict[str, Any], input_data))
        )
        # Replace state manager state for a fresh execution context
        self.state_manager.state = state
        self.state_manager.state_stream.on_next(state)
        start_stream = f"__{self.name}_node_start__"
        if start_stream in self.stream_router.streams:
            self.stream_router.send_message(start_stream, state)
        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}__"
            operators: list[dict[str, Any]] = [
                {"type": "map", "params": {"function": f"execute_node_{node_name}"}}
            ]
            self._register_node_executor(node_name)
            config = StreamConfig(
                name=stream_name, type=StreamType.COLD, operators=operators
            )
            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 _setup_node_stream_subscriptions(self) -> None:
        for node_name in self.nodes:
            stream_name = f"__{self.name}_node_{node_name}__"
            if stream_name in self.stream_router.observables:
                observable = self.stream_router.observables[stream_name]

                def on_next(_msg: Any) -> None:
                    pass

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

                def on_completed() -> None:
                    pass

                observer = Observer(
                    on_next=on_next, on_error=on_error, on_completed=on_completed
                )
                observable.subscribe(observer)

    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)
            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:
            def run_async() -> Any:
                return asyncio.run(async_executor(msg))

            with concurrent.futures.ThreadPoolExecutor() as executor:
                future = executor.submit(run_async)
                return future.result()

        setattr(self.stream_router, f"_builtin_execute_node_{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:
        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

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