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

import base64
import json
import time
from unittest import mock
import pytest

from igwn_auth_utils.error import IgwnAuthError

from ligo.gracedb.kafka.auth import (
    _decode_jwt_expiry,
    get_scitoken_for_kafka,
    create_oauth_callback,
    SciTokenRefreshManager,
)
from ligo.gracedb.kafka.exceptions import KafkaTokenError


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 TestDecodeJwtExpiry:
    """Test _decode_jwt_expiry helper."""

    def test_extracts_exp_claim(self):
        """exp claim is returned as int when present."""
        token = _make_jwt({"exp": 1700000000, "sub": "alice"})
        assert _decode_jwt_expiry(token) == 1700000000

    def test_fallback_when_no_exp(self):
        """Falls back to ~now+3600 when exp claim is missing."""
        token = _make_jwt({"sub": "alice"})
        before = int(time.time() + 3600)
        result = _decode_jwt_expiry(token)
        after = int(time.time() + 3600)
        assert before <= result <= after

    def test_fallback_on_garbage_token(self):
        """Falls back to ~now+3600 on unparseable tokens."""
        before = int(time.time() + 3600)
        result = _decode_jwt_expiry("not-a-jwt")
        after = int(time.time() + 3600)
        assert before <= result <= after

    def test_exp_returned_as_int(self):
        """exp is coerced to int even if stored as float in the JWT."""
        token = _make_jwt({"exp": 1700000000.5})
        assert _decode_jwt_expiry(token) == 1700000000
        assert isinstance(_decode_jwt_expiry(token), int)


class TestGetScitokenForKafka:
    """Test get_scitoken_for_kafka function."""

    @mock.patch('ligo.gracedb.kafka.auth.find_token')
    def test_successful_token_discovery(self, mock_find_token):
        """Test successful token discovery."""
        # Mock token object
        mock_token = mock.Mock()
        mock_token.serialize.return_value = "mock.jwt.token"
        mock_token.get.side_effect = lambda key: {
            'sub': 'test-user',
            'exp': int(time.time()) + 3600
        }.get(key)
        mock_find_token.return_value = mock_token

        # Call function
        token_str = get_scitoken_for_kafka("localhost:9092", scope="kafka.produce")

        # Verify
        assert token_str == "mock.jwt.token"
        mock_find_token.assert_called_once_with("localhost", "kafka.produce", skip_errors=True)
        mock_token.serialize.assert_called_once()

    @mock.patch('ligo.gracedb.kafka.auth.find_token')
    def test_custom_audience(self, mock_find_token):
        """Test token discovery with custom audience."""
        mock_token = mock.Mock()
        mock_token.serialize.return_value = "mock.jwt.token"
        mock_find_token.return_value = mock_token

        get_scitoken_for_kafka(
            "localhost:9092",
            scope="kafka.produce",
            audience="custom-audience"
        )

        mock_find_token.assert_called_once_with("custom-audience", "kafka.produce", skip_errors=True)

    @mock.patch('ligo.gracedb.kafka.auth._get_bearer_token_from_file', return_value=None)
    @mock.patch('ligo.gracedb.kafka.auth.find_token')
    def test_no_token_found(self, mock_find_token, mock_get_file):
        """Test when no token is found."""
        mock_find_token.return_value = None

        with pytest.raises(KafkaTokenError) as exc_info:
            get_scitoken_for_kafka("localhost:9092")

        assert "No valid SciToken found" in str(exc_info.value)
        assert "localhost" in str(exc_info.value)

    @mock.patch('ligo.gracedb.kafka.auth._get_bearer_token_from_file', return_value=None)
    @mock.patch('ligo.gracedb.kafka.auth.find_token')
    def test_igwn_auth_error(self, mock_find_token, mock_get_file):
        """Test handling of IgwnAuthError."""
        mock_find_token.side_effect = IgwnAuthError("Auth error")

        with pytest.raises(KafkaTokenError) as exc_info:
            get_scitoken_for_kafka("localhost:9092")

        assert "No valid SciToken found" in str(exc_info.value)

    @mock.patch('ligo.gracedb.kafka.auth._get_bearer_token_from_file',
                return_value=None)
    @mock.patch('ligo.gracedb.kafka.auth.find_token')
    def test_unexpected_error(self, mock_find_token, mock_get_file):
        """Test handling of unexpected errors falls through to file read."""
        mock_find_token.side_effect = RuntimeError("Unexpected")

        with pytest.raises(KafkaTokenError) as exc_info:
            get_scitoken_for_kafka("localhost:9092")

        assert "No valid SciToken found" in str(exc_info.value)

    @mock.patch('ligo.gracedb.kafka.auth.find_token')
    def test_broker_url_with_port(self, mock_find_token):
        """Test extracting audience from broker URL with port."""
        mock_token = mock.Mock()
        mock_token.serialize.return_value = "token"
        mock_find_token.return_value = mock_token

        get_scitoken_for_kafka("broker.example.com:9092")

        mock_find_token.assert_called_once_with("broker.example.com", "kafka.produce", skip_errors=True)

    @mock.patch('ligo.gracedb.kafka.auth.find_token')
    def test_broker_url_without_port(self, mock_find_token):
        """Test extracting audience from broker URL without port."""
        mock_token = mock.Mock()
        mock_token.serialize.return_value = "token"
        mock_find_token.return_value = mock_token

        get_scitoken_for_kafka("broker.example.com")

        mock_find_token.assert_called_once_with("broker.example.com", "kafka.produce", skip_errors=True)


