"""Tests for the BeaconLangGraphGovernanceHandler and GovernanceClient."""

from __future__ import annotations

from unittest.mock import MagicMock, patch
from uuid import uuid4

import httpx
import pytest

from lumenova_beacon.exceptions import (
    GovernanceError,
    GovernanceNotFoundError,
    GovernanceValidationError,
    GovernanceViolationError,
)
from lumenova_beacon.governance.client import GovernanceClient
from lumenova_beacon.governance.integrations.langchain import BeaconLangGraphGovernanceHandler


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def transport():
    """Mock HTTPTransport returned by get_transport()."""
    t = MagicMock()
    t.headers = {'Authorization': 'Bearer test-key', 'Content-Type': 'application/json'}
    t.timeout = 10.0
    t.verify = True
    return t


@pytest.fixture
def mock_beacon_client():
    """Mock BeaconClient returned by get_client() in handler/decorator init."""
    mock_client = MagicMock()
    mock_client.config.session_id = 'test-session'
    mock_client.config.environment = 'test'
    return mock_client


@pytest.fixture
def _patch_transport(transport, mock_beacon_client):
    """Patch get_transport, get_base_url, and get_client for all tests in this module."""
    with (
        patch(
            'lumenova_beacon.governance.client.get_transport',
            return_value=transport,
        ),
        patch(
            'lumenova_beacon.governance.client.get_base_url',
            return_value='https://api.test.com',
        ),
        patch(
            'lumenova_beacon.core.client.get_client',
            return_value=mock_beacon_client,
        ),
    ):
        yield


@pytest.fixture
def governance_client(_patch_transport):
    return GovernanceClient()


@pytest.fixture
def allow_response():
    """Governance allow response using the flat format."""
    return {
        'decision': 'allow',
        'message': None,
        'rules_evaluated': [{'name': 'default-allow'}],
        'latency_ms': 5.2,
    }


@pytest.fixture
def block_response():
    """Governance block response using the flat format."""
    return {
        'decision': 'block',
        'message': 'tool not permitted',
        'rules_evaluated': [{'name': 'block-dangerous-tools'}],
        'latency_ms': 3.1,
    }


@pytest.fixture
def run_id():
    return uuid4()


# ---------------------------------------------------------------------------
# GovernanceClient tests
# ---------------------------------------------------------------------------


