# -*- 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/>.
"""
Integration tests for Kafka producer with SciToken authentication.

These tests require:
- A running Kafka broker with SciToken authentication enabled
- A valid SciToken with appropriate scopes
- Network connectivity to the broker

Set environment variables to configure:
- KAFKA_BROKER: Broker address (default: localhost:9092)
- KAFKA_TOPIC: Topic name (default: test-events)
- KAFKA_CA_CERT: Path to CA certificate
- SCITOKEN_FILE: Path to SciToken file
"""

import os
import json
import time
import pytest
from unittest import mock

# Mark all tests in this module as integration tests
pytestmark = pytest.mark.integration


def pytest_configure(config):
    """Register custom test markers."""
    config.addinivalue_line(
        "markers",
        "integration: Integration tests requiring external services",
    )


@pytest.fixture
def kafka_config():
    """Get Kafka configuration from environment."""
    return {
        'broker': os.getenv('KAFKA_BROKER', 'localhost:9092'),
        'topic': os.getenv('KAFKA_TOPIC', 'test-events'),
        'ca_cert': os.getenv('KAFKA_CA_CERT'),
        'scope': os.getenv('KAFKA_SCOPE', 'kafka.produce'),
    }


@pytest.fixture
def skip_if_no_broker(kafka_config):
    """Skip test if Kafka broker is not available."""
    import socket

    host, port = kafka_config['broker'].split(':')
    port = int(port)

    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(2)
        result = sock.connect_ex((host, port))
        sock.close()

        if result != 0:
            pytest.skip(f"Kafka broker not available at {kafka_config['broker']}")
    except Exception as e:
        pytest.skip(f"Cannot check broker availability: {e}")


@pytest.fixture
def skip_if_no_token():
    """Skip test if SciToken is not available."""
    try:
        from igwn_auth_utils.scitokens import find_token
        token = find_token(audience="localhost", scope="kafka.produce")
        if token is None:
            pytest.skip("No valid SciToken found. Run: htgettoken ...")
    except ImportError:
        pytest.skip("igwn-auth-utils not installed")
    except Exception as e:
        pytest.skip(f"Cannot find SciToken: {e}")


class TestKafkaProducerIntegration:
    """Integration tests for Kafka producer."""

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_producer_creation(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test creating a Kafka producer with SciToken auth."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        producer = GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert'],
            token_scope=kafka_config['scope']
        )

        assert producer is not None
        assert producer.bootstrap_servers == kafka_config['broker']
        assert producer.default_topic == kafka_config['topic']

        producer.close()

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_produce_single_message(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test producing a single message to Kafka."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        test_event = {
            'graceid': 'TEST123',
            'group': 'Test',
            'pipeline': 'integration-test',
            'created': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
        }

        with GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert']
        ) as producer:
            # Produce message
            producer.produce(test_event, key='TEST123')

            # Verify pending messages
            assert producer.pending_messages >= 0

            # Flush to ensure delivery
            remaining = producer.flush(timeout=10)
            assert remaining == 0

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_produce_multiple_messages(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test producing multiple messages to Kafka."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        events = [
            {'graceid': f'TEST{i}', 'group': 'Test', 'index': i}
            for i in range(10)
        ]

        with GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert']
        ) as producer:
            for event in events:
                producer.produce(event, key=event['graceid'])

            # Flush all messages
            remaining = producer.flush(timeout=30)
            assert remaining == 0

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_produce_with_delivery_callback(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test producing messages with delivery callbacks."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        delivered_messages = []
        failed_messages = []

        def on_delivery(err, msg):
            if err:
                failed_messages.append((err, msg))
            else:
                delivered_messages.append(msg)

        test_event = {'graceid': 'TEST_CALLBACK', 'group': 'Test'}

        with GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert']
        ) as producer:
            producer.produce(test_event, key='TEST_CALLBACK', on_delivery=on_delivery)
            producer.flush(timeout=10)

        # Verify delivery callback was called
        assert len(delivered_messages) == 1 or len(failed_messages) == 1
        if delivered_messages:
            assert delivered_messages[0].key() == b'TEST_CALLBACK'

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_produce_with_headers(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test producing messages with custom headers."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        test_event = {'graceid': 'TEST_HEADERS', 'group': 'Test'}
        headers = {
            'source': 'integration-test',
            'version': '1.0',
            'timestamp': str(int(time.time()))
        }

        with GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert']
        ) as producer:
            producer.produce(test_event, key='TEST_HEADERS', headers=headers)
            remaining = producer.flush(timeout=10)
            assert remaining == 0

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_token_refresh(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test automatic token refresh for long-running producer."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        # Create producer with auto-refresh enabled
        producer = GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert'],
            auto_refresh_token=True
        )

        assert producer.refresh_manager is not None

        # Produce messages over time to test refresh
        for i in range(3):
            event = {'graceid': f'TEST_REFRESH_{i}', 'iteration': i}
            producer.produce(event, key=f'TEST_REFRESH_{i}')
            producer.flush(timeout=5)
            time.sleep(1)  # Small delay between messages

        producer.close()


