package org.oran.pmlog.oauth2;

import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.auth.login.AppConfigurationEntry;
import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;

class OAuthKafkaAuthenticateLoginCallbackHandlerTest {

    private OAuthKafkaAuthenticateLoginCallbackHandler callbackHandler;

    @BeforeEach
    void setUp() {
        callbackHandler = new OAuthKafkaAuthenticateLoginCallbackHandler();
    }

    @Test
    void testConfigureWithValidSaslMechanismAndConfigEntry() {
        String saslMechanism = OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
        List<AppConfigurationEntry> jaasConfigEntries = Collections.singletonList(Mockito.mock(AppConfigurationEntry.class));

        callbackHandler.configure(new HashMap<>(), saslMechanism, jaasConfigEntries);

        assertTrue(callbackHandler.isConfigured());
    }

    @SuppressWarnings("java:S5778")
    @Test
    void testConfigureWithInvalidSaslMechanism() {
        String invalidSaslMechanism = "InvalidMechanism";
        List<AppConfigurationEntry> jaasConfigEntries = Collections.singletonList(Mockito.mock(AppConfigurationEntry.class));

        assertThrows(IllegalArgumentException.class, () -> callbackHandler.configure(new HashMap<>(), invalidSaslMechanism, jaasConfigEntries));

        assertFalse(callbackHandler.isConfigured());
    }

    @SuppressWarnings("java:S5778")
    @Test
    void testConfigureWithEmptyJaasConfigEntries() {
        String saslMechanism = OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
        List<AppConfigurationEntry> emptyJaasConfigEntries = Collections.emptyList();

        assertThrows(IllegalArgumentException.class, () -> callbackHandler.configure(new HashMap<>(), saslMechanism, emptyJaasConfigEntries));

        assertFalse(callbackHandler.isConfigured());
    }

    @Test
    void testHandleSaslExtensionsCallback() throws IOException, UnsupportedCallbackException {
        String saslMechanism = OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
        List<AppConfigurationEntry> jaasConfigEntries = Collections.singletonList(Mockito.mock(AppConfigurationEntry.class));

        callbackHandler.configure(new HashMap<>(), saslMechanism, jaasConfigEntries);
        SaslExtensionsCallback callback = mock(SaslExtensionsCallback.class);

        callbackHandler.handle(new Callback[]{callback});
        verify(callback).extensions(any());
    }

    @Test
    void testHandleUnsupportedCallback() {
        Callback unsupportedCallback = mock(Callback.class);
        String saslMechanism = OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
        List<AppConfigurationEntry> jaasConfigEntries = Collections.singletonList(Mockito.mock(AppConfigurationEntry.class));

        callbackHandler.configure(new HashMap<>(), saslMechanism, jaasConfigEntries);
        assertThrows(UnsupportedCallbackException.class, () -> callbackHandler.handle(new Callback[]{unsupportedCallback}));
    }

    @Test
    void testHandleOAuthBearerTokenCallback() throws IOException, UnsupportedCallbackException {

        String saslMechanism = OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
        List<AppConfigurationEntry> jaasConfigEntries = Collections.singletonList(Mockito.mock(AppConfigurationEntry.class));
        String validJwt = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c";

        callbackHandler.configure(new HashMap<>(), saslMechanism, jaasConfigEntries);

        OAuthBearerTokenCallback oauthBearerTokenCallback = Mockito.mock(OAuthBearerTokenCallback.class);
        SecurityContext securityContextMock = Mockito.mock(SecurityContext.class);
        when(oauthBearerTokenCallback.token()).thenReturn(null); // Ensure the callback has no token initially
        when(oauthBearerTokenCallback.token()).thenAnswer(invocation -> {
            return OAuthBearerTokenJwt.create(validJwt);
        });

        when(securityContextMock.getBearerAuthToken()).thenReturn(validJwt);
        callbackHandler.handle(new Callback[]{oauthBearerTokenCallback});
        verify(oauthBearerTokenCallback).token();
    }
}