class TestGovernanceClient:
    def test_evaluate_builds_correct_payload(self, governance_client):
        """Verify the request payload matches the backend schema with nested tool."""
        mock_resp = MagicMock(spec=httpx.Response)
        mock_resp.status_code = 200
        mock_resp.text = '{}'
        mock_resp.json.return_value = {
            'decision': 'allow', 'message': None, 'rules_evaluated': [], 'latency_ms': 1.0,
        }

        with patch(
            'lumenova_beacon.governance.client.httpx.post', return_value=mock_resp
        ) as mock_post:
            governance_client.evaluate(
                phase='pre',
                type_='tool',
                stack_ids=['stack-1'],
                extra={'tool': {'name': 'search', 'parameters': {'query': 'test'}}},
            )

            mock_post.assert_called_once()
            payload = mock_post.call_args.kwargs['json']
            assert payload['stack_ids'] == ['stack-1']
            assert payload['input']['phase'] == 'pre'
            assert payload['input']['type'] == 'tool'
            # Tool data is nested under input.tool
            assert payload['input']['tool']['name'] == 'search'
            assert payload['input']['tool']['parameters'] == {'query': 'test'}

    def test_evaluate_url(self, governance_client):
        """Verify the endpoint URL is correct."""
        mock_resp = MagicMock(spec=httpx.Response)
        mock_resp.status_code = 200
        mock_resp.text = '{}'
        mock_resp.json.return_value = {'decision': 'allow', 'message': None, 'rules_evaluated': [], 'latency_ms': 0}

        with patch(
            'lumenova_beacon.governance.client.httpx.post', return_value=mock_resp
        ) as mock_post:
            governance_client.evaluate(
                phase='pre', type_='tool', stack_ids=['test'],
            )
            url = mock_post.call_args.args[0]
            assert url == 'https://api.test.com/api/v1/governance/evaluate'

    def test_evaluate_missing_stack_ids_and_tags_raises(self, governance_client):
        """When both stack_ids and tags are None, a validation error is raised."""
        with pytest.raises(GovernanceValidationError, match='(?i)at least one'):
            governance_client.evaluate(
                phase='pre', type_='llm',
                stack_ids=None, tags=None,
            )

    def test_evaluate_with_tags_only(self, governance_client):
        """When only tags are provided, stack_ids should not appear in the payload."""
        mock_resp = MagicMock(spec=httpx.Response)
        mock_resp.status_code = 200
        mock_resp.text = '{}'
        mock_resp.json.return_value = {'decision': 'allow', 'message': None, 'rules_evaluated': [], 'latency_ms': 0}

        with patch(
            'lumenova_beacon.governance.client.httpx.post', return_value=mock_resp
        ) as mock_post:
            governance_client.evaluate(
                phase='pre', type_='tool', tags=['env:prod'],
            )
            payload = mock_post.call_args.kwargs['json']
            assert payload['tags'] == ['env:prod']
            assert 'stack_ids' not in payload

    def test_evaluate_with_both_stack_ids_and_tags(self, governance_client):
        """When both stack_ids and tags are provided, both appear in the payload."""
        mock_resp = MagicMock(spec=httpx.Response)
        mock_resp.status_code = 200
        mock_resp.text = '{}'
        mock_resp.json.return_value = {'decision': 'allow', 'message': None, 'rules_evaluated': [], 'latency_ms': 0}

        with patch(
            'lumenova_beacon.governance.client.httpx.post', return_value=mock_resp
        ) as mock_post:
            governance_client.evaluate(
                phase='pre', type_='tool',
                stack_ids=['s1'], tags=['t1'],
            )
            payload = mock_post.call_args.kwargs['json']
            assert payload['stack_ids'] == ['s1']
            assert payload['tags'] == ['t1']

    def test_evaluate_404_raises_not_found(self, governance_client):
        mock_resp = MagicMock(spec=httpx.Response)
        mock_resp.status_code = 404
        mock_resp.json.return_value = {'detail': 'Not found'}
        mock_resp.text = '{"detail": "Not found"}'
        error = httpx.HTTPStatusError('Not found', request=MagicMock(), response=mock_resp)
        with patch('lumenova_beacon.governance.client.httpx.post', side_effect=error):
            with pytest.raises(GovernanceNotFoundError):
                governance_client.evaluate(
                    phase='pre', type_='tool', stack_ids=['test'],
                )

    def test_evaluate_422_raises_validation(self, governance_client):
        mock_resp = MagicMock(spec=httpx.Response)
        mock_resp.status_code = 422
        mock_resp.json.return_value = {'detail': 'Invalid input'}
        mock_resp.text = '{"detail": "Invalid input"}'
        error = httpx.HTTPStatusError('Validation', request=MagicMock(), response=mock_resp)
        with patch('lumenova_beacon.governance.client.httpx.post', side_effect=error):
            with pytest.raises(GovernanceValidationError):
                governance_client.evaluate(
                    phase='pre', type_='tool', stack_ids=['test'],
                )

    def test_evaluate_connection_error(self, governance_client):
        with patch(
            'lumenova_beacon.governance.client.httpx.post',
            side_effect=httpx.ConnectError('Connection refused'),
        ):
            with pytest.raises(GovernanceError, match='Failed to connect'):
                governance_client.evaluate(
                    phase='pre', type_='tool', stack_ids=['test'],
                )

    def test_evaluate_custom_timeout(self, governance_client):
        mock_resp = MagicMock(spec=httpx.Response)
        mock_resp.status_code = 200
        mock_resp.text = '{}'
        mock_resp.json.return_value = {'decision': 'allow', 'message': None, 'rules_evaluated': [], 'latency_ms': 0}
        with patch(
            'lumenova_beacon.governance.client.httpx.post', return_value=mock_resp
        ) as mock_post:
            governance_client.evaluate(
                phase='pre', type_='tool', stack_ids=['test'], timeout=2.0,
            )
            assert mock_post.call_args.kwargs['timeout'] == 2.0


# ---------------------------------------------------------------------------
# Default result parser tests
# ---------------------------------------------------------------------------


