# -*- coding: utf-8 -*-
# Copyright (C) Alexander Pace, Daniel Wysocki (2026)
#
# This file is part of gracedb
#
# gracedb is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# It is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with gracedb.  If not, see <http://www.gnu.org/licenses/>.
"""Tests for Kafka consumer."""

import signal
import warnings
from unittest import mock
import pytest

from ligo.gracedb.kafka.exceptions import KafkaAuthenticationError


def _make_consumer(**kwargs):
    """Create a GraceDbKafkaConsumer with mocked internals.

    Returns:
        tuple: (consumer, mock_confluent_consumer)
    """
    mock_confluent = mock.Mock()
    defaults = {
        "instance_name": "test",
        "reload_cred": False,
    }
    defaults.update(kwargs)

    with mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE', True), \
         mock.patch('ligo.gracedb.kafka.consumer.Consumer',
                    return_value=mock_confluent), \
         mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback',
                    return_value=lambda x: ("token", 1000)):
        from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer

        consumer = GraceDbKafkaConsumer(
            "localhost:9092",
            **defaults,
        )
    return consumer, mock_confluent


@mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE', True)
class TestConsumerDefaults:
    """Test consumer default configuration."""

    def test_generate_group_id_format(self):
        """_generate_group_id returns gracedb-client-<8 hex chars>."""
        import re
        from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
        gid = GraceDbKafkaConsumer._generate_group_id()
        assert re.match(r'^gracedb-client-[0-9a-f]{8}$', gid)

    def test_generate_group_id_uniqueness(self):
        """Successive calls produce different group IDs."""
        from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
        ids = {GraceDbKafkaConsumer._generate_group_id() for _ in range(10)}
        assert len(ids) == 10

    def test_config_defaults(self):
        """Verify default config values passed to confluent Consumer."""
        mock_confluent = mock.Mock()
        mock_consumer_cls = mock.Mock(return_value=mock_confluent)

        with mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE',
                        True), \
             mock.patch('ligo.gracedb.kafka.consumer.Consumer',
                        mock_consumer_cls), \
             mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback',
                        return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
            GraceDbKafkaConsumer(
                "localhost:9092", instance_name="test",
                reload_cred=False,
            )

        config = mock_consumer_cls.call_args[0][0]
        assert config['auto.offset.reset'] == 'latest'
        assert config['enable.auto.commit'] is True
        assert config['group.id'].startswith('gracedb-client-')

    def test_explicit_group_id(self):
        """Explicit group_id is used as-is."""
        mock_consumer_cls = mock.Mock(return_value=mock.Mock())

        with mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE',
                        True), \
             mock.patch('ligo.gracedb.kafka.consumer.Consumer',
                        mock_consumer_cls), \
             mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback',
                        return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
            GraceDbKafkaConsumer(
                "localhost:9092",
                instance_name="test",
                group_id="my-stable-group",
                reload_cred=False,
            )

        config = mock_consumer_cls.call_args[0][0]
        assert config['group.id'] == 'my-stable-group'

    def test_explicit_offset_reset_earliest(self):
        """Explicit auto_offset_reset='earliest' is respected."""
        mock_consumer_cls = mock.Mock(return_value=mock.Mock())

        with mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE',
                        True), \
             mock.patch('ligo.gracedb.kafka.consumer.Consumer',
                        mock_consumer_cls), \
             mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback',
                        return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
            GraceDbKafkaConsumer(
                "localhost:9092",
                instance_name="test",
                auto_offset_reset="earliest",
                reload_cred=False,
            )

        config = mock_consumer_cls.call_args[0][0]
        assert config['auto.offset.reset'] == 'earliest'

    def test_auto_commit_override(self):
        """enable.auto.commit can be overridden via kafka_config."""
        mock_consumer_cls = mock.Mock(return_value=mock.Mock())

        with mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE',
                        True), \
             mock.patch('ligo.gracedb.kafka.consumer.Consumer',
                        mock_consumer_cls), \
             mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback',
                        return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
            GraceDbKafkaConsumer(
                "localhost:9092",
                instance_name="test",
                reload_cred=False,
                **{"enable.auto.commit": False},
            )

        config = mock_consumer_cls.call_args[0][0]
        assert config['enable.auto.commit'] is False


@mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE', True)
class TestConsumerReloadCred:
    """Test consumer reload_cred / reload_buffer / deprecation."""

    @mock.patch('ligo.gracedb.kafka.consumer.Consumer')
    @mock.patch('ligo.gracedb.kafka.consumer.SciTokenRefreshManager')
    def test_reload_cred_true_creates_manager(
        self, mock_manager_cls, mock_consumer_cls
    ):
        """reload_cred=True (default) creates a SciTokenRefreshManager."""
        from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
        mock_manager = mock.Mock()
        mock_manager.create_callback.return_value = lambda x: ("t", 1000)
        mock_manager_cls.return_value = mock_manager
        mock_consumer_cls.return_value = mock.Mock()

        consumer = GraceDbKafkaConsumer("localhost:9092", instance_name="test", reload_cred=True)

        mock_manager_cls.assert_called_once_with(
            broker_url="localhost:9092",
            scope="kafka.consume",
            audience=None,
            refresh_buffer=300,
        )
        assert consumer.refresh_manager is mock_manager

    @mock.patch('ligo.gracedb.kafka.consumer.Consumer')
    @mock.patch('ligo.gracedb.kafka.consumer.SciTokenRefreshManager')
    def test_reload_buffer_forwarded(self, mock_manager_cls, mock_consumer_cls):
        """reload_buffer is forwarded to SciTokenRefreshManager."""
        from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
        mock_manager = mock.Mock()
        mock_manager.create_callback.return_value = lambda x: ("t", 1000)
        mock_manager_cls.return_value = mock_manager
        mock_consumer_cls.return_value = mock.Mock()

        GraceDbKafkaConsumer(
            "localhost:9092", instance_name="test",
            reload_cred=True, reload_buffer=600,
        )

        mock_manager_cls.assert_called_once_with(
            broker_url="localhost:9092",
            scope="kafka.consume",
            audience=None,
            refresh_buffer=600,
        )

    @mock.patch('ligo.gracedb.kafka.consumer.Consumer')
    @mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback')
    def test_reload_cred_false_uses_simple_callback(
        self, mock_callback, mock_consumer_cls
    ):
        """reload_cred=False uses simple create_oauth_callback."""
        from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
        mock_callback.return_value = lambda x: ("t", 1000)
        mock_consumer_cls.return_value = mock.Mock()

        consumer = GraceDbKafkaConsumer(
            "localhost:9092", instance_name="test",
            reload_cred=False,
        )

        mock_callback.assert_called_once_with(
            broker_url="localhost:9092",
            scope="kafka.consume",
            audience=None,
        )
        assert consumer.refresh_manager is None

    @mock.patch('ligo.gracedb.kafka.consumer.Consumer')
    @mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback')
    def test_auto_refresh_token_deprecation(
        self, mock_callback, mock_consumer_cls
    ):
        """auto_refresh_token emits DeprecationWarning, maps to reload_cred."""
        from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
        mock_callback.return_value = lambda x: ("t", 1000)
        mock_consumer_cls.return_value = mock.Mock()

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            consumer = GraceDbKafkaConsumer(
                "localhost:9092", instance_name="test",
                auto_refresh_token=False,
            )

        assert len(w) == 1
        assert issubclass(w[0].category, DeprecationWarning)
        assert "auto_refresh_token is deprecated" in str(w[0].message)
        assert consumer.refresh_manager is None


@mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE', True)
class TestConsumerShutdown:
    """Test consumer graceful shutdown machinery."""

    def _make_consumer(self):
        """Create a GraceDbKafkaConsumer with mocked internals."""
        consumer, self.mock_consumer = _make_consumer()
        return consumer

    def test_stop_sets_running_false(self):
        """stop() sets _running to False."""
        consumer = self._make_consumer()
        consumer._running = True

        consumer.stop()

        assert consumer._running is False

    def test_stream_sets_running_true_on_entry(self):
        """_running is True once stream() begins iterating."""
        consumer = self._make_consumer()

        # poll returns one message then None forever
        self.mock_consumer.poll.side_effect = [mock.Mock(
            error=mock.Mock(return_value=None),
            value=mock.Mock(return_value=b'"x"'),
            key=mock.Mock(return_value=None),
            headers=mock.Mock(return_value=None),
            topic=mock.Mock(return_value='t'),
            partition=mock.Mock(return_value=0),
            offset=mock.Mock(return_value=0),
            timestamp=mock.Mock(return_value=(0, None)),
        )]

        gen = consumer.stream(count=1)
        # Before first next(), _running should still be False
        assert consumer._running is False

        # Consume the single message — _running becomes True inside
        next(gen)
        assert consumer._running is True

    def test_stream_resets_running_on_exit(self):
        """_running is False after stream() ends (count exhaustion)."""
        consumer = self._make_consumer()
        self.mock_consumer.poll.return_value = None

        # Exhaust the generator with count=0 (yields nothing)
        list(consumer.stream(count=0))

        assert consumer._running is False

    def test_stream_exits_when_stopped(self):
        """Calling stop() during iteration causes stream() to terminate."""
        consumer = self._make_consumer()

        # poll always returns None (no messages)
        call_count = 0

        def poll_and_stop(timeout):
            nonlocal call_count
            call_count += 1
            if call_count >= 3:
                consumer.stop()
            return None

        self.mock_consumer.poll.side_effect = poll_and_stop

        messages = list(consumer.stream(timeout=0.0))

        assert messages == []
        assert call_count == 3
        assert consumer._running is False

    @mock.patch('ligo.gracedb.kafka.consumer.signal')
    def test_stream_installs_signal_handlers(self, mock_signal_mod):
        """signal.signal is called with SIGINT and SIGTERM on entry."""
        consumer = self._make_consumer()

        # Make stream exit immediately
        self.mock_consumer.poll.return_value = None
        list(consumer.stream(count=0))

        # Collect the signal.signal calls
        calls = mock_signal_mod.signal.call_args_list
        installed_signals = [c[0][0] for c in calls]

        assert mock_signal_mod.SIGINT in installed_signals
        assert mock_signal_mod.SIGTERM in installed_signals

    @mock.patch('ligo.gracedb.kafka.consumer.signal')
    def test_stream_restores_signal_handlers(self, mock_signal_mod):
        """Original handlers are restored after stream() exits."""
        consumer = self._make_consumer()

        prev_int = mock.Mock(name='prev_sigint')
        prev_term = mock.Mock(name='prev_sigterm')

        # First two calls install handlers and return previous ones;
        # next two calls restore them.
        mock_signal_mod.signal.side_effect = [
            prev_int, prev_term,
            None, None,
        ]

        self.mock_consumer.poll.return_value = None
        list(consumer.stream(count=0))

        restore_calls = mock_signal_mod.signal.call_args_list[2:]
        assert (mock_signal_mod.SIGINT, prev_int) == restore_calls[0][0]
        assert (mock_signal_mod.SIGTERM, prev_term) == restore_calls[1][0]

    @mock.patch('ligo.gracedb.kafka.consumer.signal')
    def test_stream_skips_signals_in_non_main_thread(self, mock_signal_mod):
        """When signal.signal raises ValueError, stream still works."""
        consumer = self._make_consumer()

        mock_signal_mod.signal.side_effect = ValueError(
            "signal only works in main thread"
        )

        # poll returns one message then triggers stop
        call_count = 0

        def poll_side_effect(timeout):
            nonlocal call_count
            call_count += 1
            if call_count >= 2:
                consumer.stop()
            return None

        self.mock_consumer.poll.side_effect = poll_side_effect

        # Should not raise
        messages = list(consumer.stream(timeout=0.0))

        assert messages == []
        # signal.signal was called but raised ValueError — that's fine
        assert mock_signal_mod.signal.called

    def test_handle_shutdown_signal(self):
        """_handle_shutdown_signal() sets _running to False."""
        consumer = self._make_consumer()
        consumer._running = True

        consumer._handle_shutdown_signal(signal.SIGINT, None)

        assert consumer._running is False

    def test_stream_with_count(self):
        """count parameter limits the number of yielded messages."""
        consumer = self._make_consumer()

        # Create mock messages
        def make_msg():
            m = mock.Mock()
            m.error.return_value = None
            m.value.return_value = b'"data"'
            m.key.return_value = None
            m.headers.return_value = None
            m.topic.return_value = 'test'
            m.partition.return_value = 0
            m.offset.return_value = 0
            m.timestamp.return_value = (0, None)
            return m

        self.mock_consumer.poll.side_effect = [
            make_msg(), make_msg(), make_msg(), make_msg(), make_msg(),
        ]

        messages = list(consumer.stream(count=3))

        assert len(messages) == 3
        assert consumer._running is False


@mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE', True)
class TestConsumerHeaderDecode:
    """Test header decoding with fallback on invalid UTF-8."""

    def _make_consumer(self):
        consumer, self.mock_consumer = _make_consumer()
        return consumer

    def _make_msg(self, headers):
        m = mock.Mock()
        m.error.return_value = None
        m.value.return_value = b'"data"'
        m.key.return_value = None
        m.headers.return_value = headers
        m.topic.return_value = 'test'
        m.partition.return_value = 0
        m.offset.return_value = 0
        m.timestamp.return_value = (0, None)
        return m

    def test_valid_utf8_headers(self):
        """Headers with valid UTF-8 bytes are decoded to strings."""
        consumer = self._make_consumer()
        headers = [('key1', b'value1'), ('key2', b'value2')]
        self.mock_consumer.poll.return_value = self._make_msg(headers)

        result = consumer.poll()
        assert result['headers'] == {'key1': 'value1', 'key2': 'value2'}

    def test_invalid_utf8_headers_fallback(self):
        """Headers with invalid UTF-8 fall back to raw bytes."""
        consumer = self._make_consumer()
        bad_bytes = b'\xff\xfe\x80'
        headers = [('key1', bad_bytes), ('key2', b'ok')]
        self.mock_consumer.poll.return_value = self._make_msg(headers)

        result = consumer.poll()
        assert result['headers'] == {'key1': bad_bytes, 'key2': b'ok'}

    def test_none_headers(self):
        """None headers remain None."""
        consumer = self._make_consumer()
        self.mock_consumer.poll.return_value = self._make_msg(None)

        result = consumer.poll()
        assert result['headers'] is None


@mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE', True)
class TestConsumerAuthError:
    """Test consumer authentication error detection and propagation."""

    def _make_consumer(self):
        """Create a GraceDbKafkaConsumer with mocked internals."""
        consumer, self.mock_consumer = _make_consumer()
        return consumer

    def _make_auth_error(self):
        """Create a mock Kafka error with _AUTHENTICATION code."""
        from ligo.gracedb.kafka.consumer import _AUTHENTICATION
        error = mock.Mock()
        error.code.return_value = _AUTHENTICATION
        error.__str__ = mock.Mock(return_value="Token expired")
        return error

    def _make_non_auth_error(self):
        """Create a mock Kafka error with a non-auth code."""
        error = mock.Mock()
        error.code.return_value = -1  # Generic error
        return error

    def test_error_callback_stores_auth_error(self):
        """_error_callback stores auth errors in _auth_error."""
        consumer = self._make_consumer()
        error = self._make_auth_error()

        consumer._error_callback(error)

        assert consumer._auth_error is error

    def test_error_callback_sets_running_false(self):
        """_error_callback sets _running to False on auth error."""
        consumer = self._make_consumer()
        consumer._running = True
        error = self._make_auth_error()

        consumer._error_callback(error)

        assert consumer._running is False

    def test_error_callback_ignores_non_auth_error(self):
        """_error_callback does not set _auth_error for non-auth errors."""
        consumer = self._make_consumer()
        error = self._make_non_auth_error()

        consumer._error_callback(error)

        assert consumer._auth_error is None

    def test_error_callback_does_not_stop_on_non_auth_error(self):
        """_error_callback leaves _running unchanged for non-auth errors."""
        consumer = self._make_consumer()
        consumer._running = True
        error = self._make_non_auth_error()

        consumer._error_callback(error)

        assert consumer._running is True

    def test_check_auth_raises_on_auth_error(self):
        """_check_auth raises KafkaAuthenticationError when error stored."""
        consumer = self._make_consumer()
        consumer._auth_error = self._make_auth_error()

        with pytest.raises(KafkaAuthenticationError, match="SciToken authentication failed"):
            consumer._check_auth()

    def test_check_auth_passes_when_no_error(self):
        """_check_auth does nothing when no auth error is stored."""
        consumer = self._make_consumer()
        consumer._check_auth()  # Should not raise

    def test_poll_raises_on_prior_auth_error(self):
        """poll() raises KafkaAuthenticationError if auth already failed."""
        consumer = self._make_consumer()
        consumer._auth_error = self._make_auth_error()
        self.mock_consumer.poll.return_value = None

        with pytest.raises(KafkaAuthenticationError):
            consumer.poll()

        # poll() is still called; _check_auth runs after
        self.mock_consumer.poll.assert_called_once()

    def test_poll_raises_on_auth_error_during_poll(self):
        """poll() raises if auth error arrives during poll."""
        consumer = self._make_consumer()

        def set_auth_error(timeout):
            consumer._auth_error = self._make_auth_error()
            return None

        self.mock_consumer.poll.side_effect = set_auth_error

        with pytest.raises(KafkaAuthenticationError):
            consumer.poll()

    def test_stream_exits_on_auth_error(self):
        """stream() raises KafkaAuthenticationError when error_cb fires."""
        consumer = self._make_consumer()

        call_count = 0

        def poll_then_auth_error(timeout):
            nonlocal call_count
            call_count += 1
            if call_count >= 2:
                # Simulate error_cb firing during poll
                consumer._error_callback(self._make_auth_error())
            return None

        self.mock_consumer.poll.side_effect = poll_then_auth_error

        # stream() should raise because _check_auth runs after poll
        with pytest.raises(KafkaAuthenticationError):
            list(consumer.stream(timeout=0.0))

        assert consumer._running is False

    def test_stream_raises_auth_error_on_next_poll(self):
        """stream() raises KafkaAuthenticationError when poll detects it."""
        consumer = self._make_consumer()

        # First poll succeeds, second triggers auth error
        msg = mock.Mock()
        msg.error.return_value = None
        msg.value.return_value = b'"data"'
        msg.key.return_value = None
        msg.headers.return_value = None
        msg.topic.return_value = 'test'
        msg.partition.return_value = 0
        msg.offset.return_value = 0
        msg.timestamp.return_value = (0, None)

        call_count = 0

        def poll_side_effect(timeout):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return msg
            # Simulate auth error arriving during second poll
            consumer._auth_error = self._make_auth_error()
            return None

        self.mock_consumer.poll.side_effect = poll_side_effect

        gen = consumer.stream(timeout=0.0)
        first = next(gen)  # Should succeed
        assert first['topic'] == 'test'

        # Next iteration should raise auth error
        with pytest.raises(KafkaAuthenticationError):
            next(gen)

    def test_error_cb_wired_into_config(self):
        """error_cb is set in the Kafka config and calls _error_callback."""
        mock_consumer_cls = mock.Mock(return_value=mock.Mock())

        with mock.patch('ligo.gracedb.kafka.consumer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.consumer.Consumer',
                        mock_consumer_cls), \
             mock.patch('ligo.gracedb.kafka.consumer.create_oauth_callback',
                        return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.consumer import GraceDbKafkaConsumer
            consumer = GraceDbKafkaConsumer(
                "localhost:9092", instance_name="test",
                reload_cred=False,
            )

        config = mock_consumer_cls.call_args[0][0]
        assert 'error_cb' in config
        # Invoke the callback and verify it delegates to _error_callback
        error = self._make_auth_error()
        config['error_cb'](error)
        assert consumer._auth_error is error
