"""Vorim SDK — LangChain / LangGraph Python integration.

Provides tool wrapping with permission checks + audit trail,
a callback handler for observability, and a factory for creating
Vorim-secured agents.

Requires: ``pip install vorim[langchain]``

Usage::

    from vorim import Vorim
    from vorim.integrations.langchain import vorim_tool, VorimCallbackHandler

    vorim = Vorim(api_key="agid_sk_live_...")

    @vorim_tool(vorim, agent_id="agid_acme_...", permission="agent:execute")
    def search(query: str) -> str:
        \"\"\"Search documents.\"\"\"
        return f"Results for {query}"
"""

from __future__ import annotations

import time
from functools import wraps
from typing import Any, Callable
from uuid import UUID

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.tools import BaseTool, StructuredTool
except ImportError as e:
    raise ImportError(
        "LangChain integration requires langchain-core. "
        "Install with: pip install vorim[langchain]"
    ) from e

from ..client import AsyncVorim, Vorim
from ..types import AuditEventInput, PermissionScope


# ── Tool Decorator ────────────────────────────────────────────────────────


def vorim_tool(
    vorim: Vorim,
    agent_id: str,
    permission: PermissionScope = "agent:execute",
) -> Callable:
    """Decorator that wraps a function-based LangChain tool with Vorim
    permission checks and audit logging.

    Usage::

        @vorim_tool(vorim, agent_id="agid_acme_...", permission="agent:read")
        def search_docs(query: str) -> str:
            \"\"\"Search internal documents.\"\"\"
            return do_search(query)

        # search_docs is now a LangChain StructuredTool with Vorim enforcement
    """

    def decorator(func: Callable) -> StructuredTool:
        tool_name = func.__name__

        @wraps(func)
        def guarded(*args: Any, **kwargs: Any) -> Any:
            # 1. Permission check
            result = vorim.check(agent_id, permission)
            if not result.allowed:
                vorim.emit(
                    agent_id=agent_id,
                    event_type="tool_call",
                    action=tool_name,
                    result="denied",
                    permission=permission,
                    metadata={"reason": result.reason, "framework": "langchain"},
                )
                raise PermissionError(
                    f"Vorim: permission denied for '{tool_name}' — "
                    f"scope '{permission}'"
                    f"{f': {result.reason}' if result.reason else ''}"
                )

            # 2. Execute
            start = time.monotonic()
            try:
                output = func(*args, **kwargs)
            except Exception as exc:
                vorim.emit(
                    agent_id=agent_id,
                    event_type="tool_call",
                    action=tool_name,
                    result="error",
                    permission=permission,
                    latency_ms=int((time.monotonic() - start) * 1000),
                    error_code=type(exc).__name__,
                    metadata={"error": str(exc), "framework": "langchain"},
                )
                raise

            # 3. Audit success
            vorim.emit(
                agent_id=agent_id,
                event_type="tool_call",
                action=tool_name,
                result="success",
                permission=permission,
                latency_ms=int((time.monotonic() - start) * 1000),
                metadata={"framework": "langchain"},
            )

            return output

        return StructuredTool.from_function(
            func=guarded,
            name=tool_name,
            description=func.__doc__ or f"{tool_name} tool",
        )

    return decorator


# ── Tool Wrapper ──────────────────────────────────────────────────────────


def wrap_tool(
    tool: BaseTool,
    vorim: Vorim,
    agent_id: str,
    permission: PermissionScope = "agent:execute",
) -> BaseTool:
    """Wrap an existing LangChain ``BaseTool`` with Vorim permission checks
    and audit logging.

    Returns the same tool instance with its ``_run`` method patched.

    Usage::

        from langchain_community.tools import DuckDuckGoSearchRun
        search = DuckDuckGoSearchRun()
        guarded_search = wrap_tool(search, vorim, agent_id, "agent:read")
    """
    original_run = tool._run

    def patched_run(*args: Any, **kwargs: Any) -> Any:
        result = vorim.check(agent_id, permission)
        if not result.allowed:
            vorim.emit(
                agent_id=agent_id,
                event_type="tool_call",
                action=tool.name,
                result="denied",
                permission=permission,
                metadata={"reason": result.reason, "framework": "langchain"},
            )
            raise PermissionError(
                f"Vorim: permission denied for '{tool.name}' — scope '{permission}'"
            )

        start = time.monotonic()
        try:
            output = original_run(*args, **kwargs)
        except Exception as exc:
            vorim.emit(
                agent_id=agent_id,
                event_type="tool_call",
                action=tool.name,
                result="error",
                permission=permission,
                latency_ms=int((time.monotonic() - start) * 1000),
                error_code=type(exc).__name__,
                metadata={"error": str(exc), "framework": "langchain"},
            )
            raise

        vorim.emit(
            agent_id=agent_id,
            event_type="tool_call",
            action=tool.name,
            result="success",
            permission=permission,
            latency_ms=int((time.monotonic() - start) * 1000),
            metadata={"framework": "langchain"},
        )
        return output

    tool._run = patched_run  # type: ignore[assignment]
    return tool