class TestDefaultResultParser:
    def test_decision_allow(self):
        assert BeaconLangGraphGovernanceHandler._default_result_parser(
            {'decision': 'allow'}
        ) is True

    def test_decision_block(self):
        assert BeaconLangGraphGovernanceHandler._default_result_parser(
            {'decision': 'block', 'message': 'no'}
        ) is False

    def test_decision_apply_guardrail(self):
        """apply_guardrail is not block, so it should be allowed."""
        assert BeaconLangGraphGovernanceHandler._default_result_parser(
            {'decision': 'apply_guardrail'}
        ) is True

    def test_backward_compat_nested_decision(self):
        """Old nested format should still work via backward compat."""
        assert BeaconLangGraphGovernanceHandler._default_result_parser(
            {'decision': {'action': 'allow'}}
        ) is True
        assert BeaconLangGraphGovernanceHandler._default_result_parser(
            {'decision': {'action': 'block'}}
        ) is False

    def test_empty_result_allows(self):
        assert BeaconLangGraphGovernanceHandler._default_result_parser({}) is True


# ---------------------------------------------------------------------------
# BeaconLangGraphGovernanceHandler tests
# ---------------------------------------------------------------------------


class TestBeaconLangGraphGovernanceHandler:
    def _make_handler(self, _patch_transport, **kwargs):
        kwargs.setdefault('stack_ids', ['test-stack'])
        return BeaconLangGraphGovernanceHandler(**kwargs)

    def _serialized_tool(self, name='search'):
        return {'name': name, 'id': ['langchain', 'tools', name]}

    def _serialized_llm(self, model='gpt-4'):
        return {'name': 'ChatOpenAI', 'kwargs': {'model_name': model}}

    def _human_msg(self, text='Hello'):
        msg = MagicMock()
        msg.type = 'human'
        msg.content = text
        return msg

    # -- Tool pre-execution: payload structure --

    def test_on_tool_start_payload_structure(
        self, _patch_transport, allow_response, run_id,
    ):
        """Verify the extra payload uses nested tool object matching Rego format."""
        handler = self._make_handler(_patch_transport, stack_ids=['stack-1'])
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_start(
                self._serialized_tool('send_email'),
                '',
                run_id=run_id,
                inputs={'to': 'x@y.com', 'subject': 'Hi'},
            )
            call_kwargs = mock_eval.call_args.kwargs
            assert call_kwargs['phase'] == 'pre'
            assert call_kwargs['type_'] == 'tool'
            tool = call_kwargs['extra']['tool']
            assert tool['name'] == 'send_email'
            assert tool['parameters'] == {'to': 'x@y.com', 'subject': 'Hi'}
            assert 'output' not in tool  # pre-execution has no output

    # -- Tool post-execution: payload structure --

    def test_on_tool_end_payload_structure(
        self, _patch_transport, allow_response, run_id,
    ):
        """Verify post-execution includes tool.output and tool.parameters."""
        handler = self._make_handler(_patch_transport)
        handler._runs[str(run_id)] = {
            'tool_name': 'read_file',
            'parameters': {'path': '/data/f.csv'},
        }
        output = MagicMock()
        output.content = 'Employee record: Jane Doe'
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_end(output, run_id=run_id)
            tool = mock_eval.call_args.kwargs['extra']['tool']
            assert tool['name'] == 'read_file'
            assert tool['parameters'] == {'path': '/data/f.csv'}
            assert tool['output'] == 'Employee record: Jane Doe'

    # -- Tool pre-execution: blocked --

    def test_on_tool_start_blocked(self, _patch_transport, block_response, run_id):
        handler = self._make_handler(_patch_transport, stack_ids=['stack-1'])
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError, match="tool 'search'"):
                handler.on_tool_start(
                    self._serialized_tool(),
                    '{"query": "test"}',
                    run_id=run_id,
                )

    # -- Tool pre-execution: allowed --

    def test_on_tool_start_allowed(self, _patch_transport, allow_response, run_id):
        handler = self._make_handler(_patch_transport, stack_ids=['stack-1'])
        with patch.object(handler._client, 'evaluate', return_value=allow_response):
            handler.on_tool_start(
                self._serialized_tool(),
                '{"query": "test"}',
                run_id=run_id,
            )

    # -- Tool pre-execution: fail-open --

    def test_on_tool_start_fail_open(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport, fail_open=True)
        with patch.object(
            handler._client, 'evaluate', side_effect=GovernanceError('unreachable')
        ):
            handler.on_tool_start(
                self._serialized_tool(), '{}', run_id=run_id,
            )

    # -- Tool pre-execution: fail-closed --

    def test_on_tool_start_fail_closed(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport, fail_open=False)
        with patch.object(
            handler._client, 'evaluate', side_effect=GovernanceError('unreachable')
        ):
            with pytest.raises(GovernanceError, match='unreachable'):
                handler.on_tool_start(
                    self._serialized_tool(), '{}', run_id=run_id,
                )

    # -- Tool args extraction --

    def test_on_tool_start_extracts_args_from_kwargs_inputs(
        self, _patch_transport, allow_response, run_id,
    ):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_start(
                self._serialized_tool(), '', run_id=run_id,
                inputs={'query': 'test', 'limit': 10},
            )
            tool = mock_eval.call_args.kwargs['extra']['tool']
            assert tool['parameters'] == {'query': 'test', 'limit': 10}

    def test_on_tool_start_parses_json_input_str(
        self, _patch_transport, allow_response, run_id,
    ):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_start(
                self._serialized_tool(), '{"query": "test"}', run_id=run_id,
            )
            tool = mock_eval.call_args.kwargs['extra']['tool']
            assert tool['parameters'] == {'query': 'test'}

    def test_on_tool_start_non_json_input_str(
        self, _patch_transport, allow_response, run_id,
    ):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_start(
                self._serialized_tool(), 'plain text', run_id=run_id,
            )
            tool = mock_eval.call_args.kwargs['extra']['tool']
            assert tool['parameters'] == {}

    # -- Tool post-execution --

    def test_on_tool_end_blocked(self, _patch_transport, block_response, run_id):
        handler = self._make_handler(_patch_transport)
        handler._runs[str(run_id)] = {'tool_name': 'search', 'parameters': {}}
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_end('result text', run_id=run_id)

    def test_on_tool_end_extracts_content_attribute(
        self, _patch_transport, allow_response, run_id,
    ):
        handler = self._make_handler(_patch_transport)
        handler._runs[str(run_id)] = {'tool_name': 'search', 'parameters': {}}
        output = MagicMock()
        output.content = 'actual tool result'
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_end(output, run_id=run_id)
            tool = mock_eval.call_args.kwargs['extra']['tool']
            assert tool['output'] == 'actual tool result'

    def test_on_tool_end_gets_name_from_kwargs(
        self, _patch_transport, allow_response, run_id,
    ):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_end('result', run_id=run_id, name='calc')
            tool = mock_eval.call_args.kwargs['extra']['tool']
            assert tool['name'] == 'calc'

    # -- Tool error cleans up --

    def test_on_tool_error_cleans_up(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport)
        handler._runs[str(run_id)] = {'tool_name': 'search', 'parameters': {}}
        handler.on_tool_error(RuntimeError('boom'), run_id=run_id)
        assert str(run_id) not in handler._runs

    # -- LLM pre-execution: payload structure --

    def test_on_chat_model_start_payload_structure(
        self, _patch_transport, allow_response, run_id,
    ):
        """Verify the extra payload uses nested llm object."""
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_chat_model_start(
                self._serialized_llm(),
                [[self._human_msg('Tell me a secret')]],
                run_id=run_id,
            )
            call_kwargs = mock_eval.call_args.kwargs
            assert call_kwargs['phase'] == 'pre'
            assert call_kwargs['type_'] == 'llm'
            llm = call_kwargs['extra']['llm']
            assert llm['model_name'] == 'gpt-4'
            assert llm['model_call_count'] == 1
            assert llm['last_user_content'] == 'Tell me a secret'

    # -- LLM pre-execution: blocked --

    def test_on_chat_model_start_blocked(self, _patch_transport, block_response, run_id):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError, match="LLM 'gpt-4'"):
                handler.on_chat_model_start(
                    self._serialized_llm(),
                    [[self._human_msg('Hello')]],
                    run_id=run_id,
                )

    def test_model_call_count_increments(self, _patch_transport, allow_response):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_chat_model_start(
                self._serialized_llm(), [[self._human_msg()]], run_id=uuid4(),
            )
            handler.on_chat_model_start(
                self._serialized_llm(), [[self._human_msg()]], run_id=uuid4(),
            )
            calls = mock_eval.call_args_list
            assert calls[0].kwargs['extra']['llm']['model_call_count'] == 1
            assert calls[1].kwargs['extra']['llm']['model_call_count'] == 2

    # -- on_llm_start (non-chat) --

    def test_on_llm_start_payload(self, _patch_transport, allow_response, run_id):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_llm_start(
                self._serialized_llm(), ['Hello, how are you?'], run_id=run_id,
            )
            llm = mock_eval.call_args.kwargs['extra']['llm']
            assert llm['model_call_count'] == 1
            assert llm['last_user_content'] == 'Hello, how are you?'

    # -- LLM post-execution --

    def test_on_llm_end_blocked(self, _patch_transport, block_response, run_id):
        handler = self._make_handler(_patch_transport)
        handler._runs[str(run_id)] = {'model_name': 'gpt-4'}
        mock_response = MagicMock()
        mock_response.llm_output = {'token_usage': {'prompt_tokens': 10, 'completion_tokens': 20}}
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError):
                handler.on_llm_end(mock_response, run_id=run_id)

    def test_on_llm_end_payload(self, _patch_transport, allow_response, run_id):
        handler = self._make_handler(_patch_transport)
        handler._model_call_count = 3
        handler._runs[str(run_id)] = {'model_name': 'gpt-4'}
        mock_response = MagicMock()
        mock_response.llm_output = None
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_llm_end(mock_response, run_id=run_id)
            llm = mock_eval.call_args.kwargs['extra']['llm']
            assert llm['model_call_count'] == 3
            assert llm['model_name'] == 'gpt-4'

    # -- LLM error cleans up --

    def test_on_llm_error_cleans_up(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport)
        handler._runs[str(run_id)] = {'model_name': 'gpt-4'}
        handler.on_llm_error(RuntimeError('boom'), run_id=run_id)
        assert str(run_id) not in handler._runs

    # -- Disabled hooks --

    def test_disabled_hook_skips_evaluation(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport, enabled_hooks={'pre_tool'})
        with patch.object(handler._client, 'evaluate') as mock_eval:
            handler.on_chat_model_start(
                self._serialized_llm(), [[self._human_msg()]], run_id=run_id,
            )
            mock_eval.assert_not_called()

    def test_disabled_post_tool_skips_evaluation(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport, enabled_hooks={'pre_tool'})
        handler._runs[str(run_id)] = {'tool_name': 'search', 'parameters': {}}
        with patch.object(handler._client, 'evaluate') as mock_eval:
            handler.on_tool_end('result', run_id=run_id)
            mock_eval.assert_not_called()

    # -- Custom result parser --

    def test_custom_result_parser(self, _patch_transport, run_id):
        custom_parser = lambda r: r.get('verdict') != 'deny'
        handler = self._make_handler(
            _patch_transport,
            stack_ids=['test-stack'],
            result_parser=custom_parser,
        )
        response = {'verdict': 'deny', 'latency_ms': 1.0, 'rules_evaluated': []}
        with patch.object(handler._client, 'evaluate', return_value=response):
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=run_id)

    def test_custom_result_parser_allows(self, _patch_transport, run_id):
        custom_parser = lambda r: r.get('verdict') != 'deny'
        handler = self._make_handler(
            _patch_transport,
            stack_ids=['test-stack'],
            result_parser=custom_parser,
        )
        response = {'verdict': 'ok', 'latency_ms': 1.0, 'rules_evaluated': []}
        with patch.object(handler._client, 'evaluate', return_value=response):
            handler.on_tool_start(self._serialized_tool(), '{}', run_id=run_id)

    # -- Violation error metadata --

    def test_violation_error_metadata(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport)
        response = {
            'decision': 'block',
            'message': 'not allowed',
            'rules_evaluated': [{'name': 'policy-a', 'version': 2}],
            'latency_ms': 7.5,
        }
        with patch.object(handler._client, 'evaluate', return_value=response):
            with pytest.raises(GovernanceViolationError) as exc_info:
                handler.on_tool_start(
                    self._serialized_tool('dangerous_tool'), '{}', run_id=run_id,
                )
            err = exc_info.value
            assert err.result['decision'] == 'block'
            assert err.policies == [{'name': 'policy-a', 'version': 2}]
            assert err.latency_ms == 7.5
            assert 'not allowed' in str(err)

    def test_violation_error_includes_message(self, _patch_transport, run_id):
        """The decision message should appear in the exception string."""
        handler = self._make_handler(_patch_transport)
        response = {
            'decision': 'block',
            'message': 'Cannot send internal data to external recipients',
            'rules_evaluated': [],
            'latency_ms': 2.0,
        }
        with patch.object(handler._client, 'evaluate', return_value=response):
            with pytest.raises(GovernanceViolationError, match='Cannot send internal data'):
                handler.on_tool_start(
                    self._serialized_tool('send_email'), '{}', run_id=run_id,
                )

    def test_violation_error_no_message(self, _patch_transport, run_id):
        """When decision has no message, the error should still be clear."""
        handler = self._make_handler(_patch_transport)
        response = {
            'decision': 'block',
            'message': None,
            'rules_evaluated': [],
            'latency_ms': 1.0,
        }
        with patch.object(handler._client, 'evaluate', return_value=response):
            with pytest.raises(GovernanceViolationError, match="tool 'search'") as exc_info:
                handler.on_tool_start(
                    self._serialized_tool(), '{}', run_id=run_id,
                )
            # Should NOT have a trailing ": " when no message
            assert str(exc_info.value).endswith('at pre')

    # -- Default parser: missing decision key --

    def test_default_parser_missing_decision_allows(self, _patch_transport, run_id):
        handler = self._make_handler(_patch_transport)
        response = {'some_other_key': True, 'latency_ms': 1.0, 'rules_evaluated': []}
        with patch.object(handler._client, 'evaluate', return_value=response):
            handler.on_tool_start(self._serialized_tool(), '{}', run_id=run_id)

    # -- Truncation --

    def test_large_tool_output_truncated(self, _patch_transport, allow_response, run_id):
        handler = self._make_handler(_patch_transport)
        handler._runs[str(run_id)] = {'tool_name': 'search', 'parameters': {}}
        output = MagicMock()
        output.content = 'x' * 10000
        with patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval:
            handler.on_tool_end(output, run_id=run_id)
            tool = mock_eval.call_args.kwargs['extra']['tool']
            assert len(tool['output']) == 4096

    # -- Direct kwargs --

    def test_kwargs_shorthand(self, _patch_transport):
        handler = self._make_handler(
            _patch_transport, stack_ids=['s1'], tags=['t1'], fail_open=False, timeout=3.0,
        )
        assert handler._stack_ids == ['s1']
        assert handler._tags == ['t1']
        assert handler._fail_open is False
        assert handler._timeout == 3.0

    # -- raise_error --

    def test_raise_error_attribute(self, _patch_transport):
        handler = self._make_handler(_patch_transport)
        assert handler.raise_error is True

    # -- on_violation: invalid value --

    def test_invalid_on_violation_raises(self, _patch_transport):
        with pytest.raises(ValueError, match="on_violation must be one of"):
            self._make_handler(_patch_transport, on_violation='invalid')


