"""Tests for the @governance decorator."""

from __future__ import annotations

import asyncio
from unittest.mock import MagicMock, patch

import pytest

from lumenova_beacon.exceptions import (
    GovernanceError,
    GovernanceViolationError,
)
from lumenova_beacon.governance.decorators import governance


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


@pytest.fixture
def 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_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):
    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 allow_response():
    return {
        'decision': 'allow',
        'message': None,
        'rules_evaluated': [{'name': 'default-allow'}],
        'latency_ms': 5.2,
    }


@pytest.fixture
def block_response():
    return {
        'decision': 'block',
        'message': 'not allowed',
        'rules_evaluated': [{'name': 'block-policy'}],
        'latency_ms': 3.1,
    }


# ---------------------------------------------------------------------------
# Sync decorator tests
# ---------------------------------------------------------------------------


class TestGovernanceDecoratorSync:
    def test_allowed_calls_function(self, _patch_transport, allow_response):
        @governance(stack_ids=['test-stack'])
        def my_tool(query: str) -> str:
            return f'result for {query}'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ):
            result = my_tool(query='hello')
        assert result == 'result for hello'

    def test_pre_execution_blocked_prevents_function_call(
        self, _patch_transport, block_response,
    ):
        called = False

        @governance(stack_ids=['test-stack'])
        def my_tool(query: str) -> str:
            nonlocal called
            called = True
            return 'result'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=block_response,
        ):
            with pytest.raises(GovernanceViolationError, match='not allowed'):
                my_tool(query='hello')
        assert called is False

    def test_post_execution_blocked(self, _patch_transport, allow_response, block_response):
        @governance(stack_ids=['test-stack'], enabled_hooks={'post_tool'})
        def my_tool(query: str) -> str:
            return 'sensitive data'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=block_response,
        ):
            with pytest.raises(GovernanceViolationError):
                my_tool(query='test')

    def test_payload_structure(self, _patch_transport, allow_response):
        @governance(stack_ids=['s1'], tags=['t1'])
        def send_email(to: str, subject: str) -> bool:
            return True

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ) as mock_eval:
            send_email(to='a@b.com', subject='Hi')

            # Pre-execution call
            pre_call = mock_eval.call_args_list[0]
            assert pre_call.kwargs['phase'] == 'pre'
            assert pre_call.kwargs['type_'] == 'tool'
            assert pre_call.kwargs['stack_ids'] == ['s1']
            assert pre_call.kwargs['tags'] == ['t1']
            tool = pre_call.kwargs['extra']['tool']
            assert tool['name'] == 'send_email'
            assert tool['parameters'] == {'to': 'a@b.com', 'subject': 'Hi'}

            # Post-execution call
            post_call = mock_eval.call_args_list[1]
            assert post_call.kwargs['phase'] == 'post'
            post_tool = post_call.kwargs['extra']['tool']
            assert post_tool['output'] == 'True'

    def test_custom_name(self, _patch_transport, allow_response):
        @governance(stack_ids=['s1'], name='custom_tool')
        def my_func() -> str:
            return 'ok'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ) as mock_eval:
            my_func()
            tool = mock_eval.call_args_list[0].kwargs['extra']['tool']
            assert tool['name'] == 'custom_tool'

    def test_llm_type(self, _patch_transport, allow_response):
        @governance(stack_ids=['s1'], type_='llm')
        def call_llm(prompt: str) -> str:
            return 'response'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ) as mock_eval:
            call_llm(prompt='hello')
            call_kwargs = mock_eval.call_args_list[0].kwargs
            assert call_kwargs['type_'] == 'llm'
            assert 'llm' in call_kwargs['extra']

    def test_enabled_hooks_pre_only(self, _patch_transport, allow_response):
        @governance(stack_ids=['s1'], enabled_hooks={'pre_tool'})
        def my_tool() -> str:
            return 'result'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ) as mock_eval:
            my_tool()
            assert mock_eval.call_count == 1
            assert mock_eval.call_args.kwargs['phase'] == 'pre'

    def test_enabled_hooks_post_only(self, _patch_transport, allow_response):
        @governance(stack_ids=['s1'], enabled_hooks={'post_tool'})
        def my_tool() -> str:
            return 'result'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ) as mock_eval:
            my_tool()
            assert mock_eval.call_count == 1
            assert mock_eval.call_args.kwargs['phase'] == 'post'

    def test_fail_open(self, _patch_transport):
        @governance(stack_ids=['s1'], fail_open=True)
        def my_tool() -> str:
            return 'result'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            side_effect=GovernanceError('unreachable'),
        ):
            result = my_tool()
        assert result == 'result'

    def test_fail_closed(self, _patch_transport):
        @governance(stack_ids=['s1'], fail_open=False)
        def my_tool() -> str:
            return 'result'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            side_effect=GovernanceError('unreachable'),
        ):
            with pytest.raises(GovernanceError, match='unreachable'):
                my_tool()

    def test_custom_result_parser(self, _patch_transport):
        custom_parser = lambda r: r.get('verdict') != 'deny'

        @governance(stack_ids=['s1'], result_parser=custom_parser)
        def my_tool() -> str:
            return 'result'

        response = {'verdict': 'deny', 'latency_ms': 1.0, 'rules_evaluated': []}
        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=response,
        ):
            with pytest.raises(GovernanceViolationError):
                my_tool()

    def test_bare_decorator(self, _patch_transport, allow_response):
        """@governance without parentheses."""
        @governance
        def my_tool(query: str) -> str:
            return f'result for {query}'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ):
            result = my_tool(query='hello')
        assert result == 'result for hello'

    def test_empty_parens_decorator(self, _patch_transport, allow_response):
        """@governance() with empty parentheses."""
        @governance()
        def my_tool(query: str) -> str:
            return f'result for {query}'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ):
            result = my_tool(query='hello')
        assert result == 'result for hello'

    def test_preserves_function_metadata(self, _patch_transport):
        @governance(stack_ids=['s1'])
        def my_documented_tool():
            """Tool docstring."""
            return 'ok'

        assert my_documented_tool.__name__ == 'my_documented_tool'
        assert my_documented_tool.__doc__ == 'Tool docstring.'

    def test_output_truncated(self, _patch_transport, allow_response):
        @governance(stack_ids=['s1'])
        def my_tool() -> str:
            return 'x' * 10000

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.evaluate',
            return_value=allow_response,
        ) as mock_eval:
            my_tool()
            post_call = mock_eval.call_args_list[1]
            assert len(post_call.kwargs['extra']['tool']['output']) == 4096


