"""
LangChain Integration for Agendex Governance.

Provides a simplified integration surface for LangChain agents with
automatic reasoning capture from agent actions.

Usage:
    from agendex.langchain import governed_agent

    # Line 1: Create governed context
    client, callback = governed_agent(task="monthly_report")

    # Lines 2-4: Get governed clients (as needed)
    db = client.db()
    s3 = client.s3()
    http = client.http()

    # Line 5: Attach callback to executor
    executor = AgentExecutor(agent, tools, callbacks=[callback])

    # Tool code - 0 extra lines for reasoning
    def fetch_revenue(month: str):
        return db.query("monthly_revenue", {"month": month})  # Reasoning auto-injected
"""
from __future__ import annotations

from typing import Any, Callable, Dict, Optional, Tuple

from .client import AgendexClient
from .context import clear_reasoning, set_reasoning
from .crm import GovernedCRMClient
from .db import GovernedDBClient
from .http import GovernedHTTPSession
from .s3 import GovernedS3Client


class GovernedContext:
    """
    Pre-configured governance context with auto-reasoning.

    Provides factory methods for creating governed clients that share
    the same AgendexClient and task context.

    Example:
        client, callback = governed_agent(task="monthly_report")
        db = client.db()
        s3 = client.s3()
    """

    def __init__(
        self,
        client: AgendexClient,
        task: str,
        on_event: Optional[Callable[[Dict], None]] = None,
    ):
        """
        Initialize governed context.

        Args:
            client: The AgendexClient instance
            task: Default task context for all operations
            on_event: Optional callback for governance events
        """
        self._client = client
        self._task = task
        self._on_event = on_event

    @property
    def task(self) -> str:
        """Current default task context."""
        return self._task

    @property
    def client(self) -> AgendexClient:
        """Underlying AgendexClient instance."""
        return self._client

    def db(
        self,
        task: Optional[str] = None,
        action: Optional[str] = None,
        on_event: Optional[Callable[[Dict], None]] = None,
    ) -> GovernedDBClient:
        """
        Create a governed database client.

        Args:
            task: Override task context (default: context task)
            action: Override action name
            on_event: Override event callback (default: context callback)
        """
        return GovernedDBClient(
            self._client,
            task=task or self._task,
            action=action,
            on_event=on_event or self._on_event,
        )

    def s3(
        self,
        task: Optional[str] = None,
        actions: Optional[Dict[str, str]] = None,
        on_event: Optional[Callable[[Dict], None]] = None,
    ) -> GovernedS3Client:
        """
        Create a governed S3 client.

        Args:
            task: Override task context (default: context task)
            actions: Override action names
            on_event: Override event callback (default: context callback)
        """
        return GovernedS3Client(
            self._client,
            task=task or self._task,
            actions=actions,
            on_event=on_event or self._on_event,
        )

    def http(
        self,
        task: Optional[str] = None,
        action: Optional[str] = None,
        on_event: Optional[Callable[[Dict], None]] = None,
    ) -> GovernedHTTPSession:
        """
        Create a governed HTTP session.

        Args:
            task: Override task context (default: context task)
            action: Override action name
            on_event: Override event callback (default: context callback)
        """
        return GovernedHTTPSession(
            self._client,
            task=task or self._task,
            action=action,
            on_event=on_event or self._on_event,
        )

    def crm(
        self,
        task: Optional[str] = None,
        agent_role: Optional[str] = None,
        actions: Optional[Dict[str, str]] = None,
        on_event: Optional[Callable[[Dict], None]] = None,
    ) -> GovernedCRMClient:
        """
        Create a governed CRM client.

        Args:
            task: Override task context (default: context task)
            agent_role: Optional role identifier for multi-agent scenarios
            actions: Override action names
            on_event: Override event callback (default: context callback)
        """
        return GovernedCRMClient(
            self._client,
            task=task or self._task,
            agent_role=agent_role,
            actions=actions,
            on_event=on_event or self._on_event,
        )


class AgendexReasoningCallback:
    """
    LangChain callback handler that captures reasoning into contextvars.

    When an agent action is executed, this callback extracts the agent's
    reasoning (from action.log) and stores it in a contextvar. Governed
    clients automatically read from this contextvar, so reasoning flows
    through without any extra code in tool implementations.

    Usage:
        client, callback = governed_agent(task="monthly_report")
        executor = AgentExecutor(agent, tools, callbacks=[callback])
    """

    def __init__(self, on_event: Optional[Callable[[Dict], None]] = None):
        """
        Initialize the reasoning callback.

        Args:
            on_event: Optional callback for reasoning events
        """
        self.on_event = on_event

    def on_agent_action(self, action: Any, **kwargs: Any) -> None:
        """
        Called when the agent decides on an action.

        Extracts reasoning from the action log and stores it in contextvar.
        """
        if hasattr(action, "log") and action.log:
            set_reasoning(action.log)
            if self.on_event:
                self.on_event({"type": "reasoning", "content": action.log})

    def on_tool_end(self, output: Any, **kwargs: Any) -> None:
        """
        Called when a tool finishes execution.

        Clears reasoning to prevent it from leaking to subsequent calls.
        """
        clear_reasoning()

    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
        """
        Called when a tool raises an error.

        Clears reasoning even on error to maintain clean state.
        """
        clear_reasoning()

    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """
        Called when a chain finishes.

        Clears reasoning at chain end for safety.
        """
        clear_reasoning()


def governed_agent(
    task: str,
    base_url: Optional[str] = None,
    agent_id: Optional[str] = None,
    token: Optional[str] = None,
    on_event: Optional[Callable[[Dict], None]] = None,
) -> Tuple[GovernedContext, AgendexReasoningCallback]:
    """
    Create a governed agent context with auto-reasoning capture.

    This is the main entry point for LangChain integration. It returns
    a GovernedContext for creating governed clients and a callback for
    attaching to the AgentExecutor.

    Args:
        task: Default task context for all operations
        base_url: Override Agendex URL (default: from env)
        agent_id: Override agent ID (default: from env)
        token: Override auth token (default: from env)
        on_event: Optional callback for governance events

    Returns:
        Tuple of (GovernedContext, AgendexReasoningCallback)

    Example:
        client, callback = governed_agent(task="monthly_report")
        db = client.db()
        executor = AgentExecutor(agent, tools, callbacks=[callback])

        def fetch_revenue(month: str):
            return db.query("monthly_revenue", {"month": month})
    """
    agendex = AgendexClient(
        base_url=base_url,
        agent_id=agent_id,
        token=token,
    )
    context = GovernedContext(agendex, task, on_event=on_event)
    callback = AgendexReasoningCallback(on_event=on_event)
    return context, callback


__all__ = [
    "GovernedContext",
    "AgendexReasoningCallback",
    "governed_agent",
]