# ---------------------------------------------------------------------------
# Graceful stop tests
# ---------------------------------------------------------------------------


class TestGracefulStop:
    def _make_handler(self, _patch_transport, **kwargs):
        kwargs.setdefault('stack_ids', ['test-stack'])
        return BeaconLangGraphGovernanceHandler(**kwargs)

    def _serialized_tool(self, name='search'):
        return {'name': name, 'id': ['langchain', 'tools', name]}

    def _serialized_llm(self, model='gpt-4'):
        return {'name': 'ChatOpenAI', 'kwargs': {'model_name': model}}

    def _human_msg(self, text='Hello'):
        msg = MagicMock()
        msg.type = 'human'
        msg.content = text
        return msg

    def test_on_violation_raise_is_default(self, _patch_transport):
        handler = self._make_handler(_patch_transport)
        assert handler._on_violation == 'raise'

    def test_violation_count_starts_at_zero(self, _patch_transport):
        handler = self._make_handler(_patch_transport)
        assert handler.violation_count == 0

    def test_violation_count_increments(self, _patch_transport, block_response, run_id):
        handler = self._make_handler(_patch_transport)
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=run_id)
        assert handler.violation_count == 1

    def test_on_violation_stop_sets_stopped_flag(self, _patch_transport, block_response, run_id):
        handler = self._make_handler(_patch_transport, on_violation='stop')
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=run_id)
        assert handler._stopped is True

    def test_stopped_handler_blocks_tool_without_api_call(
        self, _patch_transport, run_id,
    ):
        handler = self._make_handler(_patch_transport, on_violation='stop')
        handler._stopped = True
        with patch.object(handler._client, 'evaluate') as mock_eval:
            with pytest.raises(GovernanceViolationError, match='Agent stopped'):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=run_id)
            mock_eval.assert_not_called()

    def test_stopped_handler_blocks_llm_without_api_call(
        self, _patch_transport, run_id,
    ):
        handler = self._make_handler(_patch_transport, on_violation='stop')
        handler._stopped = True
        with patch.object(handler._client, 'evaluate') as mock_eval:
            with pytest.raises(GovernanceViolationError, match='Agent stopped'):
                handler.on_chat_model_start(
                    self._serialized_llm(),
                    [[self._human_msg()]],
                    run_id=run_id,
                )
            mock_eval.assert_not_called()

    def test_stopped_handler_blocks_on_llm_start(
        self, _patch_transport, run_id,
    ):
        handler = self._make_handler(_patch_transport, on_violation='stop')
        handler._stopped = True
        with patch.object(handler._client, 'evaluate') as mock_eval:
            with pytest.raises(GovernanceViolationError, match='Agent stopped'):
                handler.on_llm_start(
                    self._serialized_llm(), ['Hello'], run_id=run_id,
                )
            mock_eval.assert_not_called()

    def test_max_violations_escalates_to_stop(self, _patch_transport, block_response):
        handler = self._make_handler(
            _patch_transport, on_violation='raise', max_violations=2,
        )
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            # First violation: raise (not stopped)
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
            assert handler._stopped is False
            assert handler.violation_count == 1

            # Second violation: escalates to stop
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
            assert handler._stopped is True
            assert handler.violation_count == 2

    def test_max_violations_subsequent_calls_blocked(
        self, _patch_transport, block_response,
    ):
        handler = self._make_handler(
            _patch_transport, on_violation='raise', max_violations=1,
        )
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())

        # Now stopped — next call blocked without API call
        with patch.object(handler._client, 'evaluate') as mock_eval:
            with pytest.raises(GovernanceViolationError, match='Agent stopped'):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
            mock_eval.assert_not_called()

    def test_reset_clears_state(self, _patch_transport, block_response):
        handler = self._make_handler(_patch_transport, on_violation='stop')
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
        assert handler._stopped is True
        assert handler.violation_count == 1

        handler.reset()
        assert handler._stopped is False
        assert handler.violation_count == 0

    def test_raise_mode_does_not_set_stopped(self, _patch_transport, block_response):
        handler = self._make_handler(_patch_transport, on_violation='raise')
        with patch.object(handler._client, 'evaluate', return_value=block_response):
            with pytest.raises(GovernanceViolationError):
                handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
        assert handler._stopped is False
        assert handler.violation_count == 1


