# -*- 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 producer."""

import base64
import json
import warnings
from unittest import mock
import pytest

from ligo.gracedb.kafka.exceptions import (
    KafkaAuthenticationError,
    KafkaProducerError,
    KafkaConnectionError,
    KafkaMessageDeliveryError,
    KafkaTokenError,
)
from ligo.gracedb.kafka.producer import KafkaException


def _make_jwt(claims):
    """Build a minimal unsigned JWT with the given payload claims."""
    header = base64.urlsafe_b64encode(
        json.dumps({"alg": "none"}).encode()
    ).rstrip(b'=').decode()
    payload = base64.urlsafe_b64encode(
        json.dumps(claims).encode()
    ).rstrip(b'=').decode()
    return f"{header}.{payload}.signature"


class TestProducerImports:
    """Test producer module imports and availability checks."""

    @mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_imports_when_kafka_available(self, mock_callback, mock_producer_class):
        """Test that producer can be imported when Kafka is available."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_class.return_value = mock.Mock()

        producer = GraceDbKafkaProducer(
            "localhost:9092", topic="test", instance_name="test",
            reload_cred=False,
        )
        assert producer is not None

    @mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', False)
    def test_import_error_when_kafka_unavailable(self):
        """Test that ImportError is raised when Kafka is not available."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        with pytest.raises(ImportError) as exc_info:
            GraceDbKafkaProducer("localhost:9092")

        assert "confluent-kafka is required" in str(exc_info.value)
        assert "pip install ligo-gracedb[kafka]" in str(exc_info.value)


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerInitialization:
    """Test producer initialization."""

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_basic_initialization(self, mock_callback, mock_producer_class):
        """Test basic producer initialization."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_instance = mock.Mock()
        mock_producer_class.return_value = mock_producer_instance

        producer = GraceDbKafkaProducer(
            bootstrap_servers="localhost:9092",
            topic="test-topic",
            instance_name="test",
            reload_cred=False,
        )

        assert producer.bootstrap_servers == "localhost:9092"
        assert producer.instance_name == "test"
        assert producer.default_topic == "test.test-topic"
        assert producer.token_scope == "kafka.produce"
        assert producer.producer is mock_producer_instance

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_ssl_configuration(self, mock_callback, mock_producer_class):
        """Test SSL configuration in producer."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_class.return_value = mock.Mock()

        producer = GraceDbKafkaProducer(
            bootstrap_servers="localhost:9092",
            instance_name="test",
            use_ssl=True,
            ca_cert_path="/tmp/cacert.pem",
            reload_cred=False,
        )

        # Check that Producer was called with SSL config
        call_args = mock_producer_class.call_args[0][0]
        assert call_args['security.protocol'] == 'SASL_SSL'
        assert call_args['sasl.mechanism'] == 'OAUTHBEARER'
        assert call_args['ssl.ca.location'] == '/tmp/cacert.pem'

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_plaintext_configuration(self, mock_callback, mock_producer_class):
        """Test plaintext (non-SSL) configuration."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_class.return_value = mock.Mock()

        producer = GraceDbKafkaProducer(
            bootstrap_servers="localhost:9092",
            instance_name="test",
            use_ssl=False,
            reload_cred=False,
        )

        call_args = mock_producer_class.call_args[0][0]
        assert call_args['security.protocol'] == 'SASL_PLAINTEXT'
        assert call_args['sasl.mechanism'] == 'OAUTHBEARER'

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.SciTokenRefreshManager')
    def test_reload_cred_true(self, mock_refresh_manager_class, mock_producer_class):
        """Test initialization with reload_cred=True creates SciTokenRefreshManager."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_refresh_manager = mock.Mock()
        mock_refresh_manager.create_callback.return_value = lambda x: ("token", 1000)
        mock_refresh_manager_class.return_value = mock_refresh_manager
        mock_producer_class.return_value = mock.Mock()

        producer = GraceDbKafkaProducer(
            bootstrap_servers="localhost:9092",
            instance_name="test",
            reload_cred=True
        )

        # Check that refresh manager was created with default buffer
        mock_refresh_manager_class.assert_called_once_with(
            broker_url="localhost:9092",
            scope="kafka.produce",
            audience=None,
            refresh_buffer=300,
        )
        assert producer.refresh_manager is mock_refresh_manager

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.SciTokenRefreshManager')
    def test_reload_buffer_forwarded(self, mock_refresh_manager_class, mock_producer_class):
        """Test that reload_buffer is forwarded to SciTokenRefreshManager."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_refresh_manager = mock.Mock()
        mock_refresh_manager.create_callback.return_value = lambda x: ("token", 1000)
        mock_refresh_manager_class.return_value = mock_refresh_manager
        mock_producer_class.return_value = mock.Mock()

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

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

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_reload_cred_false(self, mock_callback, mock_producer_class):
        """Test initialization with reload_cred=False uses simple callback."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_class.return_value = mock.Mock()

        producer = GraceDbKafkaProducer(
            bootstrap_servers="localhost:9092",
            instance_name="test",
            reload_cred=False
        )

        # Check that simple callback was created
        mock_callback.assert_called_once_with(
            broker_url="localhost:9092",
            scope="kafka.produce",
            audience=None,
        )
        assert producer.refresh_manager is None

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_auto_refresh_token_deprecation(self, mock_callback, mock_producer_class):
        """auto_refresh_token emits DeprecationWarning and maps to reload_cred."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_class.return_value = mock.Mock()

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            producer = GraceDbKafkaProducer(
                bootstrap_servers="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 producer.refresh_manager is None

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_custom_kafka_config(self, mock_callback, mock_producer_class):
        """Test passing custom Kafka configuration."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_class.return_value = mock.Mock()

        producer = GraceDbKafkaProducer(
            bootstrap_servers="localhost:9092",
            instance_name="test",
            reload_cred=False,
            compression_type="gzip",
            batch_num_messages=1000,
            linger_ms=100,
        )

        call_args = mock_producer_class.call_args[0][0]
        assert call_args['compression_type'] == 'gzip'
        assert call_args['batch_num_messages'] == 1000
        assert call_args['linger_ms'] == 100

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_kafka_exception_during_init(self, mock_callback, mock_producer_class):
        """Test handling of KafkaException during initialization."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer_class.side_effect = KafkaException("Connection failed")

        with pytest.raises(KafkaConnectionError) as exc_info:
            GraceDbKafkaProducer("invalid:9092", instance_name="test", reload_cred=False)

        assert "Failed to connect to Kafka broker" in str(exc_info.value)


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerProduce:
    """Test producer produce method."""

    def setup_method(self):
        """Set up test producer."""
        self.mock_producer = mock.Mock()

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback', return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            self.producer = GraceDbKafkaProducer("localhost:9092", topic="test-topic", instance_name="test", reload_cred=False)

    def test_produce_dict_value(self):
        """Test producing a message with dict value."""
        event_data = {"graceid": "G12345", "group": "CBC"}

        self.producer.produce(event_data)

        # Check that producer.produce was called
        self.mock_producer.produce.assert_called_once()
        call_kwargs = self.mock_producer.produce.call_args[1]

        assert call_kwargs['topic'] == 'test.test-topic'
        assert call_kwargs['value'] == json.dumps(event_data).encode('utf-8')

    def test_produce_string_value(self):
        """Test producing a message with string value."""
        message = "Test message"

        self.producer.produce(message)

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['value'] == message.encode('utf-8')

    def test_produce_bytes_value(self):
        """Test producing a message with bytes value."""
        message = b"Test bytes"

        self.producer.produce(message)

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['value'] == message

    def test_produce_invalid_value_type(self):
        """Test producing a message with invalid value type."""
        with pytest.raises(TypeError) as exc_info:
            self.producer.produce(12345)

        assert "Value must be dict, str, or bytes" in str(exc_info.value)

    def test_produce_with_explicit_topic(self):
        """Test producing with explicit topic."""
        self.producer.produce("message", topic="other-topic")

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['topic'] == 'test.other-topic'

    def test_produce_without_default_topic(self):
        """Test producing without default topic raises error."""
        with mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback', return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            producer = GraceDbKafkaProducer("localhost:9092", instance_name="test", reload_cred=False)  # No default topic

        with pytest.raises(ValueError) as exc_info:
            producer.produce("message")

        assert "Topic must be specified" in str(exc_info.value)

    def test_produce_with_string_key(self):
        """Test producing with string key."""
        self.producer.produce("message", key="test-key")

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['key'] == b"test-key"

    def test_produce_with_bytes_key(self):
        """Test producing with bytes key."""
        self.producer.produce("message", key=b"test-key")

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['key'] == b"test-key"

    def test_produce_with_headers(self):
        """Test producing with headers."""
        headers = {"source": "gracedb", "type": "event"}

        self.producer.produce("message", headers=headers)

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['headers'] == headers

    def test_produce_with_partition(self):
        """Test producing to specific partition."""
        self.producer.produce("message", partition=2)

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['partition'] == 2

    def test_produce_with_custom_callback(self):
        """Test producing with custom delivery callback.

        The custom callback is wrapped in a tracking callback that
        decrements _pending_messages, then calls the user's callback.
        """
        custom_callback = mock.Mock()

        self.producer.produce("message", on_delivery=custom_callback)

        # The on_delivery passed to the underlying producer is the
        # tracking wrapper, not the custom callback directly.
        call_kwargs = self.mock_producer.produce.call_args[1]
        wrapper = call_kwargs['on_delivery']
        assert wrapper is not custom_callback

        # Invoke the wrapper and verify it calls the custom callback
        self.producer._pending_messages = 1
        mock_msg = mock.Mock()
        wrapper(None, mock_msg)
        custom_callback.assert_called_once_with(None, mock_msg)
        assert self.producer.pending_messages == 0

    def test_produce_increments_pending_counter(self):
        """Test that produce increments pending message counter."""
        assert self.producer.pending_messages == 0

        self.producer.produce("message1")
        assert self.producer.pending_messages == 1

        self.producer.produce("message2")
        assert self.producer.pending_messages == 2

    def test_produce_handles_buffer_error(self):
        """Test that BufferError triggers flush and retry."""
        # First call raises BufferError, second succeeds
        self.mock_producer.produce.side_effect = [BufferError(), None]
        self.mock_producer.flush.return_value = 0

        self.producer.produce("message")

        # Should have called produce twice (original + retry)
        assert self.mock_producer.produce.call_count == 2
        # Should have called flush once
        self.mock_producer.flush.assert_called_once()

    def test_produce_handles_kafka_exception(self):
        """Test that KafkaException is wrapped properly."""
        self.mock_producer.produce.side_effect = KafkaException("Production failed")

        with pytest.raises(KafkaProducerError) as exc_info:
            self.producer.produce("message")

        assert "Message production failed" in str(exc_info.value)

    def test_produce_calls_poll(self):
        """Test that produce calls poll for delivery callbacks."""
        self.producer.produce("message")

        # Should call poll(0) for non-blocking callback processing
        self.mock_producer.poll.assert_called_with(0)


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerFlush:
    """Test producer flush method."""

    def setup_method(self):
        """Set up test producer."""
        self.mock_producer = mock.Mock()

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback', return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            self.producer = GraceDbKafkaProducer("localhost:9092", topic="test-topic", instance_name="test", reload_cred=False)

    def test_successful_flush(self):
        """Test successful flush with no remaining messages."""
        self.mock_producer.flush.return_value = 0

        result = self.producer.flush()

        assert result == 0
        self.mock_producer.flush.assert_called_once_with(-1)

    def test_flush_with_timeout(self):
        """Test flush with custom timeout."""
        self.mock_producer.flush.return_value = 0

        self.producer.flush(timeout=10)

        self.mock_producer.flush.assert_called_once_with(10)

    def test_flush_with_remaining_messages(self):
        """Test flush that times out with remaining messages."""
        self.mock_producer.flush.return_value = 5

        result = self.producer.flush()

        assert result == 5

    def test_flush_handles_kafka_exception(self):
        """Test that KafkaException during flush is wrapped."""
        self.mock_producer.flush.side_effect = KafkaException("Flush failed")

        with pytest.raises(KafkaProducerError) as exc_info:
            self.producer.flush()

        assert "Flush failed" in str(exc_info.value)

    def test_flush_resets_pending_counter_on_success(self):
        """Test that successful flush resets pending counter."""
        self.producer._pending_messages = 10
        self.mock_producer.flush.return_value = 0

        self.producer.flush()

        assert self.producer.pending_messages == 0

    def test_flush_doesnt_reset_counter_with_remaining(self):
        """Test that flush doesn't reset counter if messages remain."""
        self.producer._pending_messages = 10
        self.mock_producer.flush.return_value = 3

        self.producer.flush()

        assert self.producer.pending_messages == 10  # Unchanged


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerClose:
    """Test producer close method."""

    def setup_method(self):
        """Set up test producer."""
        self.mock_producer = mock.Mock()

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback', return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            self.producer = GraceDbKafkaProducer("localhost:9092", topic="test-topic", instance_name="test", reload_cred=False)

    def test_close_calls_flush(self):
        """Test that close calls flush."""
        self.mock_producer.flush.return_value = 0

        self.producer.close()

        self.mock_producer.flush.assert_called_once_with(30)

    def test_close_with_custom_timeout(self):
        """Test close with custom timeout."""
        self.mock_producer.flush.return_value = 0

        self.producer.close(timeout=60)

        self.mock_producer.flush.assert_called_once_with(60)

    def test_close_with_remaining_messages(self):
        """Test close when messages remain after flush."""
        self.mock_producer.flush.return_value = 5

        # Should not raise exception
        self.producer.close()

    def test_close_propagates_exceptions(self):
        """Test that exceptions during close are propagated."""
        self.mock_producer.flush.side_effect = RuntimeError("Flush error")

        with pytest.raises(RuntimeError):
            self.producer.close()


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerContextManager:
    """Test producer context manager functionality."""

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_context_manager_entry(self, mock_callback, mock_producer_class):
        """Test context manager entry returns producer."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer = mock.Mock()
        mock_producer.flush.return_value = 0
        mock_producer_class.return_value = mock_producer

        with GraceDbKafkaProducer("localhost:9092", topic="test", instance_name="test", reload_cred=False) as producer:
            assert isinstance(producer, GraceDbKafkaProducer)

    @mock.patch('ligo.gracedb.kafka.producer.Producer')
    @mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback')
    def test_context_manager_exit_calls_close(self, mock_callback, mock_producer_class):
        """Test context manager exit calls close."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        mock_callback.return_value = lambda x: ("token", 1000)
        mock_producer = mock.Mock()
        mock_producer.flush.return_value = 0
        mock_producer_class.return_value = mock_producer

        with GraceDbKafkaProducer("localhost:9092", topic="test", instance_name="test", reload_cred=False) as producer:
            pass

        # Should have called flush (part of close)
        mock_producer.flush.assert_called()


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerDeliveryCallback:
    """Test producer delivery callback."""

    def setup_method(self):
        """Set up test producer."""
        self.mock_producer = mock.Mock()

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback', return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            self.producer = GraceDbKafkaProducer("localhost:9092", topic="test-topic", instance_name="test", reload_cred=False)

    def test_default_callback_on_success(self):
        """Test default delivery callback on successful delivery.

        Note: _default_delivery_callback no longer decrements
        _pending_messages; the tracking wrapper in produce() handles
        that for both default and custom callbacks.
        """
        mock_msg = mock.Mock()
        mock_msg.topic.return_value = "test-topic"
        mock_msg.partition.return_value = 0
        mock_msg.offset.return_value = 123

        self.producer._pending_messages = 5
        self.producer._default_delivery_callback(None, mock_msg)

        assert self.producer.pending_messages == 5

    def test_default_callback_on_error(self):
        """Test default delivery callback on delivery error."""
        mock_msg = mock.Mock()
        mock_msg.topic.return_value = "test-topic"

        self.producer._pending_messages = 5
        self.producer._default_delivery_callback("Delivery error", mock_msg)

        assert self.producer.pending_messages == 5

    def test_default_callback_doesnt_go_negative(self):
        """Test that pending counter doesn't go negative."""
        mock_msg = mock.Mock()
        mock_msg.topic.return_value = "test-topic"
        mock_msg.partition.return_value = 0
        mock_msg.offset.return_value = 123

        self.producer._pending_messages = 0
        self.producer._default_delivery_callback(None, mock_msg)

        assert self.producer.pending_messages == 0


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerProperties:
    """Test producer properties."""

    def setup_method(self):
        """Set up test producer."""
        self.mock_producer = mock.Mock()

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback', return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            self.producer = GraceDbKafkaProducer("localhost:9092", topic="test-topic", instance_name="test", reload_cred=False)

    def test_pending_messages_property(self):
        """Test pending_messages property."""
        assert self.producer.pending_messages == 0

        self.producer._pending_messages = 10
        assert self.producer.pending_messages == 10


