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

import hashlib
import json
import logging
import os
import signal
import socket
import time
import warnings

try:
    from confluent_kafka import Consumer, KafkaException
    from confluent_kafka import KafkaError as _KafkaError
    KAFKA_AVAILABLE = True
    _AUTHENTICATION = _KafkaError._AUTHENTICATION
except ImportError:
    Consumer = 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,
    SciTokenRefreshManager,
)
from .exceptions import (
    KafkaAuthenticationError,
    KafkaConsumerError,
    KafkaConnectionError,
)
from .utils import derive_instance_name, make_topic_name

log = logging.getLogger(__name__)


class GraceDbKafkaConsumer:
    """
    Kafka consumer with SciToken authentication for GraceDB events.

    This consumer automatically handles SciToken discovery and OAuth
    authentication for subscribing to Kafka topics and reading messages.

    By default, the consumer is configured for **live streaming**: it
    generates a unique ephemeral group ID, starts reading from the end
    of the topic (``auto.offset.reset=latest``), and automatically
    commits offsets.  This means a fresh consumer only sees messages
    produced *after* it connects.

    To persist read position across restarts, supply a stable
    ``group_id``.  Kafka will track committed offsets for that group,
    so subsequent runs resume where the previous one left off.

    To replay all historical messages, pass
    ``auto_offset_reset="earliest"`` (useful with a new ``group_id``
    that has no committed offsets yet).

    Args:
        bootstrap_servers (str): Kafka broker address
            (e.g., "localhost:9092")
        topics (str or list, optional): Topic(s) to subscribe to
            immediately.  These should be pipeline names (e.g.,
            ``"gstlal"``); they are automatically namespaced with
            the instance name.
        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/"``).
        group_id (str, optional): Consumer group ID.  When ``None``
            (the default), a unique ephemeral ID is generated from
            the hostname, PID, and current time so that each consumer
            instance operates independently.  Provide a stable string
            to share offsets across restarts or multiple consumers.
        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.consume")
        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``.
        auto_offset_reset (str): Where to start reading when there is
            no committed offset for the consumer group.
            Default: ``"latest"`` (only new messages).  Set to
            ``"earliest"`` to replay from the beginning of the topic.
        **kafka_config: Additional confluent_kafka.Consumer
            configuration options.  For example, pass
            ``**{"enable.auto.commit": False}`` to manage offset
            commits manually.

    Example:
        Stream only new messages (default behaviour)::

            with GraceDbKafkaConsumer(
                bootstrap_servers="kafka-dev.ligo.org:9092",
                topics="test",
                ca_cert_path="/tmp/cacert.pem",
            ) as consumer:
                for msg in consumer.stream():
                    print(msg)

        Resume from last committed offset across restarts::

            with GraceDbKafkaConsumer(
                bootstrap_servers="kafka-dev.ligo.org:9092",
                topics="test",
                group_id="my-pipeline-consumer",
                ca_cert_path="/tmp/cacert.pem",
            ) as consumer:
                for msg in consumer.stream():
                    process(msg)

        Replay all messages from the beginning of a topic::

            with GraceDbKafkaConsumer(
                bootstrap_servers="kafka-dev.ligo.org:9092",
                topics="test",
                group_id="my-replay-group",
                auto_offset_reset="earliest",
                ca_cert_path="/tmp/cacert.pem",
            ) as consumer:
                for msg in consumer.stream():
                    print(msg)
    """

    @staticmethod
    def _generate_group_id():
        """Generate a unique ephemeral consumer group ID.

        The ID is derived from the hostname, PID, and current time so
        that each consumer instance gets its own group.  This ensures
        a fresh consumer with ``auto.offset.reset=latest`` will only
        see messages produced after it connects.

        Returns:
            str: A group ID of the form ``gracedb-client-<8-char hex>``.
        """
        raw = f"{socket.gethostname()}-{os.getpid()}-{time.time()}"
        digest = hashlib.sha256(raw.encode()).hexdigest()[:8]
        return f"gracedb-client-{digest}"

    def __init__(self, bootstrap_servers, topics=None,
                 instance_name=None, service_url=None,
                 group_id=None, use_ssl=True, ca_cert_path=None,
                 skip_hostname_verification=False,
                 token_scope="kafka.consume", token_audience=None,
                 reload_cred=True, reload_buffer=300,
                 auto_refresh_token=None, auto_offset_reset="latest",
                 **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.token_scope = token_scope
        self.token_audience = token_audience

        if group_id is None:
            group_id = self._generate_group_id()
            log.info(f"Using ephemeral consumer group: {group_id}")

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

        # Build consumer configuration.
        # enable.auto.commit defaults to True (confluent-kafka default)
        # but can be overridden via kafka_config.
        config = {
            'bootstrap.servers': bootstrap_servers,
            'client.id': 'ligo-gracedb-consumer',
            'group.id': group_id,
            'auto.offset.reset': auto_offset_reset,
            'enable.auto.commit': True,
        }

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

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

            if skip_hostname_verification:
                config['ssl.endpoint.identification.algorithm'] = 'none'
                log.warning(
                    "SSL hostname verification disabled - "
                    "not recommended for production"
                )
        else:
            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:
            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:
            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 consumer
        try:
            self.consumer = Consumer(config)
            log.info("Kafka consumer created successfully")
        except KafkaException as e:
            log.error(f"Failed to create Kafka consumer: {e}")
            raise KafkaConnectionError(
                f"Failed to connect to Kafka broker: {e}"
            ) from e

        self._running = False

        # Subscribe to topics if provided
        if topics is not None:
            self.subscribe(topics)

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

        Detects authentication failures (e.g. expired SciToken) and
        stores them so the next poll call raises a clean Python
        exception instead of spinning on retries.
        """
        if error.code() == _AUTHENTICATION:
            self._auth_error = error
            self._running = False
            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 subscribe(self, topics):
        """
        Subscribe to one or more Kafka topics.

        Topic names are automatically namespaced with the instance name.
        For example, subscribing to ``"gstlal"`` with instance name
        ``"gracedb-dev"`` subscribes to ``"gracedb-dev.gstlal"``.

        Args:
            topics (str or list): Pipeline name or list of pipeline names

        Raises:
            KafkaConsumerError: If subscription fails
        """
        if isinstance(topics, str):
            topics = [topics]

        topics = [make_topic_name(self.instance_name, t) for t in topics]

        try:
            self.consumer.subscribe(topics)
            log.info(f"Subscribed to topics: {topics}")
        except KafkaException as e:
            log.error(f"Failed to subscribe to topics: {e}")
            raise KafkaConsumerError(
                f"Failed to subscribe to topics: {e}"
            ) from e

    def poll(self, timeout=1.0):
        """
        Poll for a single message.

        Args:
            timeout (float): Maximum time to wait for a message in seconds
                (default: 1.0)

        Returns:
            dict or None: Message dict with keys ``topic``, ``partition``,
            ``offset``, ``key``, ``value``, ``timestamp``, ``headers``,
            or None if no message is available.

        Raises:
            KafkaConsumerError: If a consumer error occurs
        """
        msg = self.consumer.poll(timeout)
        self._check_auth()

        if msg is None:
            return None

        err = msg.error()
        if err:
            raise KafkaConsumerError(f"Consumer error: {err}")

        # Decode value
        raw_value = msg.value()
        if raw_value is not None:
            try:
                value = json.loads(raw_value)
            except (json.JSONDecodeError, UnicodeDecodeError):
                # Fall back to string, then raw bytes
                try:
                    value = raw_value.decode('utf-8')
                except (UnicodeDecodeError, AttributeError):
                    value = raw_value
        else:
            value = None

        # Decode key
        raw_key = msg.key()
        if raw_key is not None:
            try:
                raw_key = raw_key.decode('utf-8')
            except (UnicodeDecodeError, AttributeError):
                pass

        # Extract headers
        headers = msg.headers()
        if headers is not None:
            try:
                headers = {
                    k: v.decode('utf-8') if isinstance(v, bytes) else v
                    for k, v in headers
                }
            except (UnicodeDecodeError, AttributeError):
                headers = {k: v for k, v in headers}

        # Extract timestamp
        ts_type, ts_value = msg.timestamp()
        timestamp = ts_value if ts_type != 0 else None

        return {
            'topic': msg.topic(),
            'partition': msg.partition(),
            'offset': msg.offset(),
            'key': raw_key,
            'value': value,
            'timestamp': timestamp,
            'headers': headers,
        }

    def stop(self):
        """Signal the stream() loop to stop at the next poll iteration."""
        log.info("Consumer stop requested")
        self._running = False

    def _handle_shutdown_signal(self, signum, frame):
        """Signal handler for graceful shutdown during stream()."""
        sig_name = signal.Signals(signum).name
        log.info(f"Received {sig_name}, shutting down consumer...")
        self._running = False

    def stream(self, timeout=1.0, count=None):
        """
        Generator that yields messages from subscribed topics.

        Installs signal handlers for SIGINT and SIGTERM so that the
        first signal causes the loop to exit cleanly (within at most
        *timeout* seconds).  Original handlers are restored when the
        generator exits.

        Args:
            timeout (float): Poll timeout in seconds per iteration
                (default: 1.0)
            count (int, optional): Maximum number of messages to yield.
                None means stream indefinitely.

        Yields:
            dict: Message dicts (same format as :meth:`poll`)
        """
        self._running = True
        yielded = 0

        # Install signal handlers for clean shutdown (main thread only)
        prev_sigint = prev_sigterm = None
        try:
            prev_sigint = signal.signal(
                signal.SIGINT, self._handle_shutdown_signal)
            prev_sigterm = signal.signal(
                signal.SIGTERM, self._handle_shutdown_signal)
        except ValueError:
            # Not in main thread; signal handlers can't be installed
            pass

        try:
            while self._running and (count is None or yielded < count):
                msg = self.poll(timeout=timeout)
                if msg is not None:
                    yield msg
                    yielded += 1
        finally:
            self._running = False
            if prev_sigint is not None:
                signal.signal(signal.SIGINT, prev_sigint)
            if prev_sigterm is not None:
                signal.signal(signal.SIGTERM, prev_sigterm)

    def close(self):
        """
        Unsubscribe from topics and close the consumer.
        """
        log.info("Closing Kafka consumer...")
        try:
            self.consumer.close()
            log.info("Kafka consumer closed")
        except Exception as e:
            log.error(f"Error closing consumer: {e}")
            raise

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

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