# ---------------------------------------------------------------------------
# Trace context propagation tests
# ---------------------------------------------------------------------------


class TestTraceContextPropagation:
    def _make_handler(self, _patch_transport, **kwargs):
        kwargs.setdefault('stack_ids', ['test-stack'])
        return BeaconLangGraphGovernanceHandler(**kwargs)

    def _serialized_tool(self, name='search'):
        return {'name': name, 'id': ['langchain', 'tools', name]}

    def _serialized_llm(self, model='gpt-4'):
        return {'name': 'ChatOpenAI', 'kwargs': {'model_name': model}}

    def _human_msg(self, text='Hello'):
        msg = MagicMock()
        msg.type = 'human'
        msg.content = text
        return msg

    def test_trace_context_forwarded_on_tool_start(
        self, _patch_transport, allow_response,
    ):
        """When tracing is active, trace_id and span_id are sent to the API."""
        handler = self._make_handler(_patch_transport)
        mock_span = MagicMock()
        mock_span.span_id = 'span-xyz'
        with (
            patch('lumenova_beacon.governance.integrations.langchain.get_current_trace_id', return_value='trace-abc'),
            patch('lumenova_beacon.governance.integrations.langchain.get_current_span', return_value=mock_span),
            patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval,
        ):
            handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
            assert mock_eval.call_args.kwargs['trace_id'] == 'trace-abc'
            assert mock_eval.call_args.kwargs['span_id'] == 'span-xyz'

    def test_trace_context_forwarded_on_tool_end(
        self, _patch_transport, allow_response,
    ):
        handler = self._make_handler(_patch_transport)
        run_id = uuid4()
        handler._runs[str(run_id)] = {'tool_name': 'search', 'parameters': {}}
        mock_span = MagicMock()
        mock_span.span_id = 'span-456'
        with (
            patch('lumenova_beacon.governance.integrations.langchain.get_current_trace_id', return_value='trace-123'),
            patch('lumenova_beacon.governance.integrations.langchain.get_current_span', return_value=mock_span),
            patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval,
        ):
            handler.on_tool_end('result', run_id=run_id)
            assert mock_eval.call_args.kwargs['trace_id'] == 'trace-123'
            assert mock_eval.call_args.kwargs['span_id'] == 'span-456'

    def test_trace_context_forwarded_on_llm_start(
        self, _patch_transport, allow_response,
    ):
        handler = self._make_handler(_patch_transport)
        mock_span = MagicMock()
        mock_span.span_id = 'span-llm'
        with (
            patch('lumenova_beacon.governance.integrations.langchain.get_current_trace_id', return_value='trace-llm'),
            patch('lumenova_beacon.governance.integrations.langchain.get_current_span', return_value=mock_span),
            patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval,
        ):
            handler.on_chat_model_start(
                self._serialized_llm(), [[self._human_msg()]], run_id=uuid4(),
            )
            assert mock_eval.call_args.kwargs['trace_id'] == 'trace-llm'
            assert mock_eval.call_args.kwargs['span_id'] == 'span-llm'

    def test_none_when_no_tracing_active(
        self, _patch_transport, allow_response,
    ):
        """When no tracing handler is active, trace_id and span_id are None."""
        handler = self._make_handler(_patch_transport)
        with (
            patch('lumenova_beacon.governance.integrations.langchain.get_current_trace_id', return_value=None),
            patch('lumenova_beacon.governance.integrations.langchain.get_current_span', return_value=None),
            patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval,
        ):
            handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
            assert mock_eval.call_args.kwargs['trace_id'] is None
            assert mock_eval.call_args.kwargs['span_id'] is None

    def test_trace_context_read_at_evaluation_time(
        self, _patch_transport, allow_response,
    ):
        """Context is read at evaluation time, not at handler construction."""
        # Construct handler with no trace context.
        with (
            patch('lumenova_beacon.governance.integrations.langchain.get_current_trace_id', return_value=None),
            patch('lumenova_beacon.governance.integrations.langchain.get_current_span', return_value=None),
        ):
            handler = self._make_handler(_patch_transport)

        # Now set trace context and trigger evaluation.
        mock_span = MagicMock()
        mock_span.span_id = 'span-late'
        with (
            patch('lumenova_beacon.governance.integrations.langchain.get_current_trace_id', return_value='trace-late'),
            patch('lumenova_beacon.governance.integrations.langchain.get_current_span', return_value=mock_span),
            patch.object(handler._client, 'evaluate', return_value=allow_response) as mock_eval,
        ):
            handler.on_tool_start(self._serialized_tool(), '{}', run_id=uuid4())
            assert mock_eval.call_args.kwargs['trace_id'] == 'trace-late'
            assert mock_eval.call_args.kwargs['span_id'] == 'span-late'