class TestCreateOauthCallback:
    """Test create_oauth_callback function."""

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_callback_returns_token_and_expiry(self, mock_get_token):
        """Test that callback returns token and real exp from JWT."""
        exp_time = int(time.time()) + 7200
        mock_get_token.return_value = _make_jwt({"exp": exp_time})

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

        token, expiry = callback(None)

        assert token == mock_get_token.return_value
        assert isinstance(expiry, int)
        assert expiry == exp_time

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_callback_with_custom_audience(self, mock_get_token):
        """Test callback with custom audience."""
        mock_get_token.return_value = "mock.jwt.token"

        callback = create_oauth_callback(
            "localhost:9092",
            scope="kafka.consume",
            audience="custom-audience"
        )
        callback(None)

        mock_get_token.assert_called_once_with(
            "localhost:9092",
            "kafka.consume",
            "custom-audience"
        )

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_callback_propagates_token_error(self, mock_get_token):
        """Test that callback propagates KafkaTokenError."""
        mock_get_token.side_effect = KafkaTokenError("No token")

        callback = create_oauth_callback("localhost:9092")

        with pytest.raises(KafkaTokenError):
            callback(None)

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_callback_wraps_unexpected_error(self, mock_get_token):
        """Test that callback wraps unexpected errors."""
        mock_get_token.side_effect = RuntimeError("Unexpected")

        callback = create_oauth_callback("localhost:9092")

        with pytest.raises(KafkaTokenError) as exc_info:
            callback(None)

        assert "OAuth callback failed" in str(exc_info.value)


