"""LangChain/LangGraph governance callback handler.

This module provides a callback handler that enforces governance policies
on LangChain agent actions by calling the Beacon governance evaluation
endpoint before and/or after tool and LLM invocations.

Usage:
    from lumenova_beacon import BeaconClient, BeaconLangGraphHandler
    from lumenova_beacon.governance import BeaconLangGraphGovernanceHandler

    client = BeaconClient()

    tracer = BeaconLangGraphHandler()
    governance = BeaconLangGraphGovernanceHandler(stack_ids=['my-stack-id'])

    llm = ChatOpenAI()
    response = llm.invoke("Hello", config={"callbacks": [tracer, governance]})

Input format sent to the governance API (matches Rego policy expectations)::

    # Tool pre-execution:
    {"phase": "pre", "type": "tool",
     "source": {"agent_id": "my-agent", "framework": "langchain"},
     "context": {"session_id": "sess-1", "environment": "production"},
     "tool": {"name": "send_email", "parameters": {"to": "x@y.com"}}}

    # Tool post-execution:
    {"phase": "post", "type": "tool",
     "source": {"agent_id": "my-agent"},
     "tool": {"name": "read_file", "parameters": {"path": "/data/f.csv"},
              "output": "file contents..."}}

    # LLM pre-execution:
    {"phase": "pre", "type": "llm",
     "source": {"agent_id": "my-agent"},
     "context": {"model_call_count": 2},
     "llm": {"model_name": "gpt-4", "last_user_content": "..."}}

    # LLM post-execution:
    {"phase": "post", "type": "llm",
     "context": {"model_call_count": 2},
     "llm": {"model_name": "gpt-4"}}

OPA/Rego policies return ``{"decision": "allow"}`` or
``{"decision": "block", "message": "..."}``.
"""

from __future__ import annotations

import json
import logging

from collections import deque
from collections.abc import Callable
from typing import Any

from lumenova_beacon.exceptions import GovernanceViolationError
from lumenova_beacon.governance._helpers import (
    default_result_parser,
    evaluate_or_fail_open,
    get_sdk_version,
    is_allowed,
    truncate,
)
from lumenova_beacon.governance.client import GovernanceClient
from lumenova_beacon.tracing.trace import get_current_span, get_current_trace_id


try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.messages import BaseMessage
    from langchain_core.outputs import LLMResult
except ImportError:
    raise ImportError(
        'langchain-core is required for the governance callback handler. '
        'Install it with: pip install langchain-core'
    )

logger = logging.getLogger(__name__)

_DEFAULT_HOOKS = frozenset({'pre_tool', 'post_tool', 'pre_llm', 'post_llm'})
_VALID_ON_VIOLATION = ('raise', 'stop')
_INVOCATION_CHAIN_MAX = 10


