# -*- 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/>.
"""
SciToken authentication for Kafka producers.

This module provides functions for discovering and using SciTokens
with Kafka's OAUTHBEARER SASL mechanism.
"""

import base64
import json
import logging
import time
import os
from igwn_auth_utils.scitokens import find_token
from igwn_auth_utils.error import IgwnAuthError

from .exceptions import KafkaTokenError

log = logging.getLogger(__name__)


def _decode_jwt_expiry(token_str):
    """Extract the ``exp`` claim from a JWT without signature verification.

    Args:
        token_str (str): Serialized JWT string (header.payload.signature).

    Returns:
        int: The ``exp`` timestamp from the token payload, or
        ``int(time.time() + 3600)`` if no ``exp`` claim is present.
    """
    try:
        payload = token_str.split('.')[1]
        # Restore base64url padding
        payload += '=' * (4 - len(payload) % 4)
        claims = json.loads(base64.urlsafe_b64decode(payload))
        return int(claims.get('exp', int(time.time() + 3600)))
    except Exception:
        return int(time.time() + 3600)


def _get_bearer_token_from_file():
    """
    Get bearer token using WLCG Bearer Token Discovery.

    Follows the discovery order:
    1. $BEARER_TOKEN environment variable (token content)
    2. $BEARER_TOKEN_FILE environment variable (path to file)
    3. ${XDG_RUNTIME_DIR}/bt_u<uid> or /tmp/bt_u<uid>

    Returns:
        str: Token string or None if not found
    """
    # Try $BEARER_TOKEN
    if 'BEARER_TOKEN' in os.environ:
        log.debug("Using token from $BEARER_TOKEN")
        return os.environ['BEARER_TOKEN']

    # Try $BEARER_TOKEN_FILE
    if 'BEARER_TOKEN_FILE' in os.environ:
        token_file = os.environ['BEARER_TOKEN_FILE']
        if os.path.isfile(token_file):
            log.debug(f"Reading token from $BEARER_TOKEN_FILE: {token_file}")
            with open(token_file, 'r') as f:
                return f.read().strip()

    # Try default location
    uid = os.getuid()
    runtime_dir = os.environ.get('XDG_RUNTIME_DIR', '/tmp')
    token_file = os.path.join(runtime_dir, f'bt_u{uid}')

    if os.path.isfile(token_file):
        log.debug(f"Reading token from {token_file}")
        with open(token_file, 'r') as f:
            return f.read().strip()

    return None


def get_scitoken_for_kafka(broker_url, scope="kafka.produce", audience=None):
    """
    Find and return a valid SciToken for Kafka authentication.

    This function discovers a SciToken using igwn-auth-utils and returns
    the serialized JWT string suitable for Kafka OAUTHBEARER authentication.

    Args:
        broker_url (str): Kafka broker URL (e.g., "localhost:9092")
        scope (str): Required token scope (default: "kafka.produce")
        audience (str, optional): Required audience. If None, extracted from broker_url

    Returns:
        str: Serialized JWT token string

    Raises:
        KafkaTokenError: If no valid token can be found
    """
    # Extract hostname from broker URL for audience
    if audience is None:
        # broker_url format: "hostname:port" or "hostname"
        audience = broker_url.split(':')[0]

    log.debug(f"Searching for SciToken with audience='{audience}', scope='{scope}'")

    try:
        # Try multiple strategies to find a token
        token = None

        # Strategy 1: Try exact match with specified audience and scope
        if scope and audience:
            try:
                token = find_token(audience, scope, skip_errors=True)
            except Exception:
                pass

        # Strategy 2: Try "ANY" audience (common for development/testing)
        if token is None:
            try:
                # Try with ANY audience and any scope
                token = find_token("ANY", "", skip_errors=True)
                if token:
                    log.info("Found token with audience 'ANY'")
            except Exception:
                pass

        # Strategy 3: Try specified audience with any scope
        if token is None and audience:
            try:
                token = find_token(audience, "", skip_errors=True)
            except Exception:
                pass

        # If igwn-auth-utils returned a token object, try to serialize it.
        # Serialization can fail when the token was deserialized without the
        # private key (e.g. read from a file by igwn-auth-utils internals).
        if token is not None:
            try:
                token_str = token.serialize()
                log.info(f"Found valid SciToken for audience '{audience}'")
                log.debug(f"Token subject: {token.get('sub')}, expires: {token.get('exp')}")
                return token_str
            except Exception:
                log.debug("Token object found but serialize() failed; "
                          "falling back to direct file read")

        # Strategy 4: Fallback to reading token directly from file.
        # This is useful when the token has non-standard audience/scope
        # or when igwn-auth-utils cannot serialize the token object.
        log.info("Trying direct token file read as fallback")
        token_str = _get_bearer_token_from_file()
        if token_str:
            log.info("Found token via direct file read")
            return token_str

        raise KafkaTokenError(
            f"No valid SciToken found. Please obtain a token using htgettoken. "
            f"Tried audience: '{audience}' or 'ANY', and direct file read."
        )

    except IgwnAuthError as e:
        raise KafkaTokenError(f"Error finding SciToken: {e}") from e
    except Exception as e:
        raise KafkaTokenError(f"Unexpected error obtaining SciToken: {e}") from e