class TestKafkaAuthenticationIntegration:
    """Integration tests for SciToken authentication."""

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_token_discovery(self, skip_if_no_token):
        """Test SciToken discovery from file system."""
        from ligo.gracedb.kafka.auth import get_scitoken_for_kafka

        token = get_scitoken_for_kafka("localhost:9092", scope="kafka.produce")

        assert token is not None
        assert isinstance(token, str)
        assert len(token) > 0
        # JWT tokens have three parts separated by dots
        assert token.count('.') == 2

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_oauth_callback_creation(self, skip_if_no_token):
        """Test OAuth callback creation for Kafka."""
        from ligo.gracedb.kafka.auth import create_oauth_callback

        callback = create_oauth_callback("localhost:9092", scope="kafka.produce")

        assert callback is not None
        assert callable(callback)

        # Test calling the callback
        token, expiry = callback(None)

        assert token is not None
        assert isinstance(token, str)
        assert isinstance(expiry, float)
        assert expiry > time.time()

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_token_refresh_manager(self, skip_if_no_token):
        """Test SciToken refresh manager."""
        from ligo.gracedb.kafka.auth import SciTokenRefreshManager

        manager = SciTokenRefreshManager(
            broker_url="localhost:9092",
            scope="kafka.produce",
            refresh_buffer=300
        )

        # Get token
        token1 = manager.get_token()
        assert token1 is not None

        # Get token again - should use cache
        token2 = manager.get_token()
        assert token2 == token1

        # Create callback
        callback = manager.create_callback()
        assert callable(callback)


class TestKafkaErrorHandling:
    """Integration tests for error handling."""

    def test_invalid_broker_address(self):
        """Test handling of invalid broker address."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer, KafkaConnectionError

        # This should fail to connect but not raise immediately
        # (confluent-kafka connects lazily)
        try:
            producer = GraceDbKafkaProducer(
                bootstrap_servers="invalid-broker:9999",
                topic="test-topic"
            )
            # Attempting to produce may trigger connection error
            producer.produce({"test": "data"})
            producer.flush(timeout=5)
            producer.close()
        except Exception as e:
            # Expected to fail
            assert True

    def test_missing_topic(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test handling of missing topic."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        # Producer should create topic if auto.create.topics.enable=true
        # Otherwise it will queue messages until timeout
        with GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic="non-existent-topic-12345",
            ca_cert_path=kafka_config['ca_cert']
        ) as producer:
            producer.produce({"test": "data"}, key="test")
            # Flush may timeout if topic cannot be created
            producer.flush(timeout=5)

    def test_invalid_token_scope(self, kafka_config, skip_if_no_broker):
        """Test handling of invalid token scope."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer, KafkaTokenError

        # Requesting a scope that doesn't exist should fail
        try:
            producer = GraceDbKafkaProducer(
                bootstrap_servers=kafka_config['broker'],
                topic=kafka_config['topic'],
                ca_cert_path=kafka_config['ca_cert'],
                token_scope="invalid.scope.does.not.exist"
            )
            producer.close()
        except KafkaTokenError:
            # Expected to fail
            assert True
        except Exception:
            # May also fail for other reasons
            pass


class TestGraceDbClientIntegration:
    """Integration tests for GraceDB client with Kafka."""

    @pytest.mark.skipif(
        'KAFKA_INTEGRATION_TESTS' not in os.environ,
        reason="Integration tests disabled. Set KAFKA_INTEGRATION_TESTS=1 to enable."
    )
    def test_create_producer_from_client(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test creating Kafka producer from GraceDB client."""
        from unittest.mock import Mock

        # Create mock GraceDB client
        mock_client = Mock()
        mock_client.create_kafka_producer = Mock()

        # Import the real function
        from ligo.gracedb.rest import GraceDb
        real_client = GraceDb()

        # Test producer creation
        producer = real_client.create_kafka_producer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert']
        )

        assert producer is not None
        producer.close()


# Performance test (optional, can be slow)
class TestKafkaPerformance:
    """Performance tests for Kafka integration."""

    @pytest.mark.skipif(
        'KAFKA_PERFORMANCE_TESTS' not in os.environ,
        reason="Performance tests disabled. Set KAFKA_PERFORMANCE_TESTS=1 to enable."
    )
    def test_throughput(self, kafka_config, skip_if_no_broker, skip_if_no_token):
        """Test message throughput."""
        from ligo.gracedb.kafka import GraceDbKafkaProducer

        num_messages = 1000
        start_time = time.time()

        with GraceDbKafkaProducer(
            bootstrap_servers=kafka_config['broker'],
            topic=kafka_config['topic'],
            ca_cert_path=kafka_config['ca_cert']
        ) as producer:
            for i in range(num_messages):
                event = {'graceid': f'PERF_{i}', 'index': i}
                producer.produce(event, key=f'PERF_{i}')

            producer.flush(timeout=60)

        elapsed = time.time() - start_time
        throughput = num_messages / elapsed

        print(f"\nProduced {num_messages} messages in {elapsed:.2f} seconds")
        print(f"Throughput: {throughput:.2f} messages/second")

        # Expect at least 100 messages/second for simple events
        assert throughput > 10  # Conservative threshold
