"""Advanced governance configuration: custom result parsers and fail-closed mode.

Demonstrates how to customize the governance handler for different OPA policy
output shapes and strict enforcement requirements.

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

from typing import Annotated, Any, 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 BeaconLangGraphHandler, BeaconClient
from lumenova_beacon.exceptions import GovernanceError, GovernanceViolationError
from lumenova_beacon.governance import BeaconLangGraphGovernanceHandler

import dotenv

dotenv.load_dotenv()

client = BeaconClient()


# --- Custom result parsers ---

def strict_parser(result: dict[str, Any]) -> bool:
    """Parser for OPA policies that use a custom nested decision format.

    Expected OPA output shape:
        {"decision": "block", "message": "reason", "reasons": [...]}

    This demonstrates how to write a custom parser when your OPA policies
    return additional fields beyond the standard decision/message format.
    """
    decision = result.get('decision', 'allow')
    allowed = decision == 'allow'
    if not allowed:
        reasons = result.get('reasons', [])
        if reasons:
            print(f'  Denial reasons: {reasons}')
    return allowed


def confidence_parser(result: dict[str, Any]) -> bool:
    """Parser for OPA policies that return a confidence score.

    Expected OPA output shape:
        {"decision": "allow"/"block", "confidence": 0.0-1.0}

    Only allows if confidence exceeds threshold.
    """
    if result.get('decision') != 'allow':
        return False
    confidence = result.get('confidence', 0.0)
    threshold = 0.8
    return confidence >= threshold


# --- Tools ---

@tool
def search_database(query: str) -> str:
    """Search the internal database."""
    return f'Results for: {query}'


@tool
def delete_record(record_id: str) -> str:
    """Delete a record from the database."""
    return f'Deleted record {record_id}'


# --- Agent ---

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


tools = [search_database, delete_record]
tool_node = ToolNode(tools)
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
llm_with_tools = llm.bind_tools(tools)


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


def should_continue(state: AgentState):
    last = state['messages'][-1]
    if hasattr(last, 'tool_calls') and last.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()

tracer = BeaconLangGraphHandler(agent_name='strict-agent')


# === Example 1: Custom result parser ===
print('--- Example 1: Custom nested-decision parser ---')

governance_strict = BeaconLangGraphGovernanceHandler(
    stack_ids=['strict-agent-001'],
    result_parser=strict_parser,
    fail_open=False,  # Block if governance API is unreachable
)

try:
    result = app.invoke(
        {'messages': [HumanMessage(content='Delete record ABC-123')]},
        config={'callbacks': [tracer, governance_strict]},
    )
    print(f'Response: {result["messages"][-1].content}')
except GovernanceViolationError as e:
    print(f'Blocked: {e}')
except GovernanceError as e:
    # With fail_open=False, API errors also block execution
    print(f'Governance API error (action blocked): {e}')


# === Example 2: Confidence-based parser ===
print('\n--- Example 2: Confidence-based parser ---')

governance_confidence = BeaconLangGraphGovernanceHandler(
    stack_ids=['strict-agent-001'],
    result_parser=confidence_parser,
    enabled_hooks={'pre_tool'},  # Only check before tool calls
)

try:
    result = app.invoke(
        {'messages': [HumanMessage(content='Search for active users')]},
        config={'callbacks': [tracer, governance_confidence]},
    )
    print(f'Response: {result["messages"][-1].content}')
except GovernanceViolationError as e:
    print(f'Low confidence - blocked: {e}')


# === Example 3: Fail-closed with custom timeout ===
print('\n--- Example 3: Fail-closed with tight timeout ---')

governance_failclosed = BeaconLangGraphGovernanceHandler(
    stack_ids=['strict-agent-001'],
    fail_open=False,  # Block if API unreachable
    timeout=2.0,      # Tight timeout - fail fast
    enabled_hooks={'pre_tool'},
)

try:
    result = app.invoke(
        {'messages': [HumanMessage(content='Search for user records')]},
        config={'callbacks': [tracer, governance_failclosed]},
    )
    print(f'Response: {result["messages"][-1].content}')
except GovernanceViolationError as e:
    print(f'Policy violation: {e}')
except GovernanceError as e:
    print(f'Governance unavailable, action blocked (fail-closed): {e}')


# Summary of configuration patterns:
#
# | Pattern           | fail_open | result_parser    | Use case                        |
# |-------------------|-----------|------------------|---------------------------------|
# | Permissive        | True      | default          | Development, non-critical paths |
# | Strict            | False     | default          | Production, audit-required      |
# | Custom OPA shape  | any       | custom function  | Non-standard Rego output        |
# | Minimal latency   | True      | default          | enabled_hooks={'pre_tool'}      |
