"""Bridge between RxPy streams and LangGraph execution (actor-first)."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import weakref
from collections.abc import Callable
from typing import Any, cast

import rx  # type: ignore
from rx import operators as ops  # type: ignore[attr-defined]

from cleveragents.langgraph.graph import GraphConfig, LangGraph
from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
from cleveragents.langgraph.state import GraphState
from cleveragents.reactive.stream_router import (
    ReactiveStreamRouter,
    StreamConfig,
    StreamMessage,
    StreamType,
)


class RxPyLangGraphBridge:
    """Allows LangGraph nodes as RxPy operators and streams triggering graphs."""

    def __init__(self, stream_router: ReactiveStreamRouter):
        self.stream_router = stream_router
        self.logger = logging.getLogger(__name__)
        self.graphs: dict[str, LangGraph] = {}
        self._active_tasks: set[asyncio.Task[Any]] = set()
        self.cancellation_reasons: weakref.WeakKeyDictionary[asyncio.Task[Any], str] = (
            weakref.WeakKeyDictionary()
        )
        self._register_langgraph_operators()

    def __del__(self) -> None:
        with contextlib.suppress(AttributeError):
            self.cleanup_tasks()

    def cleanup_tasks(self) -> None:
        """Cancel all active tasks immediately (sync, best-effort)."""
        for task in list(self._active_tasks):
            if not task.done():
                task.cancel()
        self._active_tasks.clear()

    async def cleanup_tasks_async(self, timeout: float = 5.0) -> None:
        """Cancel active tasks and await their completion within *timeout*.

        Tasks that do not finish within the deadline are re-cancelled and
        forcefully discarded.  Tasks added to ``_active_tasks`` during the
        await window are also cancelled to prevent silent loss.

        Args:
            timeout: Maximum seconds to wait for each task to finish.
        """
        # Snapshot and clear: any tasks added during the await are picked
        # up in a second pass below.
        tasks = list(self._active_tasks)
        if not tasks:
            return

        for task in tasks:
            if not task.done():
                task.cancel()

        done, pending = await asyncio.wait(tasks, timeout=timeout)

        for task in done:
            self.logger.info(
                "Task completed during cleanup: %s",
                task.get_name(),
            )

        for task in pending:
            task.cancel()
            self.logger.warning(
                "Task did not complete within %.1f s timeout: %s",
                timeout,
                task.get_name(),
            )

        # Cancel any tasks added while we were awaiting.
        late_tasks = self._active_tasks - set(tasks)
        for task in late_tasks:
            if not task.done():
                task.cancel()
            self.logger.warning(
                "Late task cancelled during cleanup: %s",
                task.get_name(),
            )

        self._active_tasks.clear()

    async def cancel_task_with_reason(
        self,
        task: asyncio.Task[Any],
        reason: str,
    ) -> None:
        """Cancel *task* and record the *reason* for the cancellation.

        Args:
            task: The asyncio task to cancel.
            reason: Human-readable explanation for the cancellation.

        Raises:
            ValueError: If *task* is ``None`` or *reason* is empty.
        """
        if task is None:
            raise ValueError("task must not be None")
        if not reason:
            raise ValueError("reason must be a non-empty string")
        self.cancellation_reasons[task] = reason
        task.cancel()
        self.logger.info(
            "Cancelled task %s with reason: %s",
            task.get_name(),
            reason,
        )

    def _run_async_safely(self, coro: Any) -> Any:
        return coro

    def _register_langgraph_operators(self) -> None:
        self._register_operator("graph_execute", self._create_graph_executor)
        self._register_operator("state_update", self._create_state_updater)
        self._register_operator("state_checkpoint", self._create_state_checkpointer)
        self._register_operator("langgraph_node", self._create_node_operator)
        self._register_operator("conditional_route", self._create_conditional_router)

    def _register_operator(self, name: str, factory_func: Callable[..., Any]) -> None:
        setattr(self, f"_operator_{name}", factory_func)

    def create_graph_from_config(self, config: dict[str, Any]) -> LangGraph:
        graph_config = GraphConfig(
            name=config.get("name", "default"),
            entry_point=config.get("entry_point", "start"),
            checkpointing=config.get("checkpointing", False),
            enable_time_travel=config.get("enable_time_travel", False),
            parallel_execution=config.get("parallel_execution", True),
        )
        for node_name, node_data in config.get("nodes", {}).items():
            metadata = node_data.get("metadata", {})
            node_type = NodeType(node_data.get("type", "function"))
            if node_type == NodeType.MESSAGE_ROUTER and "rules" in node_data:
                metadata["rules"] = node_data["rules"]
            node_config = NodeConfig(
                name=node_name,
                type=node_type,
                agent=node_data.get("agent"),
                function=node_data.get("function"),
                tools=node_data.get("tools", []),
                retry_policy=node_data.get("retry_policy"),
                timeout=node_data.get("timeout"),
                parallel=node_data.get("parallel", False),
                condition=node_data.get("condition"),
                subgraph=node_data.get("subgraph"),
                metadata=metadata,
            )
            graph_config.nodes[node_name] = node_config
        for edge_data in config.get("edges", []):
            edge = Edge(
                source=edge_data["source"],
                target=edge_data["target"],
                condition=edge_data.get("condition"),
                metadata=edge_data.get("metadata", {}),
            )
            graph_config.edges.append(edge)
        graph = LangGraph(
            config=graph_config,
            agents=self.stream_router.agents,
            stream_router=self.stream_router,
            scheduler=self.stream_router.scheduler,
        )
        self.graphs[graph.name] = graph
        return graph

    def create_graph_stream(self, graph_name: str) -> StreamConfig:
        if graph_name not in self.graphs:
            raise ValueError(f"Graph '{graph_name}' not found")
        return StreamConfig(
            name=f"graph_{graph_name}",
            type=StreamType.COLD,
            operators=[{"type": "graph_execute", "params": {"graph": graph_name}}],
            publications=["__output__"],
        )

    def _create_graph_executor(self, params: dict[str, Any]) -> Any:
        graph_name = params.get("graph")
        if not graph_name or graph_name not in self.graphs:
            raise ValueError(f"Invalid graph name: {graph_name}")
        graph = cast(Any, self.graphs[graph_name])
        execution_lock = asyncio.Lock()

        async def execute_graph(msg: StreamMessage) -> StreamMessage:
            base_content = msg.content
            if isinstance(base_content, str):
                input_data: dict[str, Any] = {
                    "messages": [{"role": "user", "content": base_content}]
                }
            elif isinstance(base_content, dict):
                input_data = {**base_content}
            else:
                input_data = {"content": base_content}
            if msg.metadata:
                input_data["metadata"] = msg.metadata
            async with execution_lock:
                final_state: GraphState = await graph.execute(input_data)
            if final_state.messages:
                result: Any = final_state.messages[-1].get("content", "")
            else:
                result = final_state.to_dict()
            return StreamMessage(
                content=result,
                metadata={
                    **msg.metadata,
                    "graph": graph_name,
                    "execution_history": graph.get_execution_history(),
                    "final_state": final_state.to_dict(),
                },
            )

        def create_future_task(msg: StreamMessage) -> asyncio.Future[StreamMessage]:
            loop = asyncio.get_event_loop()
            task: asyncio.Future[StreamMessage] = asyncio.ensure_future(
                execute_graph(msg), loop=loop
            )
            self._active_tasks.add(task)
            task.add_done_callback(
                lambda t: self._active_tasks.discard(cast(asyncio.Task[Any], t))
            )
            return task

        return ops.flat_map(lambda msg: rx.from_future(create_future_task(msg)))

    def _create_state_updater(self, params: dict[str, Any]) -> Any:
        graph_name = params.get("graph")
        if not graph_name or graph_name not in self.graphs:
            raise ValueError(f"Invalid graph name: {graph_name}")
        graph = self.graphs[graph_name]

        def update_state(msg: StreamMessage) -> StreamMessage:
            updates = (
                msg.content if isinstance(msg.content, dict) else {"data": msg.content}
            )
            graph.state_manager.update_state(updates)
            return msg.copy_with(
                metadata={**msg.metadata, "state_updated": True, "graph": graph_name}
            )

        return ops.map(update_state)

    def _create_state_checkpointer(self, params: dict[str, Any]) -> Any:
        graph_name = params.get("graph")
        if not graph_name or graph_name not in self.graphs:
            raise ValueError(f"Invalid graph name: {graph_name}")
        graph = self.graphs[graph_name]

        def checkpoint(msg: StreamMessage) -> StreamMessage:
            if graph.state_manager.checkpoint_dir:
                graph.state_manager._save_checkpoint()  # pylint: disable=protected-access
            return msg.copy_with(
                metadata={**msg.metadata, "checkpointed": True, "graph": graph_name}
            )

        return ops.map(checkpoint)

    def _create_node_operator(self, params: dict[str, Any]) -> Any:
        node_name = params.get("node")
        graph_name = params.get("graph")
        if not graph_name or graph_name not in self.graphs:
            raise ValueError(f"Invalid graph name: {graph_name}")
        graph = self.graphs[graph_name]
        if not node_name or node_name not in graph.nodes:
            raise ValueError(f"Invalid node name: {node_name}")
        node = graph.nodes[node_name]

        async def execute_node(msg: StreamMessage) -> StreamMessage:
            state = graph.state_manager.get_state()
            updates = await node.execute(state)
            graph.state_manager.update_state(updates, node_id=node_name)
            return msg.copy_with(
                metadata={
                    **msg.metadata,
                    "node": node_name,
                    "graph": graph_name,
                    "state": graph.state_manager.get_state().to_dict(),
                }
            )

        def create_future_task(msg: StreamMessage) -> Any:
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
            return loop.create_task(execute_node(msg))

        return ops.flat_map(lambda msg: rx.from_future(create_future_task(msg)))

    def _create_conditional_router(self, params: dict[str, Any]) -> Any:
        rules = params.get("rules", [])

        def route_message(msg: StreamMessage) -> StreamMessage:
            for rule in rules:
                condition = rule.get("condition")
                target = rule.get("target")
                if condition and target and self._evaluate_condition(msg, condition):
                    return StreamMessage(
                        content=msg.content,
                        metadata={**msg.metadata, "routed_to": target},
                    )
            return msg

        return ops.map(route_message)

    def _evaluate_condition(
        self, msg: StreamMessage, condition: dict[str, Any]
    ) -> bool:
        if "equals" in condition:
            return msg.content == condition["equals"]
        if "field" in condition and isinstance(msg.content, dict):
            field = condition["field"]
            return msg.content.get(field) == condition.get("value")
        return False

    def create_message_router_node(self, rules: list[dict[str, Any]]) -> NodeConfig:
        node_name = f"message_router_{len(rules)}"
        metadata = {"rules": rules}
        return NodeConfig(
            name=node_name, type=NodeType.MESSAGE_ROUTER, metadata=metadata
        )

    def create_conditional_edges(self, rules: list[dict[str, Any]]) -> list[Edge]:
        edges: list[Edge] = []
        for rule in rules:
            target = rule.get("target")
            if target:
                edges.append(
                    Edge(source="start", target=target, condition=rule.get("condition"))
                )
        edges.append(Edge(source="start", target="end"))
        return edges