def wrap_tools(
    tools: list[BaseTool],
    vorim: Vorim,
    agent_id: str,
    permission_map: dict[str, PermissionScope] | None = None,
    default_permission: PermissionScope = "agent:execute",
) -> list[BaseTool]:
    """Wrap a list of tools with Vorim permission checks."""
    pm = permission_map or {}
    return [
        wrap_tool(t, vorim, agent_id, pm.get(t.name, default_permission))
        for t in tools
    ]


# ── Callback Handler ─────────────────────────────────────────────────────


class VorimCallbackHandler(BaseCallbackHandler):
    """LangChain callback handler that emits Vorim audit events for every
    tool invocation. Non-intrusive observability layer.

    Usage::

        handler = VorimCallbackHandler(vorim, agent_id)
        agent.invoke({"input": "..."}, config={"callbacks": [handler]})
    """

    def __init__(self, vorim: Vorim, agent_id: str) -> None:
        super().__init__()
        self._vorim = vorim
        self._agent_id = agent_id
        self._runs: dict[UUID, dict[str, Any]] = {}

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        inputs: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        tool_name = serialized.get("name", "unknown")
        self._runs[run_id] = {
            "tool": tool_name,
            "input": input_str,
            "start": time.monotonic(),
        }

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(run_id, None)
        if not run:
            return
        try:
            self._vorim.emit(
                agent_id=self._agent_id,
                event_type="tool_call",
                action=run["tool"],
                resource=_truncate(run["input"], 500),
                result="success",
                latency_ms=int((time.monotonic() - run["start"]) * 1000),
                metadata={"framework": "langchain"},
            )
        except Exception:
            pass  # never throw from callback

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(run_id, None)
        if not run:
            return
        try:
            self._vorim.emit(
                agent_id=self._agent_id,
                event_type="tool_call",
                action=run["tool"],
                resource=_truncate(run["input"], 500),
                result="error",
                latency_ms=int((time.monotonic() - run["start"]) * 1000),
                error_code=type(error).__name__,
                metadata={"error": str(error), "framework": "langchain"},
            )
        except Exception:
            pass


# ── Agent Factory ─────────────────────────────────────────────────────────


def create_vorim_agent(
    vorim: Vorim,
    *,
    name: str,
    capabilities: list[str],
    scopes: list[PermissionScope],
    tools: list[BaseTool],
    description: str | None = None,
    permission_map: dict[str, PermissionScope] | None = None,
    default_permission: PermissionScope = "agent:execute",
) -> dict[str, Any]:
    """Register a Vorim agent and wrap LangChain tools.

    Returns a dict with ``agent_id``, ``registration``, ``tools``,
    ``callback_handler``, and ``private_key``.

    Usage::

        from langchain_openai import ChatOpenAI
        from langgraph.prebuilt import create_react_agent

        result = create_vorim_agent(
            vorim,
            name="research-agent",
            capabilities=["search"],
            scopes=["agent:read", "agent:execute"],
            tools=[search_tool],
        )

        agent = create_react_agent(
            ChatOpenAI(model="gpt-4o"),
            result["tools"],
        )
        agent.invoke(
            {"messages": [("user", "Find AI papers")]},
            config={"callbacks": [result["callback_handler"]]},
        )
    """
    reg = vorim.register(
        name=name,
        description=description,
        capabilities=capabilities,
        scopes=scopes,
    )
    agent_id = reg.agent.agent_id

    wrapped = wrap_tools(
        tools, vorim, agent_id, permission_map, default_permission,
    )
    handler = VorimCallbackHandler(vorim, agent_id)

    return {
        "agent_id": agent_id,
        "registration": reg,
        "tools": wrapped,
        "callback_handler": handler,
        "private_key": reg.private_key,
    }


# ── Internal ──────────────────────────────────────────────────────────────


def _truncate(s: str, max_len: int) -> str:
    return s[:max_len] + "…" if len(s) > max_len else s
