"""Basic governance enforcement with BeaconLangGraphGovernanceHandler.

Demonstrates how to add governance policy checks to LangChain tool and LLM
calls. The handler intercepts callbacks before/after execution and calls the
Beacon governance API to decide whether each action is allowed.

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

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

from lumenova_beacon import BeaconClient, BeaconLangGraphHandler
from lumenova_beacon.exceptions import GovernanceViolationError
from lumenova_beacon.governance import BeaconLangGraphGovernanceHandler

import dotenv

dotenv.load_dotenv()

# Initialize BeaconClient (required for transport/auth)
client = BeaconClient()

# Tracing handler (existing - traces all operations to the Beacon dashboard)
tracer = BeaconLangGraphHandler()

# Governance handler (new - enforces policies before/after execution)
governance = BeaconLangGraphGovernanceHandler(
    stack_ids=['basic-chain-001'],
    fail_open=True,  # Allow if governance API is unreachable
)

# Use both handlers together in the callbacks list
callbacks = [tracer, governance]


# === Example 1: Simple LLM call with governance ===
print('--- Example 1: LLM call with governance ---')

llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)

try:
    response = llm.invoke(
        'What is the capital of France?',
        config={'callbacks': callbacks},
    )
    print(f'Response: {response.content}')
except GovernanceViolationError as e:
    print(f'Blocked by governance: {e}')
    print(f'  Policies: {e.policies}')
    print(f'  Latency: {e.latency_ms:.1f}ms')


# === Example 2: Chain with governance ===
print('\n--- Example 2: Chain with governance ---')

prompt = ChatPromptTemplate.from_messages([
    ('system', 'You are a helpful assistant.'),
    ('user', '{question}'),
])

chain = prompt | llm | StrOutputParser()

try:
    result = chain.invoke(
        {'question': 'Explain quantum computing in one sentence.'},
        config={'callbacks': callbacks},
    )
    print(f'Result: {result}')
except GovernanceViolationError as e:
    print(f'Blocked by governance: {e}')


# === Example 3: Selective hooks (pre-execution only) ===
print('\n--- Example 3: Pre-execution checks only ---')

# Only check governance BEFORE tool/LLM execution (skip post-execution).
# This reduces latency by cutting governance API calls in half.
governance_pre_only = BeaconLangGraphGovernanceHandler(
    stack_ids=['basic-chain-001'],
    enabled_hooks={'pre_tool', 'pre_llm'},
)

try:
    response = llm.invoke(
        'What is 2 + 2?',
        config={'callbacks': [tracer, governance_pre_only]},
    )
    print(f'Response: {response.content}')
except GovernanceViolationError as e:
    print(f'Blocked by governance: {e}')


# Check your Beacon dashboard to see:
# - Traces with governance evaluation latency
# - Blocked actions recorded as error spans by the tracing handler
# - Governance enforcement logs in the audit trail