# ---------------------------------------------------------------------------
# Async decorator tests
# ---------------------------------------------------------------------------


class TestGovernanceDecoratorAsync:
    def test_async_allowed(self, _patch_transport, allow_response):
        @governance(stack_ids=['test-stack'])
        async def my_async_tool(query: str) -> str:
            return f'async result for {query}'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.aevaluate',
            return_value=allow_response,
        ):
            result = asyncio.get_event_loop().run_until_complete(
                my_async_tool(query='hello')
            )
        assert result == 'async result for hello'

    def test_async_blocked(self, _patch_transport, block_response):
        called = False

        @governance(stack_ids=['test-stack'])
        async def my_async_tool(query: str) -> str:
            nonlocal called
            called = True
            return 'result'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.aevaluate',
            return_value=block_response,
        ):
            with pytest.raises(GovernanceViolationError):
                asyncio.get_event_loop().run_until_complete(
                    my_async_tool(query='hello')
                )
        assert called is False

    def test_async_fail_open(self, _patch_transport):
        @governance(stack_ids=['s1'], fail_open=True)
        async def my_async_tool() -> str:
            return 'result'

        with patch(
            'lumenova_beacon.governance.client.GovernanceClient.aevaluate',
            side_effect=GovernanceError('unreachable'),
        ):
            result = asyncio.get_event_loop().run_until_complete(my_async_tool())
        assert result == 'result'


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


class TestGovernanceDecoratorTraceContext:
    def test_sync_forwards_trace_context(self, _patch_transport, allow_response):
        """Sync decorator reads and forwards trace_id/span_id."""
        @governance(stack_ids=['s1'])
        def my_tool() -> str:
            return 'ok'

        mock_span = MagicMock()
        mock_span.span_id = 'span-dec'
        with (
            patch('lumenova_beacon.governance.decorators.get_current_trace_id', return_value='trace-dec'),
            patch('lumenova_beacon.governance.decorators.get_current_span', return_value=mock_span),
            patch(
                'lumenova_beacon.governance.client.GovernanceClient.evaluate',
                return_value=allow_response,
            ) as mock_eval,
        ):
            my_tool()
            # Both pre and post calls should include trace context.
            for call in mock_eval.call_args_list:
                assert call.kwargs['trace_id'] == 'trace-dec'
                assert call.kwargs['span_id'] == 'span-dec'

    def test_async_forwards_trace_context(self, _patch_transport, allow_response):
        """Async decorator reads and forwards trace_id/span_id."""
        @governance(stack_ids=['s1'])
        async def my_async_tool() -> str:
            return 'ok'

        mock_span = MagicMock()
        mock_span.span_id = 'span-async'
        with (
            patch('lumenova_beacon.governance.decorators.get_current_trace_id', return_value='trace-async'),
            patch('lumenova_beacon.governance.decorators.get_current_span', return_value=mock_span),
            patch(
                'lumenova_beacon.governance.client.GovernanceClient.aevaluate',
                return_value=allow_response,
            ) as mock_eval,
        ):
            asyncio.get_event_loop().run_until_complete(my_async_tool())
            for call in mock_eval.call_args_list:
                assert call.kwargs['trace_id'] == 'trace-async'
                assert call.kwargs['span_id'] == 'span-async'

    def test_no_trace_context_passes_none(self, _patch_transport, allow_response):
        """When no tracing is active, None values are passed (and omitted from payload)."""
        @governance(stack_ids=['s1'])
        def my_tool() -> str:
            return 'ok'

        with (
            patch('lumenova_beacon.governance.decorators.get_current_trace_id', return_value=None),
            patch('lumenova_beacon.governance.decorators.get_current_span', return_value=None),
            patch(
                'lumenova_beacon.governance.client.GovernanceClient.evaluate',
                return_value=allow_response,
            ) as mock_eval,
        ):
            my_tool()
            for call in mock_eval.call_args_list:
                assert call.kwargs['trace_id'] is None
                assert call.kwargs['span_id'] is None
