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

import base64
import json
import logging
import time
import warnings

try:
    from confluent_kafka import Producer, KafkaException
    from confluent_kafka import KafkaError as _KafkaError
    KAFKA_AVAILABLE = True
    _AUTHENTICATION = _KafkaError._AUTHENTICATION
except ImportError:
    Producer = None
    KafkaException = type('KafkaException', (Exception,), {})
    _KafkaError = None
    KAFKA_AVAILABLE = False
    _AUTHENTICATION = 58  # confluent-kafka error code for SASL auth failure

from .auth import (
    create_oauth_callback,
    get_scitoken_for_kafka,
    SciTokenRefreshManager,
)
from .exceptions import (
    KafkaAuthenticationError,
    KafkaProducerError,
    KafkaConnectionError,
    KafkaMessageDeliveryError,
)
from .utils import derive_instance_name, make_topic_name

log = logging.getLogger(__name__)


class GraceDbKafkaProducer:
    """
    Kafka producer with SciToken authentication for GraceDB events.

    This producer automatically handles SciToken discovery and OAuth
    authentication for publishing messages to Kafka brokers.

    Args:
        bootstrap_servers (str): Kafka broker address (e.g., "localhost:9092")
        topic (str, optional): Default topic name for messages (pipeline name,
            e.g., ``"gstlal"``).  Will be automatically namespaced with the
            instance name to produce the actual Kafka topic
            (e.g., ``"gracedb-dev.gstlal"``).  When omitted, a topic must be
            specified at :meth:`produce` time.  When using the producer with
            :meth:`~ligo.gracedb.rest.GraceDb.create_event`, the topic is
            determined automatically from the event's ``group`` and
            ``pipeline`` arguments, so setting a default here is not required.
        instance_name (str, optional): GraceDB instance name used to namespace
            topics (e.g., ``"gracedb-dev"``).  Exactly one of *instance_name*
            or *service_url* must be provided.
        service_url (str, optional): GraceDB service URL from which the
            instance name is derived (e.g.,
            ``"https://gracedb-dev.ligo.org/api/"``).
        use_ssl (bool): Enable SSL/TLS encryption (default: True)
        ca_cert_path (str, optional): Path to CA certificate for SSL verification
        skip_hostname_verification (bool): Skip SSL hostname verification (default: False).
            Use for development/testing when cert CN doesn't match broker hostname.
        token_scope (str): Required SciToken scope (default: "kafka.produce")
        token_audience (str, optional): Required audience. Defaults to broker hostname
        reload_cred (bool): Proactively reload SciTokens before expiry (default: True).
            When enabled, a ``SciTokenRefreshManager`` monitors token expiration and
            re-reads the token from disk before it expires.
        reload_buffer (int): Seconds before token expiry to trigger a reload
            (default: 300).  Only used when ``reload_cred=True``.
        flush_timeout (int): Seconds to wait for message delivery when
            :meth:`~ligo.gracedb.rest.GraceDb.create_event` flushes the
            producer after publishing (default: 10).
        guaranteed_delivery (bool): When ``True`` (default), configures the
            producer for reliable delivery: ``acks=all``, retries on
            transient failures, and delivery verification via callback.
            Set to ``False`` for higher throughput at the risk of
            silent message loss.
        **kafka_config: Additional confluent_kafka.Producer configuration options

    Example:

        Using with ``create_event`` (topic determined automatically)::

            producer = GraceDbKafkaProducer(
                bootstrap_servers="gracedb-dev.ligo.org:9092",
                service_url="https://gracedb-dev.ligo.org/api/",
                ca_cert_path="/tmp/cacert.pem",
            )
            gdb.create_event(
                group="CBC", pipeline="gstlal", filename="coinc.xml",
                filecontents=data, kafka=producer,
            )
            # Message is published to gracedb-dev.gstlal

        Using with a default topic for direct ``produce()`` calls::

            producer = GraceDbKafkaProducer(
                bootstrap_servers="gracedb-dev.ligo.org:9092",
                topic="test",
                service_url="https://gracedb-dev.ligo.org/api/",
                ca_cert_path="/tmp/cacert.pem",
            )
            producer.produce(value={"graceid": "G12345"})
            # Message is published to gracedb-dev.test
    """

    def __init__(self, bootstrap_servers, topic=None,
                 instance_name=None, service_url=None,
                 use_ssl=True, ca_cert_path=None,
                 skip_hostname_verification=False,
                 token_scope="kafka.produce", token_audience=None,
                 reload_cred=True, reload_buffer=300,
                 flush_timeout=10,
                 guaranteed_delivery=True,
                 auto_refresh_token=None, **kafka_config):

        if not KAFKA_AVAILABLE:
            raise ImportError(
                "confluent-kafka is required for Kafka functionality. "
                "Install it with: pip install ligo-gracedb[kafka]"
            )

        # Handle deprecated auto_refresh_token parameter
        if auto_refresh_token is not None:
            warnings.warn(
                "auto_refresh_token is deprecated; use reload_cred instead",
                DeprecationWarning,
                stacklevel=2,
            )
            reload_cred = auto_refresh_token

        # Derive instance name for topic namespacing
        if instance_name is not None:
            self.instance_name = instance_name
        elif service_url is not None:
            self.instance_name = derive_instance_name(service_url)
        else:
            raise ValueError(
                "instance_name or service_url is required for topic namespacing"
            )

        self.bootstrap_servers = bootstrap_servers
        self.default_topic = (
            make_topic_name(self.instance_name, topic) if topic else None
        )
        self.token_scope = token_scope
        self.token_audience = token_audience
        self.flush_timeout = flush_timeout

        log.info(f"Initializing GraceDbKafkaProducer for {bootstrap_servers}")

        # Build producer configuration
        config = {
            'bootstrap.servers': bootstrap_servers,
            'client.id': 'ligo-gracedb-producer',
            'message.max.bytes': 10485760,
            'compression.type': 'zstd',
        }

        if guaranteed_delivery:
            # Require broker acknowledgment and retry transient failures.
            config.update({
                'acks': 'all',
                'retries': 5,
                'retry.backoff.ms': 200,
                'delivery.timeout.ms': 30000,
                'request.timeout.ms': 10000,
            })

        # Configure security protocol based on SSL setting
        if use_ssl:
            config['security.protocol'] = 'SASL_SSL'
            config['sasl.mechanism'] = 'OAUTHBEARER'

            # Add SSL CA certificate if provided
            if ca_cert_path:
                config['ssl.ca.location'] = ca_cert_path
                log.debug(f"Using CA certificate: {ca_cert_path}")

            # Skip hostname verification if requested (for dev/testing)
            if skip_hostname_verification:
                config['ssl.endpoint.identification.algorithm'] = 'none'
                log.warning("SSL hostname verification disabled - not recommended for production")
        else:
            # Plaintext with OAUTHBEARER (not recommended for production)
            config['security.protocol'] = 'SASL_PLAINTEXT'
            config['sasl.mechanism'] = 'OAUTHBEARER'
            log.warning("Using plaintext connection (not recommended for production)")

        # Create OAuth callback for SciToken authentication
        if reload_cred:
            # Use refresh manager for proactive token renewal
            self.refresh_manager = SciTokenRefreshManager(
                broker_url=bootstrap_servers,
                scope=token_scope,
                audience=token_audience,
                refresh_buffer=reload_buffer,
            )
            config['oauth_cb'] = self.refresh_manager.create_callback()
            log.debug("Using automatic token refresh")
        else:
            # Simple callback without proactive refresh
            config['oauth_cb'] = create_oauth_callback(
                broker_url=bootstrap_servers,
                scope=token_scope,
                audience=token_audience,
            )
            self.refresh_manager = None

        # Error callback for detecting authentication failures
        self._auth_error = None
        config['error_cb'] = self._error_callback

        # Merge user-provided Kafka configuration
        config.update(kafka_config)

        # Create producer
        try:
            self.producer = Producer(config)
            log.info("Kafka producer created successfully")
        except KafkaException as e:
            log.error(f"Failed to create Kafka producer: {e}")
            raise KafkaConnectionError(f"Failed to connect to Kafka broker: {e}") from e

        # Track pending messages for graceful shutdown
        self._pending_messages = 0

    def _error_callback(self, error):
        """Handle Kafka errors from librdkafka.

        Detects authentication failures (e.g. expired SciToken) and
        stores them so the next produce/flush call raises a clean
        Python exception instead of spinning on retries.
        """
        if error.code() == _AUTHENTICATION:
            self._auth_error = error
            log.error(f"Authentication failed (token expired?): {error}")
        else:
            log.warning(f"Kafka error: {error}")

    def _check_auth(self):
        """Raise KafkaAuthenticationError if an auth failure was detected."""
        if self._auth_error is not None:
            raise KafkaAuthenticationError(
                f"SciToken authentication failed: {self._auth_error}"
            )

    def produce(self, value, topic=None, key=None, headers=None,
                partition=None, on_delivery=None):
        """
        Produce a message to Kafka.

        Args:
            value: Message value. Can be:
                - dict: Automatically serialized to JSON
                - str: Encoded to UTF-8 bytes
                - bytes: Sent as-is
            topic (str, optional): Pipeline or topic name (e.g., ``"gstlal"``).
                Automatically namespaced with the instance name (e.g.,
                ``"gracedb-dev.gstlal"``).  Falls back to the default topic
                set at construction if not specified.
            key (str or bytes, optional): Message key for partitioning
            headers (dict, optional): Message headers as key-value pairs
            partition (int, optional): Specific partition to send to
            on_delivery (callable, optional): Callback for delivery reports
                Signature: on_delivery(err, msg)

        Raises:
            ValueError: If topic is not specified and no default was set
            KafkaProducerError: If message production fails

        Example:
            >>> producer.produce(
            ...     value={"graceid": "G12345"},
            ...     topic="gstlal",
            ...     key="G12345",
            ... )
        """
        self._check_auth()

        # Determine topic — namespace if caller passed a raw pipeline name
        if topic is not None:
            topic = make_topic_name(self.instance_name, topic)
        else:
            topic = self.default_topic
        if topic is None:
            raise ValueError("Topic must be specified (no default topic configured)")

        # Serialize value based on type
        if isinstance(value, dict):
            value = json.dumps(value).encode('utf-8')
        elif isinstance(value, str):
            value = value.encode('utf-8')
        elif not isinstance(value, bytes):
            raise TypeError(
                f"Value must be dict, str, or bytes, got {type(value).__name__}"
            )

        # Encode key if it's a string
        if isinstance(key, str):
            key = key.encode('utf-8')

        # Wrap callback to always track pending message count
        def _tracking_callback(err, msg):
            if self._pending_messages > 0:
                self._pending_messages -= 1
            if on_delivery is not None:
                on_delivery(err, msg)
            else:
                self._default_delivery_callback(err, msg)

        callback = _tracking_callback

        # Build produce kwargs - only include non-None optional parameters
        # (confluent-kafka doesn't accept None for partition/headers)
        produce_kwargs = {
            'topic': topic,
            'value': value,
            'on_delivery': callback
        }
        if key is not None:
            produce_kwargs['key'] = key
        if headers is not None:
            produce_kwargs['headers'] = headers
        if partition is not None:
            produce_kwargs['partition'] = partition

        try:
            # Produce message
            self.producer.produce(**produce_kwargs)

            self._pending_messages += 1

            # Poll to handle delivery callbacks
            # Non-blocking poll to avoid blocking the caller
            self.producer.poll(0)

            log.debug(f"Message queued for topic '{topic}'")

        except BufferError:
            # Queue is full, need to flush
            log.warning("Kafka producer queue full, flushing...")
            self.flush()
            # Retry once
            self.producer.produce(**produce_kwargs)
            self._pending_messages += 1
        except KafkaException as e:
            log.error(f"Failed to produce message: {e}")
            raise KafkaProducerError(f"Message production failed: {e}") from e

    def flush(self, timeout=-1):
        """
        Wait for all messages in the producer queue to be delivered.

        Args:
            timeout (float): Maximum time to wait in seconds.
                -1 = wait indefinitely (default)
                0 = non-blocking check
                >0 = wait up to timeout seconds

        Returns:
            int: Number of messages still in queue after timeout

        Raises:
            KafkaProducerError: If flush encounters an error
        """
        log.debug("Flushing Kafka producer queue...")
        self._check_auth()

        try:
            remaining = self.producer.flush(timeout)
            self._check_auth()

            if remaining > 0:
                log.warning(f"{remaining} messages still in queue after flush")
            else:
                log.debug("All messages flushed successfully")
                self._pending_messages = 0

            return remaining

        except KafkaException as e:
            log.error(f"Error during flush: {e}")
            raise KafkaProducerError(f"Flush failed: {e}") from e

    def close(self, timeout=30):
        """
        Close the producer and release resources.

        This method flushes any pending messages before closing.

        Args:
            timeout (float): Maximum time to wait for flush in seconds
        """
        log.info("Closing Kafka producer...")

        try:
            # Flush pending messages
            remaining = self.flush(timeout)

            if remaining > 0:
                log.warning(
                    f"Producer closed with {remaining} messages still pending"
                )

            log.info("Kafka producer closed")

        except Exception as e:
            log.error(f"Error closing producer: {e}")
            raise

    def _default_delivery_callback(self, err, msg):
        """
        Default delivery report callback.

        Args:
            err: Delivery error (None if successful)
            msg: Message object
        """
        if err:
            log.error(
                f"Message delivery failed for topic '{msg.topic()}': {err}"
            )
        else:
            log.debug(
                f"Message delivered to {msg.topic()} "
                f"[partition {msg.partition()}, offset {msg.offset()}]"
            )

    @staticmethod
    def _decode_jwt_claims(token_str):
        """Decode the payload of a JWT without signature verification."""
        payload = token_str.split('.')[1]
        # Restore base64url padding
        payload += '=' * (4 - len(payload) % 4)
        return json.loads(base64.urlsafe_b64decode(payload))

    def ping(self, topic=None, timeout=10):
        """
        Send a test message to verify broker connectivity and authentication.

        Produces a small JSON message whose body includes the ``sub`` and
        ``scope`` claims from the SciToken used to authenticate, then waits
        for delivery confirmation.  This exercises the full path: token
        discovery, OAUTHBEARER handshake, SSL connection, and message
        delivery.

        Args:
            topic (str, optional): Topic to ping. Uses default_topic,
                or ``'gracedb-test'`` if no default is set.
            timeout (float): Seconds to wait for delivery confirmation
                (default: 10).

        Returns:
            dict: Delivery details with keys ``'topic'``, ``'partition'``,
            ``'offset'``, and ``'message'`` (the test payload that was sent).

        Raises:
            KafkaProducerError: If the message cannot be delivered.
            KafkaConnectionError: If the broker is unreachable.
        """
        if topic is not None:
            topic = make_topic_name(self.instance_name, topic)
        else:
            topic = self.default_topic or make_topic_name(
                self.instance_name, 'test'
            )

        # Obtain the current token and extract identity claims
        token_str = get_scitoken_for_kafka(
            self.bootstrap_servers,
            scope=self.token_scope,
            audience=self.token_audience,
        )
        claims = self._decode_jwt_claims(token_str)

        test_message = {
            'type': 'ping',
            'timestamp': time.time(),
            'client': 'ligo-gracedb',
            'sub': claims.get('sub'),
            'scope': claims.get('scope'),
        }

        # Use a mutable container so the callback closure can write results
        result = {}
        error = [None]

        def _ping_callback(err, msg):
            if err:
                error[0] = err
            else:
                result['topic'] = msg.topic()
                result['partition'] = msg.partition()
                result['offset'] = msg.offset()

        try:
            self.producer.produce(
                topic=topic,
                value=json.dumps(test_message).encode('utf-8'),
                on_delivery=_ping_callback,
            )
            remaining = self.producer.flush(timeout=timeout)
        except KafkaException as e:
            raise KafkaConnectionError(f"Ping failed: {e}") from e

        if error[0]:
            raise KafkaProducerError(f"Ping delivery failed: {error[0]}")

        if remaining > 0:
            raise KafkaProducerError(
                "Ping timed out waiting for delivery confirmation"
            )

        result['message'] = test_message
        log.info(
            f"Ping successful: topic={result['topic']}, "
            f"partition={result['partition']}, "
            f"offset={result['offset']}, "
            f"sub={test_message['sub']}"
        )
        return result

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - automatically close producer."""
        self.close()
        return False

    @property
    def pending_messages(self):
        """Number of messages currently pending delivery."""
        return self._pending_messages