def create_oauth_callback(broker_url, scope="kafka.produce", audience=None):
    """
    Create an OAuth callback function for confluent_kafka Producer.

    The callback is called by the Kafka client to obtain the OAuth token
    before establishing a connection.

    Args:
        broker_url (str): Kafka broker URL
        scope (str): Required token scope
        audience (str, optional): Required audience

    Returns:
        callable: OAuth callback function that returns (token_str, expiry_time)
    """
    def oauth_cb(oauth_config):
        """
        OAuth callback for confluent-kafka Producer.

        Args:
            oauth_config (str): OAuth configuration string (not used)

        Returns:
            tuple: (token_string, expiration_timestamp)

        Raises:
            KafkaTokenError: If token cannot be obtained
        """
        log.debug("OAuth callback invoked by Kafka client")

        try:
            # Get the SciToken
            token_str = get_scitoken_for_kafka(broker_url, scope, audience)

            # Parse the real exp claim from the JWT so that
            # confluent-kafka knows when to request a new token.
            expiry_time = _decode_jwt_expiry(token_str)

            log.debug("OAuth callback returning token")
            return token_str, expiry_time

        except KafkaTokenError:
            # Re-raise KafkaTokenError as-is
            raise
        except Exception as e:
            # Wrap unexpected errors
            raise KafkaTokenError(f"OAuth callback failed: {e}") from e

    return oauth_cb


class SciTokenRefreshManager:
    """
    Manages automatic SciToken refresh for long-running Kafka producers.

    This class monitors token expiration and triggers callbacks when
    tokens need to be refreshed.
    """

    def __init__(self, broker_url, scope="kafka.produce", audience=None,
                 refresh_buffer=300):
        """
        Initialize the token refresh manager.

        Args:
            broker_url (str): Kafka broker URL
            scope (str): Required token scope
            audience (str, optional): Required audience
            refresh_buffer (int): Seconds before expiry to trigger refresh (default: 300)
        """
        self.broker_url = broker_url
        self.scope = scope
        self.audience = audience or broker_url.split(':')[0]
        self.refresh_buffer = refresh_buffer
        self.current_token = None
        self.token_expiry = 0
        log.debug(f"SciTokenRefreshManager initialized for {self.audience}")

    def get_token(self):
        """
        Get current token, refreshing if necessary.

        Returns:
            str: Current valid token string

        Raises:
            KafkaTokenError: If token cannot be obtained
        """
        current_time = time.time()

        # Check if we need to refresh
        if self.current_token is None or \
           (self.token_expiry - current_time) <= self.refresh_buffer:
            log.info("Refreshing SciToken")
            self._refresh_token()

        return self.current_token

    def _refresh_token(self):
        """Refresh the current token."""
        try:
            token_str = get_scitoken_for_kafka(
                self.broker_url,
                self.scope,
                self.audience
            )

            self.current_token = token_str
            self.token_expiry = _decode_jwt_expiry(token_str)

            log.info(f"Token refreshed, expires at {self.token_expiry}")

        except Exception as e:
            log.error(f"Failed to refresh token: {e}")
            raise

    def create_callback(self):
        """
        Create an OAuth callback that uses this manager.

        Returns:
            callable: OAuth callback function
        """
        def oauth_cb(oauth_config):
            try:
                token_str = self.get_token()
                return token_str, self.token_expiry
            except Exception as e:
                raise KafkaTokenError(f"OAuth callback failed: {e}") from e

        return oauth_cb
