"""Governance enforcement on a LangGraph agent with tools.

Demonstrates how the governance handler blocks disallowed tool calls
before they execute. When a tool is blocked, LangChain propagates the
GovernanceViolationError through on_tool_error. LangGraph agents with
error handling can gracefully recover.

Requirements:
    pip install lumenova-beacon[langchain]
    pip install langchain-openai langgraph
"""

from typing import Annotated, Sequence, TypedDict

from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph, add_messages
from langgraph.prebuilt import ToolNode

from lumenova_beacon import BeaconClient
from lumenova_beacon.exceptions import GovernanceViolationError
from lumenova_beacon.governance import GovernanceConfig
from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphConfig

import dotenv

dotenv.load_dotenv()

# Initialize client
client = BeaconClient()


# --- Tools ---

@tool
def get_weather(location: str) -> str:
    """Get the current weather for a location."""
    weather_data = {
        'new york': 'Sunny, 72F',
        'london': 'Cloudy, 61F',
        'tokyo': 'Rainy, 68F',
    }
    return weather_data.get(location.lower(), f'No data for {location}')


@tool
def execute_query(sql: str) -> str:
    """Execute a SQL query against the database."""
    # In a real app, this would run a query.
    # Governance policies might block certain queries (e.g., DROP, DELETE).
    return f'Query executed: {sql}'


@tool
def send_email(to: str, subject: str, body: str) -> str:
    """Send an email to a recipient."""
    # Governance policies might restrict who can be emailed.
    return f'Email sent to {to}: {subject}'


# --- Agent graph ---

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]


tools = [get_weather, execute_query, send_email]
tool_node = ToolNode(tools)

llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
llm_with_tools = llm.bind_tools(tools)


def call_model(state: AgentState):
    response = llm_with_tools.invoke(state['messages'])
    return {'messages': [response]}


def should_continue(state: AgentState):
    last_message = state['messages'][-1]
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        return 'tools'
    return END


workflow = StateGraph(AgentState)
workflow.add_node('agent', call_model)
workflow.add_node('tools', tool_node)
workflow.set_entry_point('agent')
workflow.add_conditional_edges('agent', should_continue, {'tools': 'tools', END: END})
workflow.add_edge('tools', 'agent')

app = workflow.compile()


# --- Beacon config with integrated governance ---

beacon = BeaconLangGraphConfig(
    graph=app,
    thread_id='governance-demo',
    agent_name='governed-agent',
    environment='development',
    governance=GovernanceConfig(
        stack_ids=['governed-agent-001'],
        fail_open=True,
        enabled_hooks={'pre_tool', 'post_tool', 'pre_llm', 'post_llm'},
    ),
)


# --- Execute ---

if __name__ == '__main__':
    # Query that triggers tool calls -- governance evaluates each one
    query = "What's the weather in London and then send an email to admin@example.com about it"
    print(f'Query: {query}')

    config = beacon.get_config()

    try:
        result = app.invoke(
            {'messages': [HumanMessage(content=query)]},
            config=config,
        )
        print(f'Response: {result["messages"][-1].content}')
    except GovernanceViolationError as e:
        # If the agent doesn't handle tool errors internally, the violation
        # propagates here. Inspect which policy blocked the action.
        print(f'\nBlocked by governance policy:')
        print(f'  Message:  {e}')
        print(f'  Result:   {e.result}')
        print(f'  Policies: {e.policies}')
        print(f'  Latency:  {e.latency_ms:.1f}ms')

    # Inspect governance state after execution
    if beacon.governance_handler:
        print(f'\nGovernance violations: {beacon.governance_handler.violation_count}')

# What happens when governance blocks a tool:
#
# 1. Agent decides to call "send_email"
# 2. BeaconLangGraphGovernanceHandler.on_tool_start fires BEFORE execution
# 3. Handler calls POST /api/v1/governance/evaluate with:
#    {"stack_ids": ["governed-agent-001"],
#     "input": {"phase": "pre", "type": "tool",
#               "tool": {"name": "send_email", "parameters": {...}}}}
# 4. OPA policy returns {"decision": "block", "message": "email not permitted"}
# 5. Handler raises GovernanceViolationError
# 6. LangChain does NOT execute the tool
# 7. BeaconLangGraphHandler.on_tool_error records an error span
# 8. The error propagates to the agent which can retry or respond
