/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.connect.util;

import org.apache.kafka.connect.util.LoggingContext.Scope;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class LoggingContextTest {

    private static final Logger log = LoggerFactory.getLogger(LoggingContextTest.class);

    private static final String CONNECTOR_NAME = "MyConnector";
    private static final ConnectorTaskId TASK_ID1 = new ConnectorTaskId(CONNECTOR_NAME, 1);
    private static final String EXTRA_KEY1 = "extra.key.1";
    private static final String EXTRA_VALUE1 = "value1";
    private static final String EXTRA_KEY2 = "extra.key.2";
    private static final String EXTRA_VALUE2 = "value2";
    private static final String EXTRA_KEY3 = "extra.key.3";
    private static final String EXTRA_VALUE3 = "value3";

    private Map<String, String> mdc;

    @BeforeEach
    public void setup() {
        mdc = new HashMap<>();
        Map<String, String> existing = MDC.getCopyOfContextMap();
        if (existing != null) {
            mdc.putAll(existing);
        }
        MDC.put(EXTRA_KEY1, EXTRA_VALUE1);
        MDC.put(EXTRA_KEY2, EXTRA_VALUE2);
    }

    @AfterEach
    public void tearDown() {
        MDC.clear();
        MDC.setContextMap(mdc);
    }

    @Test
    public void shouldNotAllowNullConnectorNameForConnectorContext() {
        assertThrows(NullPointerException.class, () -> LoggingContext.forConnector(null));
    }

    @Test
    public void shouldNotAllowNullTaskIdForTaskContext() {
        assertThrows(NullPointerException.class, () -> LoggingContext.forTask(null));
    }

    @Test
    public void shouldNotAllowNullTaskIdForOffsetContext() {
        assertThrows(NullPointerException.class, () -> LoggingContext.forOffsets(null));
    }

    @Test
    public void shouldCreateAndCloseLoggingContextEvenWithNullContextMap() {
        MDC.clear();
        assertMdc(null, null, null);
        try (LoggingContext loggingContext = LoggingContext.forConnector(CONNECTOR_NAME)) {
            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
            log.info("Starting Connector");
        }
        assertMdc(null, null, null);
    }

    @Test
    public void shouldCreateConnectorLoggingContext() {
        assertMdcExtrasUntouched();
        assertMdc(null, null, null);

        try (LoggingContext loggingContext = LoggingContext.forConnector(CONNECTOR_NAME)) {
            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
            log.info("Starting Connector");
        }

        assertMdcExtrasUntouched();
        assertMdc(null, null, null);
    }

    @Test
    public void shouldCreateTaskLoggingContext() {
        assertMdcExtrasUntouched();
        try (LoggingContext loggingContext = LoggingContext.forTask(TASK_ID1)) {
            assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.TASK);
            log.info("Running task");
        }

        assertMdcExtrasUntouched();
        assertMdc(null, null, null);
    }

    @Test
    public void shouldCreateOffsetsLoggingContext() {
        assertMdcExtrasUntouched();
        try (LoggingContext loggingContext = LoggingContext.forOffsets(TASK_ID1)) {
            assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.OFFSETS);
            log.info("Running task");
        }

        assertMdcExtrasUntouched();
        assertMdc(null, null, null);
    }

    @Test
    public void shouldAllowNestedLoggingContexts() {
        assertMdcExtrasUntouched();
        assertMdc(null, null, null);
        try (LoggingContext loggingContext1 = LoggingContext.forConnector(CONNECTOR_NAME)) {
            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
            log.info("Starting Connector");
            // Set the extra MDC parameter, as if the connector were
            MDC.put(EXTRA_KEY3, EXTRA_VALUE3);
            assertConnectorMdcSet();

            try (LoggingContext loggingContext2 = LoggingContext.forTask(TASK_ID1)) {
                assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.TASK);
                log.info("Starting task");
                // The extra connector-specific MDC parameter should still be set
                assertConnectorMdcSet();

                try (LoggingContext loggingContext3 = LoggingContext.forOffsets(TASK_ID1)) {
                    assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.OFFSETS);
                    assertConnectorMdcSet();
                    log.info("Offsets for task");
                }

                assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.TASK);
                log.info("Stopping task");
                // The extra connector-specific MDC parameter should still be set
                assertConnectorMdcSet();
            }

            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
            log.info("Stopping Connector");
            // The extra connector-specific MDC parameter should still be set
            assertConnectorMdcSet();
        }
        assertMdcExtrasUntouched();
        assertMdc(null, null, null);

        // The extra connector-specific MDC parameter should still be set
        assertConnectorMdcSet();

        LoggingContext.clear();
        assertConnectorMdcUnset();
    }

    protected void assertMdc(String connectorName, Integer taskId, Scope scope) {
        String context = MDC.get(LoggingContext.CONNECTOR_CONTEXT);
        if (context != null) {
            assertEquals(
                connectorName != null,
                context.startsWith("[" + connectorName),
                "Context should begin with connector name when the connector name is non-null"
            );
            if (scope != null) {
                assertTrue(context.contains(scope.toString()), "Context should contain the scope");
            }
            if (taskId != null) {
                assertTrue(context.contains(taskId.toString()), "Context should contain the taskId");
            }
        } else {
            assertNull(connectorName, "No logging context found, expected null connector name");
            assertNull(taskId, "No logging context found, expected null task ID");
            assertNull(scope, "No logging context found, expected null scope");
        }
    }

    protected void assertMdcExtrasUntouched() {
        assertEquals(EXTRA_VALUE1, MDC.get(EXTRA_KEY1));
        assertEquals(EXTRA_VALUE2, MDC.get(EXTRA_KEY2));
    }

    protected void assertConnectorMdcSet() {
        assertEquals(EXTRA_VALUE3, MDC.get(EXTRA_KEY3));
    }

    protected void assertConnectorMdcUnset() {
        assertNull(MDC.get(EXTRA_KEY3));
    }
}