"""
LockStock LangGraph Middleware
-------------------------------
Middleware layer for LangGraph that enforces LockStock authorization.

Uses the thin client pattern - all verification happens server-side via /v1/stamp.

Usage:
    from langgraph.graph import StateGraph
    from lockstock_langgraph import lockstock_middleware, lockstock_node_wrapper

    # Option 1: Wrap entire graph
    graph = StateGraph(AgentState)
    # ... add nodes ...
    app = lockstock_middleware(graph.compile(), agent_id="agent_abc123", api_key="lsk_...")

    # Option 2: Wrap individual nodes
    @lockstock_node_wrapper(agent_id="agent_abc123", api_key="lsk_...", capability="DEPLOY")
    def deploy_node(state):
        # This node requires DEPLOY capability
        ...
"""

from functools import wraps
from typing import Any, Callable, Dict, Optional, TypeVar
import asyncio

from lockstock_core import LockStockClient
from lockstock_core.types import VerifyStatus


T = TypeVar("T")


class LockStockMiddleware:
    """
    Middleware that wraps a LangGraph compiled graph.

    Intercepts node transitions and verifies authorization via /v1/stamp.
    """

    def __init__(
        self,
        app: Any,
        agent_id: str,
        api_key: str,
        endpoint: str = "https://lockstock-api-i9kp.onrender.com",
        node_capability_map: Optional[Dict[str, str]] = None
    ):
        """
        Initialize LockStock middleware.

        Args:
            app: The compiled LangGraph application
            agent_id: The agent's unique identifier
            api_key: API key for authentication (from dashboard)
            endpoint: LockStock API endpoint
            node_capability_map: Mapping of node names to required capabilities
        """
        self.app = app
        self.client = LockStockClient(
            agent_id=agent_id,
            api_key=api_key,
            endpoint=endpoint
        )
        self.agent_id = agent_id
        self.node_capability_map = node_capability_map or {}

    async def invoke(
        self,
        input_state: Dict[str, Any],
        config: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Invoke the graph with LockStock authorization.

        Args:
            input_state: Initial state
            config: Optional configuration

        Returns:
            Final state after graph execution
        """
        # Stamp the graph invocation
        result = await self.client.stamp("graph_invoke")
        if not result.authorized:
            raise PermissionError(
                f"LockStock DENIED: graph_invoke. Reason: {result.reason}"
            )

        try:
            # Get the original invoke method
            if asyncio.iscoroutinefunction(self.app.invoke):
                result = await self.app.invoke(input_state, config)
            else:
                result = self.app.invoke(input_state, config)

            return result

        except Exception as e:
            raise

    async def astream(
        self,
        input_state: Dict[str, Any],
        config: Optional[Dict[str, Any]] = None
    ):
        """
        Stream graph execution with LockStock authorization.

        Args:
            input_state: Initial state
            config: Optional configuration

        Yields:
            State updates from graph execution
        """
        # Stamp the stream start
        stamp_result = await self.client.stamp("graph_stream")
        if not stamp_result.authorized:
            raise PermissionError(
                f"LockStock DENIED: graph_stream. Reason: {stamp_result.reason}"
            )

        try:
            async for update in self.app.astream(input_state, config):
                # Log each state transition
                if isinstance(update, dict):
                    node_name = update.get("__node__", "unknown")

                    # Check if this node requires authorization
                    if node_name in self.node_capability_map:
                        capability = self.node_capability_map[node_name]
                        result = await self.client.stamp(task=capability)

                        if not result.authorized:
                            raise PermissionError(
                                f"LockStock DENIED: Node '{node_name}' requires "
                                f"capability '{capability}'"
                            )

                yield update

        except Exception as e:
            raise

    def __getattr__(self, name: str) -> Any:
        """Proxy other attributes to the wrapped app."""
        return getattr(self.app, name)


def lockstock_middleware(
    app: Any,
    agent_id: str,
    api_key: str,
    endpoint: str = "https://lockstock-api-i9kp.onrender.com",
    node_capability_map: Optional[Dict[str, str]] = None
) -> LockStockMiddleware:
    """
    Wrap a LangGraph compiled app with LockStock middleware.

    Args:
        app: The compiled LangGraph application
        agent_id: The agent's unique identifier
        api_key: API key for authentication (from dashboard)
        endpoint: LockStock API endpoint
        node_capability_map: Mapping of node names to required capabilities

    Returns:
        Wrapped application with LockStock authorization
    """
    return LockStockMiddleware(
        app=app,
        agent_id=agent_id,
        api_key=api_key,
        endpoint=endpoint,
        node_capability_map=node_capability_map
    )


def lockstock_node_wrapper(
    agent_id: str,
    capability: str,
    api_key: str,
    endpoint: str = "https://lockstock-api-i9kp.onrender.com"
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """
    Decorator to wrap a LangGraph node with LockStock authorization.

    Args:
        agent_id: The agent's unique identifier
        capability: Required capability for this node
        api_key: API key for authentication (from dashboard)
        endpoint: LockStock API endpoint

    Returns:
        Decorator function

    Example:
        @lockstock_node_wrapper(agent_id="agent_abc", api_key="lsk_...", capability="DEPLOY")
        def deploy_node(state):
            # This requires DEPLOY capability
            return {"deployed": True}
    """
    def decorator(node_func: Callable[..., T]) -> Callable[..., T]:
        # Create client (will be shared across calls)
        client = LockStockClient(
            agent_id=agent_id,
            api_key=api_key,
            endpoint=endpoint
        )

        @wraps(node_func)
        async def async_wrapper(*args, **kwargs) -> T:
            # Stamp capability
            result = await client.stamp(task=capability)

            if not result.authorized:
                raise PermissionError(
                    f"LockStock DENIED: Node '{node_func.__name__}' "
                    f"requires capability '{capability}'. "
                    f"Reason: {result.reason}"
                )

            # Call the original function
            if asyncio.iscoroutinefunction(node_func):
                return await node_func(*args, **kwargs)
            else:
                return node_func(*args, **kwargs)

        @wraps(node_func)
        def sync_wrapper(*args, **kwargs) -> T:
            # For sync nodes, run async verification in event loop
            loop = asyncio.get_event_loop()

            async def verify_and_run():
                result = await client.stamp(task=capability)

                if not result.authorized:
                    raise PermissionError(
                        f"LockStock DENIED: Node '{node_func.__name__}' "
                        f"requires capability '{capability}'. "
                        f"Reason: {result.reason}"
                    )

                return node_func(*args, **kwargs)

            return loop.run_until_complete(verify_and_run())

        # Return appropriate wrapper based on original function type
        if asyncio.iscoroutinefunction(node_func):
            return async_wrapper
        else:
            return sync_wrapper

    return decorator