class TestProducerDecodeJwt:
    """Test the static _decode_jwt_claims method."""

    def test_decode_valid_jwt(self):
        """Test decoding a JWT with known claims."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        claims = {"sub": "alice", "scope": "kafka.produce"}
        token = _make_jwt(claims)
        result = GraceDbKafkaProducer._decode_jwt_claims(token)

        assert result["sub"] == "alice"
        assert result["scope"] == "kafka.produce"

    def test_decode_jwt_with_padding(self):
        """Test decoding a JWT whose payload needs base64url padding."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        # Use a payload whose base64url encoding length is not a multiple of 4
        claims = {"x": "a"}
        token = _make_jwt(claims)
        result = GraceDbKafkaProducer._decode_jwt_claims(token)

        assert result["x"] == "a"

    def test_decode_jwt_missing_claims(self):
        """Test decoding a JWT that has no sub or scope claims."""
        from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

        claims = {"iss": "https://example.com", "aud": "kafka"}
        token = _make_jwt(claims)
        result = GraceDbKafkaProducer._decode_jwt_claims(token)

        assert "sub" not in result
        assert "scope" not in result
        assert result["iss"] == "https://example.com"


@mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True)
class TestProducerPing:
    """Test producer ping method."""

    def setup_method(self):
        """Set up test producer with mocked internals."""
        self.mock_producer = mock.Mock()

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback',
                        return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer

            self.producer = GraceDbKafkaProducer(
                "localhost:9092", topic="default-topic",
                instance_name="test",
                reload_cred=False,
            )

        self.test_claims = {"sub": "alice", "scope": "kafka.produce"}
        self.test_token = _make_jwt(self.test_claims)

    def _simulate_successful_delivery(self):
        """Configure mock producer to simulate a successful delivery.

        Captures the on_delivery callback from produce() and invokes it
        with a mock message when flush() is called.
        """
        captured = {}

        def capture_produce(**kwargs):
            captured['callback'] = kwargs.get('on_delivery')

        self.mock_producer.produce.side_effect = capture_produce

        def invoke_callback(timeout=None):
            mock_msg = mock.Mock()
            mock_msg.topic.return_value = "the-topic"
            mock_msg.partition.return_value = 1
            mock_msg.offset.return_value = 42
            captured['callback'](None, mock_msg)
            return 0

        self.mock_producer.flush.side_effect = invoke_callback

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_successful(self, mock_get_token):
        """Test full success path: returns dict with delivery details."""
        mock_get_token.return_value = self.test_token
        self._simulate_successful_delivery()

        result = self.producer.ping(topic="the-topic")

        assert result['topic'] == "the-topic"
        assert result['partition'] == 1
        assert result['offset'] == 42
        assert 'message' in result

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_message_contains_identity_claims(self, mock_get_token):
        """Test that the ping message includes sub and scope from the token."""
        mock_get_token.return_value = self.test_token
        self._simulate_successful_delivery()

        result = self.producer.ping(topic="t")

        assert result['message']['sub'] == "alice"
        assert result['message']['scope'] == "kafka.produce"
        assert result['message']['type'] == "ping"
        assert result['message']['client'] == "ligo-gracedb"

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_uses_explicit_topic(self, mock_get_token):
        """Test that an explicit topic argument is namespaced."""
        mock_get_token.return_value = self.test_token
        self._simulate_successful_delivery()

        self.producer.ping(topic="custom-topic")

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['topic'] == "test.custom-topic"

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_uses_default_topic(self, mock_get_token):
        """Test that the producer's default_topic is used when no topic arg."""
        mock_get_token.return_value = self.test_token
        self._simulate_successful_delivery()

        self.producer.ping()

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['topic'] == "test.default-topic"

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_falls_back_to_instance_test(self, mock_get_token):
        """Test fallback to '{instance}.test' when no topic is configured."""
        mock_get_token.return_value = self.test_token
        self._simulate_successful_delivery()

        self.producer.default_topic = None
        self.producer.ping()

        call_kwargs = self.mock_producer.produce.call_args[1]
        assert call_kwargs['topic'] == "test.test"

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_passes_token_scope_and_audience(self, mock_get_token):
        """Test that get_scitoken_for_kafka is called with producer's
        token_scope and token_audience."""
        mock_get_token.return_value = self.test_token
        self._simulate_successful_delivery()

        self.producer.token_scope = "write"
        self.producer.token_audience = "my-audience"
        self.producer.ping(topic="t")

        mock_get_token.assert_called_once_with(
            "localhost:9092",
            scope="write",
            audience="my-audience",
        )

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_no_scitoken_available(self, mock_get_token):
        """Test that KafkaTokenError from get_scitoken_for_kafka propagates."""
        mock_get_token.side_effect = KafkaTokenError("No token found")

        with pytest.raises(KafkaTokenError, match="No token found"):
            self.producer.ping(topic="t")

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_broker_unreachable(self, mock_get_token):
        """Test that KafkaException from produce raises KafkaConnectionError."""
        mock_get_token.return_value = self.test_token
        self.mock_producer.produce.side_effect = KafkaException(
            "broker unavailable"
        )

        with pytest.raises(KafkaConnectionError, match="Ping failed"):
            self.producer.ping(topic="t")

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_flush_raises_kafka_exception(self, mock_get_token):
        """Test that KafkaException from flush raises KafkaConnectionError."""
        mock_get_token.return_value = self.test_token
        self.mock_producer.flush.side_effect = KafkaException("flush error")

        with pytest.raises(KafkaConnectionError, match="Ping failed"):
            self.producer.ping(topic="t")

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_delivery_error(self, mock_get_token):
        """Test that a delivery error in the callback raises
        KafkaProducerError."""
        mock_get_token.return_value = self.test_token

        captured = {}

        def capture_produce(**kwargs):
            captured['callback'] = kwargs.get('on_delivery')

        self.mock_producer.produce.side_effect = capture_produce

        def invoke_callback_with_error(timeout=None):
            captured['callback']("delivery failed", None)
            return 0

        self.mock_producer.flush.side_effect = invoke_callback_with_error

        with pytest.raises(KafkaProducerError, match="Ping delivery failed"):
            self.producer.ping(topic="t")

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_timeout(self, mock_get_token):
        """Test that remaining > 0 after flush raises KafkaProducerError
        with 'timed out'."""
        mock_get_token.return_value = self.test_token
        self.mock_producer.flush.return_value = 1

        with pytest.raises(KafkaProducerError, match="timed out"):
            self.producer.ping(topic="t")

    @mock.patch('ligo.gracedb.kafka.producer.get_scitoken_for_kafka')
    def test_ping_custom_timeout(self, mock_get_token):
        """Test that timeout argument is forwarded to flush."""
        mock_get_token.return_value = self.test_token
        self._simulate_successful_delivery()

        self.producer.ping(topic="t", timeout=30)

        self.mock_producer.flush.assert_called_once_with(timeout=30)


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

    def setup_method(self):
        """Set up test producer."""
        self.mock_producer = mock.Mock()

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', return_value=self.mock_producer), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback', return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            self.producer = GraceDbKafkaProducer("localhost:9092", topic="test-topic", instance_name="test", reload_cred=False)

    def _make_auth_error(self):
        """Create a mock Kafka error with _AUTHENTICATION code."""
        from ligo.gracedb.kafka.producer 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."""
        error = self._make_auth_error()

        self.producer._error_callback(error)

        assert self.producer._auth_error is error

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

        self.producer._error_callback(error)

        assert self.producer._auth_error is None

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

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

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

    def test_produce_raises_on_prior_auth_error(self):
        """produce() raises KafkaAuthenticationError if auth already failed."""
        self.producer._auth_error = self._make_auth_error()

        with pytest.raises(KafkaAuthenticationError):
            self.producer.produce("message")

        # Should not have attempted to produce
        self.mock_producer.produce.assert_not_called()

    def test_flush_raises_on_prior_auth_error(self):
        """flush() raises KafkaAuthenticationError if auth already failed."""
        self.producer._auth_error = self._make_auth_error()

        with pytest.raises(KafkaAuthenticationError):
            self.producer.flush()

        # Should not have attempted to flush
        self.mock_producer.flush.assert_not_called()

    def test_flush_raises_on_auth_error_during_flush(self):
        """flush() raises if auth error arrives while flushing."""
        def set_auth_error(timeout):
            self.producer._auth_error = self._make_auth_error()
            return 0

        self.mock_producer.flush.side_effect = set_auth_error

        with pytest.raises(KafkaAuthenticationError):
            self.producer.flush()

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

        with mock.patch('ligo.gracedb.kafka.producer.KAFKA_AVAILABLE', True), \
             mock.patch('ligo.gracedb.kafka.producer.Producer', mock_producer_cls), \
             mock.patch('ligo.gracedb.kafka.producer.create_oauth_callback',
                        return_value=lambda x: ("token", 1000)):
            from ligo.gracedb.kafka.producer import GraceDbKafkaProducer
            producer = GraceDbKafkaProducer("localhost:9092", instance_name="test", reload_cred=False)

        config = mock_producer_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 producer._auth_error is error