class TestSciTokenRefreshManager:
    """Test SciTokenRefreshManager class."""

    def test_initialization(self):
        """Test manager initialization."""
        manager = SciTokenRefreshManager(
            "localhost:9092",
            scope="kafka.produce",
            audience="custom-audience",
            refresh_buffer=600
        )

        assert manager.broker_url == "localhost:9092"
        assert manager.scope == "kafka.produce"
        assert manager.audience == "custom-audience"
        assert manager.refresh_buffer == 600
        assert manager.current_token is None
        assert manager.token_expiry == 0

    def test_default_audience_extraction(self):
        """Test default audience extraction from broker URL."""
        manager = SciTokenRefreshManager("broker.example.com:9092")

        assert manager.audience == "broker.example.com"

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_get_token_refreshes_on_first_call(self, mock_get_token):
        """Test that get_token refreshes on first call."""
        exp = int(time.time()) + 7200
        fresh_token = _make_jwt({"exp": exp})
        mock_get_token.return_value = fresh_token

        manager = SciTokenRefreshManager("localhost:9092")
        token = manager.get_token()

        assert token == fresh_token
        assert manager.current_token == fresh_token
        assert manager.token_expiry == exp
        mock_get_token.assert_called_once()

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_get_token_returns_cached_if_valid(self, mock_get_token):
        """Test that get_token returns cached token if still valid."""
        exp = int(time.time()) + 7200
        mock_get_token.return_value = _make_jwt({"exp": exp})

        manager = SciTokenRefreshManager("localhost:9092", refresh_buffer=300)

        # First call - should refresh
        token1 = manager.get_token()

        # Second call immediately - should use cache
        token2 = manager.get_token()

        assert token1 == token2
        assert mock_get_token.call_count == 1  # Only called once

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_get_token_refreshes_when_near_expiry(self, mock_get_token):
        """Test that get_token refreshes when token is about to expire."""
        # First token expires 100s from now (inside 300s buffer)
        exp1 = int(time.time()) + 100
        # Second token expires far in the future
        exp2 = int(time.time()) + 7200
        mock_get_token.side_effect = [
            _make_jwt({"exp": exp1}),
            _make_jwt({"exp": exp2}),
        ]

        manager = SciTokenRefreshManager("localhost:9092", refresh_buffer=300)

        # First call — gets token1
        token1 = manager.get_token()
        # token_expiry - time.time() <= 300, so second call triggers refresh
        token2 = manager.get_token()

        assert token1 != token2
        assert mock_get_token.call_count == 2
        assert manager.token_expiry == exp2

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_create_callback_uses_manager(self, mock_get_token):
        """Test that create_callback creates a callback using the manager."""
        exp = int(time.time()) + 7200
        jwt = _make_jwt({"exp": exp})
        mock_get_token.return_value = jwt

        manager = SciTokenRefreshManager("localhost:9092")
        callback = manager.create_callback()

        token, expiry = callback(None)

        assert token == jwt
        assert manager.current_token == jwt
        assert expiry == exp

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_create_callback_propagates_errors(self, mock_get_token):
        """Test that callback from manager propagates errors."""
        mock_get_token.side_effect = KafkaTokenError("Token error")

        manager = SciTokenRefreshManager("localhost:9092")
        callback = manager.create_callback()

        with pytest.raises(KafkaTokenError):
            callback(None)

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_refresh_uses_real_exp(self, mock_get_token):
        """Test that _refresh_token parses the real exp claim."""
        exp = int(time.time()) + 9000
        mock_get_token.return_value = _make_jwt({"exp": exp})

        manager = SciTokenRefreshManager("localhost:9092")
        manager._refresh_token()

        assert manager.current_token == mock_get_token.return_value
        assert manager.token_expiry == exp

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_refresh_fallback_when_no_exp(self, mock_get_token):
        """Test that _refresh_token falls back when JWT has no exp."""
        mock_get_token.return_value = _make_jwt({"sub": "alice"})

        manager = SciTokenRefreshManager("localhost:9092")
        before = int(time.time() + 3600)
        manager._refresh_token()
        after = int(time.time() + 3600)

        assert before <= manager.token_expiry <= after

    @mock.patch('ligo.gracedb.kafka.auth.get_scitoken_for_kafka')
    def test_refresh_propagates_exceptions(self, mock_get_token):
        """Test that _refresh_token propagates exceptions."""
        mock_get_token.side_effect = KafkaTokenError("Refresh failed")

        manager = SciTokenRefreshManager("localhost:9092")

        with pytest.raises(KafkaTokenError):
            manager._refresh_token()