class BeaconLangGraphGovernanceHandler(BaseCallbackHandler):
    """LangChain callback handler that enforces governance policies.

    Intercepts tool and LLM lifecycle events, calls the Beacon governance
    evaluation endpoint (``POST /api/v1/governance/evaluate``), and raises
    :class:`~lumenova_beacon.exceptions.GovernanceViolationError` to block
    actions that violate governance policies.

    The handler follows the same pattern as
    :class:`~lumenova_beacon.tracing.integrations.langchain.BeaconLangGraphHandler`:
    constructor params allow per-handler overrides, but defaults for
    ``session_id`` and ``environment`` are pulled from the active
    :class:`~lumenova_beacon.core.client.BeaconClient` config.

    Args:
        stack_ids: Explicit policy stack IDs to evaluate against. At least
            one of ``stack_ids`` or ``tags`` must be provided.
        tags: Tag-based policy stack discovery.
        fail_open: When ``True`` (default), allow the action if the governance
            API is unreachable or returns an unexpected error.
        enabled_hooks: Set of hook names that should trigger evaluation.
            Valid values: ``'pre_tool'``, ``'post_tool'``, ``'pre_llm'``,
            ``'post_llm'``. Defaults to all four.
        result_parser: Optional callable ``(result_dict) -> bool`` that
            determines whether the response allows the action.
        timeout: Override the transport timeout (in seconds) for governance
            evaluation calls.
        debug: When ``True``, enable DEBUG-level logging for the governance
            handler and client.
        on_violation: Strategy when a policy blocks an action.
            ``'raise'`` (default) raises ``GovernanceViolationError``.
            ``'stop'`` sets a stopped flag and raises, preventing all
            subsequent tool/LLM calls from proceeding.
        max_violations: After this many violations, escalate to ``'stop'``
            regardless of the ``on_violation`` setting. ``None`` disables.
        agent_id: Identifier for the agent, included in ``input.source``.
        agent_name: Human-readable agent name, included in ``input.source``.
        session_id: Override session ID. Defaults to ``BeaconClient.config.session_id``.
        environment: Override environment. Defaults to ``BeaconClient.config.environment``.
        metadata: Custom metadata dict included in ``input.context.metadata``.
    """

    # Tell LangChain to propagate exceptions raised in callbacks.
    raise_error = True

    def __init__(
        self,
        *,
        stack_ids: list[str] | None = None,
        tags: list[str] | None = None,
        fail_open: bool = True,
        enabled_hooks: set[str] | None = None,
        result_parser: Callable[[dict[str, Any]], bool] | None = None,
        timeout: float | None = None,
        debug: bool = False,
        on_violation: str = 'raise',
        max_violations: int | None = None,
        agent_id: str | None = None,
        agent_name: str | None = None,
        session_id: str | None = None,
        environment: str | None = None,
        metadata: dict[str, Any] | None = None,
    ):
        super().__init__()

        if on_violation not in _VALID_ON_VIOLATION:
            raise ValueError(
                f'on_violation must be one of {_VALID_ON_VIOLATION!r}, got {on_violation!r}'
            )

        self._stack_ids = stack_ids
        self._tags = tags
        self._fail_open = fail_open
        self._enabled_hooks: set[str] = (
            set(enabled_hooks) if enabled_hooks is not None else set(_DEFAULT_HOOKS)
        )
        self._result_parser = result_parser
        self._timeout = timeout
        self._debug = debug
        self._on_violation = on_violation
        self._max_violations = max_violations

        # Source context
        self._agent_id = agent_id
        self._agent_name = agent_name

        # Context — defaults from BeaconClient (same pattern as BeaconLangGraphHandler)
        from lumenova_beacon.core.client import get_client

        client = get_client()
        self._session_id = session_id if session_id is not None else client.config.session_id
        self._environment = environment if environment is not None else client.config.environment
        self._metadata = metadata
        self._user_id = client.config.user_id
        self._user_email = client.config.user_email
        self._user_name = client.config.user_name

        self._client = GovernanceClient()
        # Minimal run tracking so post-execution hooks can reference the
        # tool name and parameters (not available in on_tool_end signature).
        self._runs: dict[str, dict[str, Any]] = {}
        # Track model call count across the handler lifetime (useful for
        # policies that limit the number of LLM round-trips).
        self._model_call_count: int = 0
        # Graceful stop state.
        self._violation_count: int = 0
        self._stopped: bool = False
        # Invocation chain — bounded list of recent actions for sequence-aware policies.
        self._invocation_chain: deque[dict[str, str]] = deque(maxlen=_INVOCATION_CHAIN_MAX)

        if self._debug:
            self._enable_debug_logging()

        logger.debug(
            'Initialized BeaconLangGraphGovernanceHandler: stack_ids=%s tags=%s fail_open=%s '
            'enabled_hooks=%s timeout=%s agent_id=%s agent_name=%s session_id=%s '
            'environment=%s on_violation=%s max_violations=%s',
            self._stack_ids,
            self._tags,
            self._fail_open,
            self._enabled_hooks,
            self._timeout,
            self._agent_id,
            self._agent_name,
            self._session_id,
            self._environment,
            self._on_violation,
            self._max_violations,
        )

    # ------------------------------------------------------------------
    # Public properties / methods
    # ------------------------------------------------------------------

    @property
    def violation_count(self) -> int:
        """Number of governance violations encountered so far."""
        return self._violation_count

    def reset(self) -> None:
        """Reset violation count and stopped state for handler reuse."""
        self._violation_count = 0
        self._stopped = False

    # ------------------------------------------------------------------
    # Source & context builders
    # ------------------------------------------------------------------

    def _build_source(self) -> dict[str, Any]:
        """Build the ``input.source`` dict for governance evaluation."""
        source: dict[str, Any] = {'framework': 'langchain'}
        if self._agent_id:
            source['agent_id'] = self._agent_id
        if self._agent_name:
            source['agent_name'] = self._agent_name
        source['sdk_version'] = get_sdk_version()
        return source

    def _build_context(self) -> dict[str, Any]:
        """Build the ``input.context`` dict for governance evaluation."""
        ctx: dict[str, Any] = {}
        if self._session_id:
            ctx['session_id'] = self._session_id
        if self._environment:
            ctx['environment'] = self._environment
        if self._user_id:
            ctx['user_id'] = self._user_id
        if self._user_email:
            ctx['user_email'] = self._user_email
        if self._user_name:
            ctx['user_name'] = self._user_name
        ctx['model_call_count'] = self._model_call_count
        if self._metadata:
            ctx['metadata'] = self._metadata
        if self._invocation_chain:
            ctx['invocation_chain'] = list(self._invocation_chain)
        return ctx

    def _build_extra(self, action_data: dict[str, Any]) -> dict[str, Any]:
        """Build the full extra dict with source, context, and action data."""
        extra: dict[str, Any] = dict(action_data)
        extra['source'] = self._build_source()
        extra['context'] = self._build_context()
        return extra

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _enable_debug_logging() -> None:
        """Configure governance loggers to emit DEBUG output to stderr."""
        governance_logger = logging.getLogger('lumenova_beacon.governance')
        if not governance_logger.handlers:
            handler = logging.StreamHandler()
            handler.setFormatter(logging.Formatter('[governance] %(levelname)s %(message)s'))
            governance_logger.addHandler(handler)
        governance_logger.setLevel(logging.DEBUG)

    @staticmethod
    def _default_result_parser(result: dict[str, Any]) -> bool:
        """Default result parser.

        Delegates to :func:`~lumenova_beacon.governance._helpers.default_result_parser`.
        """
        return default_result_parser(result)

    def _is_allowed(self, response: dict[str, Any]) -> bool:
        return is_allowed(response, self._result_parser)

    def _evaluate_or_fail_open(
        self,
        phase: str,
        type_: str,
        extra: dict[str, Any],
    ) -> dict[str, Any] | None:
        """Call the governance API; return the response or ``None`` on fail-open."""
        trace_id = get_current_trace_id()
        span = get_current_span()
        span_id = span.span_id if span is not None else None
        return evaluate_or_fail_open(
            self._client,
            phase,
            type_,
            self._stack_ids,
            self._tags,
            extra,
            self._timeout,
            self._fail_open,
            trace_id=trace_id,
            span_id=span_id,
        )

    def _enforce(
        self,
        response: dict[str, Any] | None,
        phase: str,
        context_msg: str,
    ) -> None:
        """Raise if the response blocks the action, respecting the violation strategy."""
        # If already stopped, block immediately without calling the API.
        if self._stopped:
            raise GovernanceViolationError(
                message='Agent stopped due to governance violations',
                result={},
                policies=[],
                latency_ms=0.0,
            )

        if response is None:
            logger.debug('Governance enforce skipped: no response (fail-open path)')
            return

        if not self._is_allowed(response):
            self._violation_count += 1
            rego_message = response.get('message', '')

            msg = f'Governance policy blocked {context_msg} at {phase}'
            if rego_message:
                msg = f'{msg}: {rego_message}'

            logger.debug(
                'Governance BLOCKED: %s at %s | decision=%s rules=%s',
                context_msg,
                phase,
                response.get('decision'),
                response.get('rules_evaluated'),
            )

            # Determine effective strategy.
            strategy = self._on_violation
            if self._max_violations is not None and self._violation_count >= self._max_violations:
                strategy = 'stop'

            if strategy == 'stop':
                self._stopped = True

            raise GovernanceViolationError(
                message=msg,
                result=response,
                policies=response.get('rules_evaluated', []),
                latency_ms=response.get('latency_ms', 0.0),
            )
        logger.debug('Governance ALLOWED: %s at %s', context_msg, phase)

    @staticmethod
    def _truncate(value: str) -> str:
        """Truncate a string for inclusion in evaluation payloads."""
        return truncate(value)

    @staticmethod
    def _get_parameters(input_str: str, kwargs: dict[str, Any]) -> dict[str, Any]:
        """Extract tool arguments as a dict from callback parameters.

        LangChain passes tool args in several ways depending on version
        and tool type.  Try ``kwargs['inputs']``, ``kwargs['tool_input']``,
        then fall back to JSON-parsing ``input_str``.
        """
        parameters = kwargs.get('inputs') or kwargs.get('tool_input')
        if isinstance(parameters, dict) and parameters:
            return parameters
        # Fall back to JSON-parsing the input_str positional argument.
        if isinstance(input_str, str) and input_str:
            try:
                parsed = json.loads(input_str)
                if isinstance(parsed, dict):
                    return parsed
            except (json.JSONDecodeError, TypeError):
                pass
        return {}

    @staticmethod
    def _extract_output_str(output: Any) -> str:
        """Extract a string from a tool output.

        LangChain tool outputs may be ``ToolMessage`` objects (with a
        ``.content`` attribute), plain strings, or arbitrary objects.
        """
        if hasattr(output, 'content'):
            return str(output.content)
        return str(output)

    @staticmethod
    def _extract_last_user_content(messages: list[list[BaseMessage]]) -> str:
        """Walk message batches in reverse to find the last human message."""
        if not messages:
            return ''
        for msg in reversed(messages[0]):
            if msg.type == 'human':
                content = msg.content
                return str(content) if content else ''
        return ''

    # ------------------------------------------------------------------
    # Tool callbacks
    # ------------------------------------------------------------------

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Pre-execution governance check for tool invocations."""
        # Block immediately if stopped.
        if self._stopped:
            raise GovernanceViolationError(
                message='Agent stopped due to governance violations',
                result={},
                policies=[],
                latency_ms=0.0,
            )

        run_id = kwargs.get('run_id')
        name = serialized.get('name') or kwargs.get('name', 'unknown')
        parameters = self._get_parameters(input_str, kwargs)

        if run_id is not None:
            self._runs[str(run_id)] = {'tool_name': name, 'parameters': parameters}
        logger.debug(
            'on_tool_start: tool=%s run_id=%s parameters=%s',
            name,
            run_id,
            parameters,
        )

        if 'pre_tool' not in self._enabled_hooks:
            logger.debug('on_tool_start: pre_tool hook disabled, skipping')
            return

        tool_data: dict[str, Any] = {
            'name': name,
            'parameters': parameters,
        }
        extra = self._build_extra({'tool': tool_data})
        resp = self._evaluate_or_fail_open('pre', 'tool', extra)
        self._enforce(resp, 'pre', f"tool '{name}'")

    def on_tool_end(
        self,
        output: Any,
        **kwargs: Any,
    ) -> None:
        """Post-execution governance check for tool invocations."""
        run_id = kwargs.get('run_id')
        run_data = self._runs.pop(str(run_id), {}) if run_id else {}
        tool_name = kwargs.get('name') or run_data.get('tool_name', 'unknown')
        output_str = self._extract_output_str(output)

        logger.debug(
            'on_tool_end: tool=%s run_id=%s output_len=%d',
            tool_name,
            run_id,
            len(output_str),
        )

        # Track in invocation chain
        self._invocation_chain.append({'action': tool_name, 'type': 'tool'})

        if 'post_tool' not in self._enabled_hooks:
            logger.debug('on_tool_end: post_tool hook disabled, skipping')
            return

        tool_data: dict[str, Any] = {
            'name': tool_name,
            'parameters': run_data.get('parameters', {}),
            'output': self._truncate(output_str),
        }
        extra = self._build_extra({'tool': tool_data})
        resp = self._evaluate_or_fail_open('post', 'tool', extra)
        self._enforce(resp, 'post', f"tool '{tool_name}'")

    def on_tool_error(
        self,
        error: BaseException,
        **kwargs: Any,
    ) -> None:
        """Clean up run tracking on tool error."""
        run_id = kwargs.get('run_id')
        logger.debug('on_tool_error: run_id=%s error=%s', run_id, error)
        if run_id is not None:
            self._runs.pop(str(run_id), None)

    # ------------------------------------------------------------------
    # LLM callbacks
    # ------------------------------------------------------------------

    def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        **kwargs: Any,
    ) -> None:
        """Pre-execution governance check for non-chat LLM calls."""
        if self._stopped:
            raise GovernanceViolationError(
                message='Agent stopped due to governance violations',
                result={},
                policies=[],
                latency_ms=0.0,
            )

        run_id = kwargs.get('run_id')
        model_name = self._extract_model_name(serialized, **kwargs)
        self._model_call_count += 1
        if run_id is not None:
            self._runs[str(run_id)] = {'model_name': model_name}
        logger.debug(
            'on_llm_start: model=%s run_id=%s call_count=%d',
            model_name,
            run_id,
            self._model_call_count,
        )

        if 'pre_llm' not in self._enabled_hooks:
            logger.debug('on_llm_start: pre_llm hook disabled, skipping')
            return

        last_prompt = prompts[-1] if prompts else ''
        llm_data: dict[str, Any] = {
            'model_name': model_name,
            'last_user_content': self._truncate(last_prompt),
        }
        extra = self._build_extra({'llm': llm_data})
        resp = self._evaluate_or_fail_open('pre', 'llm', extra)
        self._enforce(resp, 'pre', f"LLM '{model_name}'")

    def on_chat_model_start(
        self,
        serialized: dict[str, Any],
        messages: list[list[BaseMessage]],
        **kwargs: Any,
    ) -> None:
        """Pre-execution governance check for chat model calls."""
        if self._stopped:
            raise GovernanceViolationError(
                message='Agent stopped due to governance violations',
                result={},
                policies=[],
                latency_ms=0.0,
            )

        run_id = kwargs.get('run_id')
        model_name = self._extract_model_name(serialized, **kwargs)
        self._model_call_count += 1
        if run_id is not None:
            self._runs[str(run_id)] = {'model_name': model_name}

        last_user_content = self._extract_last_user_content(messages)
        logger.debug(
            'on_chat_model_start: model=%s call_count=%d run_id=%s last_user_content=%s',
            model_name,
            self._model_call_count,
            run_id,
            last_user_content[:200] if last_user_content else '(none)',
        )

        if 'pre_llm' not in self._enabled_hooks:
            logger.debug('on_chat_model_start: pre_llm hook disabled, skipping')
            return

        llm_data: dict[str, Any] = {
            'model_name': model_name,
            'last_user_content': self._truncate(last_user_content),
        }
        extra = self._build_extra({'llm': llm_data})
        resp = self._evaluate_or_fail_open('pre', 'llm', extra)
        self._enforce(resp, 'pre', f"LLM '{model_name}'")

    def on_llm_end(
        self,
        response: LLMResult,
        **kwargs: Any,
    ) -> None:
        """Post-execution governance check for LLM calls."""
        run_id = kwargs.get('run_id')
        run_data = self._runs.pop(str(run_id), {}) if run_id else {}

        # Resolve model name with fallbacks (mirrors observability handler).
        model_name = run_data.get('model_name')

        if not model_name or model_name == 'unknown':
            if response.llm_output and isinstance(response.llm_output, dict):
                model_name = response.llm_output.get('model_name') or response.llm_output.get(
                    'model'
                )

        if not model_name or model_name == 'unknown':
            if response.generations and response.generations[0]:
                gen = response.generations[0][0]
                if hasattr(gen, 'message') and hasattr(gen.message, 'response_metadata'):
                    meta = gen.message.response_metadata
                    if meta:
                        model_name = meta.get('model_name') or meta.get('model')

        if not model_name:
            model_name = 'unknown'

        logger.debug(
            'on_llm_end: model=%s run_id=%s call_count=%d',
            model_name,
            run_id,
            self._model_call_count,
        )

        # Track in invocation chain
        self._invocation_chain.append({'action': model_name, 'type': 'llm'})

        if 'post_llm' not in self._enabled_hooks:
            logger.debug('on_llm_end: post_llm hook disabled, skipping')
            return

        llm_data: dict[str, Any] = {
            'model_name': model_name,
        }

        # Include token usage if available.
        if response.llm_output and isinstance(response.llm_output, dict):
            usage = response.llm_output.get('token_usage') or {}
            if usage:
                llm_data['token_usage'] = {
                    'input_tokens': usage.get('prompt_tokens'),
                    'output_tokens': usage.get('completion_tokens'),
                    'total_tokens': usage.get('total_tokens'),
                }

        extra = self._build_extra({'llm': llm_data})
        resp = self._evaluate_or_fail_open('post', 'llm', extra)
        self._enforce(resp, 'post', f"LLM '{model_name}'")

    def on_llm_error(
        self,
        error: BaseException,
        **kwargs: Any,
    ) -> None:
        """Clean up run tracking on LLM error."""
        run_id = kwargs.get('run_id')
        logger.debug('on_llm_error: run_id=%s error=%s', run_id, error)
        if run_id is not None:
            self._runs.pop(str(run_id), None)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _extract_model_name(serialized: dict[str, Any], **kwargs: Any) -> str:
        """Extract the model name from serialized LangChain kwargs or invocation_params."""
        model_kwargs = serialized.get('kwargs', {})
        model_name = (
            model_kwargs.get('model_name')
            or model_kwargs.get('model')
            or model_kwargs.get('deployment_name')
        )
        if not model_name:
            invocation_params = kwargs.get('invocation_params', {})
            model_name = (
                invocation_params.get('model_name')
                or invocation_params.get('model')
                or invocation_params.get('azure_deployment')
            )
        return model_name or 'unknown'
