/*
 * 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.image;

import org.apache.kafka.common.Endpoint;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.metadata.BrokerRegistrationChangeRecord;
import org.apache.kafka.common.metadata.FenceBrokerRecord;
import org.apache.kafka.common.metadata.RegisterBrokerRecord;
import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpoint;
import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpointCollection;
import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerFeature;
import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerFeatureCollection;
import org.apache.kafka.common.metadata.RegisterControllerRecord;
import org.apache.kafka.common.metadata.RegisterControllerRecord.ControllerEndpoint;
import org.apache.kafka.common.metadata.RegisterControllerRecord.ControllerEndpointCollection;
import org.apache.kafka.common.metadata.UnfenceBrokerRecord;
import org.apache.kafka.common.metadata.UnregisterBrokerRecord;
import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.image.writer.ImageWriterOptions;
import org.apache.kafka.image.writer.RecordListWriter;
import org.apache.kafka.metadata.BrokerRegistration;
import org.apache.kafka.metadata.BrokerRegistrationInControlledShutdownChange;
import org.apache.kafka.metadata.ControllerRegistration;
import org.apache.kafka.metadata.RecordTestUtils;
import org.apache.kafka.metadata.VersionRange;
import org.apache.kafka.server.common.ApiMessageAndVersion;
import org.apache.kafka.server.common.MetadataVersion;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;

import static org.apache.kafka.common.metadata.MetadataRecordType.BROKER_REGISTRATION_CHANGE_RECORD;
import static org.apache.kafka.common.metadata.MetadataRecordType.FENCE_BROKER_RECORD;
import static org.apache.kafka.common.metadata.MetadataRecordType.REGISTER_BROKER_RECORD;
import static org.apache.kafka.common.metadata.MetadataRecordType.UNFENCE_BROKER_RECORD;
import static org.junit.jupiter.api.Assertions.assertEquals;


@Timeout(value = 40)
public class ClusterImageTest {

    public static final ClusterImage IMAGE1;

    static final List<ApiMessageAndVersion> DELTA1_RECORDS;

    static final ClusterDelta DELTA1;

    static final ClusterImage IMAGE2;

    static final List<ApiMessageAndVersion> DELTA2_RECORDS;

    static final ClusterDelta DELTA2;

    static final ClusterImage IMAGE3;

    static {
        Map<Integer, BrokerRegistration> map1 = new HashMap<>();
        map1.put(0, new BrokerRegistration.Builder().
            setId(0).
            setEpoch(1000).
            setIncarnationId(Uuid.fromString("vZKYST0pSA2HO5x_6hoO2Q")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9092))).
            setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 3))).
            setRack(Optional.empty()).
            setFenced(true).
            setInControlledShutdown(false).build());
        map1.put(1, new BrokerRegistration.Builder().
            setId(1).
            setEpoch(1001).
            setIncarnationId(Uuid.fromString("U52uRe20RsGI0RvpcTx33Q")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9093))).
            setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 3))).
            setRack(Optional.empty()).
            setFenced(false).
            setInControlledShutdown(false).build());
        map1.put(2, new BrokerRegistration.Builder().
            setId(2).
            setEpoch(123).
            setIncarnationId(Uuid.fromString("hr4TVh3YQiu3p16Awkka6w")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9094))).
            setSupportedFeatures(Map.of()).
            setRack(Optional.of("arack")).
            setFenced(false).
            setInControlledShutdown(false).build());
        Map<Integer, ControllerRegistration> cmap1 = new HashMap<>();
        cmap1.put(1000, new ControllerRegistration.Builder().
            setId(1000).
            setIncarnationId(Uuid.fromString("9ABu6HEgRuS-hjHLgC4cHw")).
            setZkMigrationReady(false).
            setListeners(Map.of("PLAINTEXT",
                    new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 19092))).
            setSupportedFeatures(Map.of()).build());
        IMAGE1 = new ClusterImage(map1, cmap1);

        DELTA1_RECORDS = new ArrayList<>();
        // unfence b0
        DELTA1_RECORDS.add(new ApiMessageAndVersion(new UnfenceBrokerRecord().
            setId(0).setEpoch(1000), UNFENCE_BROKER_RECORD.highestSupportedVersion()));
        // fence b1
        DELTA1_RECORDS.add(new ApiMessageAndVersion(new FenceBrokerRecord().
            setId(1).setEpoch(1001), FENCE_BROKER_RECORD.highestSupportedVersion()));
        // mark b0 in controlled shutdown
        DELTA1_RECORDS.add(new ApiMessageAndVersion(new BrokerRegistrationChangeRecord().
            setBrokerId(0).setBrokerEpoch(1000).setInControlledShutdown(
                BrokerRegistrationInControlledShutdownChange.IN_CONTROLLED_SHUTDOWN.value()),
            BROKER_REGISTRATION_CHANGE_RECORD.highestSupportedVersion()));
        // unregister b2
        DELTA1_RECORDS.add(new ApiMessageAndVersion(new UnregisterBrokerRecord().
            setBrokerId(2).setBrokerEpoch(123),
            (short) 0));

        ControllerEndpointCollection endpointsFor1001 = new ControllerEndpointCollection();
        endpointsFor1001.add(new ControllerEndpoint().
            setHost("localhost").
            setName("PLAINTEXT").
            setPort(19093).
            setSecurityProtocol(SecurityProtocol.PLAINTEXT.id));
        DELTA1_RECORDS.add(new ApiMessageAndVersion(new RegisterControllerRecord().
            setControllerId(1001).
            setIncarnationId(Uuid.fromString("FdEHF-IqScKfYyjZ1CjfNQ")).
            setZkMigrationReady(true).
            setEndPoints(endpointsFor1001),
            (short) 0));

        DELTA1 = new ClusterDelta(IMAGE1);
        RecordTestUtils.replayAll(DELTA1, DELTA1_RECORDS);

        Map<Integer, BrokerRegistration> map2 = new HashMap<>();
        map2.put(0, new BrokerRegistration.Builder().
            setId(0).
            setEpoch(1000).
            setIncarnationId(Uuid.fromString("vZKYST0pSA2HO5x_6hoO2Q")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9092))).
            setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 3))).
            setRack(Optional.empty()).
            setFenced(false).
            setInControlledShutdown(true).build());
        map2.put(1, new BrokerRegistration.Builder().
            setId(1).
            setEpoch(1001).
            setIncarnationId(Uuid.fromString("U52uRe20RsGI0RvpcTx33Q")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9093))).
            setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 3))).
            setRack(Optional.empty()).
            setFenced(true).
            setInControlledShutdown(false).build());
        Map<Integer, ControllerRegistration> cmap2 = new HashMap<>(cmap1);
        cmap2.put(1001, new ControllerRegistration.Builder().
            setId(1001).
            setIncarnationId(Uuid.fromString("FdEHF-IqScKfYyjZ1CjfNQ")).
            setZkMigrationReady(true).
            setListeners(Map.of("PLAINTEXT",
                new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 19093))).
            setSupportedFeatures(Map.of()).build());
        IMAGE2 = new ClusterImage(map2, cmap2);

        DELTA2_RECORDS = new ArrayList<>(DELTA1_RECORDS);
        // fence b0
        DELTA2_RECORDS.add(new ApiMessageAndVersion(new FenceBrokerRecord().
            setId(0).setEpoch(1000), FENCE_BROKER_RECORD.highestSupportedVersion()));
        // unfence b1
        DELTA2_RECORDS.add(new ApiMessageAndVersion(new UnfenceBrokerRecord().
            setId(1).setEpoch(1001), UNFENCE_BROKER_RECORD.highestSupportedVersion()));
        // mark b0 as not in controlled shutdown
        DELTA2_RECORDS.add(new ApiMessageAndVersion(new BrokerRegistrationChangeRecord().
            setBrokerId(0).setBrokerEpoch(1000).setInControlledShutdown(
                BrokerRegistrationInControlledShutdownChange.NONE.value()),
            BROKER_REGISTRATION_CHANGE_RECORD.highestSupportedVersion()));
        // re-register b2
        DELTA2_RECORDS.add(new ApiMessageAndVersion(new RegisterBrokerRecord().
            setBrokerId(2).setIsMigratingZkBroker(true).setIncarnationId(Uuid.fromString("Am5Yse7GQxaw0b2alM74bP")).
            setBrokerEpoch(1002).setEndPoints(new BrokerEndpointCollection(
                List.of(new BrokerEndpoint().setName("PLAINTEXT").setHost("localhost").
                    setPort(9094).setSecurityProtocol((short) 0)).iterator())).
            setFeatures(new BrokerFeatureCollection(
                Set.of(new BrokerFeature().
                    setName(MetadataVersion.FEATURE_NAME).
                    setMinSupportedVersion(MetadataVersion.MINIMUM_VERSION.featureLevel()).
                    setMaxSupportedVersion(MetadataVersion.IBP_3_6_IV0.featureLevel())).iterator())).
            setRack("rack3"),
            REGISTER_BROKER_RECORD.highestSupportedVersion()));

        DELTA2 = new ClusterDelta(IMAGE2);
        RecordTestUtils.replayAll(DELTA2, DELTA2_RECORDS);

        Map<Integer, BrokerRegistration> map3 = new HashMap<>();
        map3.put(0, new BrokerRegistration.Builder().
            setId(0).
            setEpoch(1000).
            setIncarnationId(Uuid.fromString("vZKYST0pSA2HO5x_6hoO2Q")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9092))).
            setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 3))).
            setRack(Optional.empty()).
            setFenced(true).
            setInControlledShutdown(true).build());
        map3.put(1, new BrokerRegistration.Builder().
            setId(1).
            setEpoch(1001).
            setIncarnationId(Uuid.fromString("U52uRe20RsGI0RvpcTx33Q")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9093))).
            setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 3))).
            setRack(Optional.empty()).
            setFenced(false).
            setInControlledShutdown(false).build());
        map3.put(2, new BrokerRegistration.Builder().
            setId(2).
            setEpoch(1002).
            setIncarnationId(Uuid.fromString("Am5Yse7GQxaw0b2alM74bP")).
            setListeners(List.of(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9094))).
            setSupportedFeatures(Map.of("metadata.version",
                VersionRange.of(MetadataVersion.MINIMUM_VERSION.featureLevel(), MetadataVersion.IBP_3_6_IV0.featureLevel()))).
            setRack(Optional.of("rack3")).
            setFenced(true).
            setIsMigratingZkBroker(true).build());

        IMAGE3 = new ClusterImage(map3, cmap2);
    }

    @Test
    public void testEmptyImageRoundTrip() {
        testToImage(ClusterImage.EMPTY);
    }

    @Test
    public void testImage1RoundTrip() {
        testToImage(IMAGE1);
    }

    @Test
    public void testApplyDelta1() {
        assertEquals(IMAGE2, DELTA1.apply());
        // check image1 + delta1 = image2, since records for image1 + delta1 might differ from records from image2
        List<ApiMessageAndVersion> records = getImageRecords(IMAGE1);
        records.addAll(DELTA1_RECORDS);
        testToImage(IMAGE2, records);
    }

    @Test
    public void testImage2RoundTrip() {
        testToImage(IMAGE2);
    }

    @Test
    public void testApplyDelta2() {
        assertEquals(IMAGE3, DELTA2.apply());
        // check image2 + delta2 = image3, since records for image2 + delta2 might differ from records from image3
        List<ApiMessageAndVersion> records = getImageRecords(IMAGE2);
        records.addAll(DELTA2_RECORDS);
        testToImage(IMAGE3, records);
    }

    @Test
    public void testImage3RoundTrip() {
        testToImage(IMAGE3);
    }

    private static void testToImage(ClusterImage image) {
        testToImage(image, Optional.empty());
    }

    private static void testToImage(ClusterImage image, Optional<List<ApiMessageAndVersion>> fromRecords) {
        testToImage(image, fromRecords.orElseGet(() -> getImageRecords(image)));
    }

    private static void testToImage(ClusterImage image, List<ApiMessageAndVersion> fromRecords) {
        // test from empty image stopping each of the various intermediate images along the way
        new RecordTestUtils.TestThroughAllIntermediateImagesLeadingToFinalImageHelper<>(
            () -> ClusterImage.EMPTY,
            ClusterDelta::new
        ).test(image, fromRecords);
    }

    private static List<ApiMessageAndVersion> getImageRecords(ClusterImage image) {
        RecordListWriter writer = new RecordListWriter();
        image.write(writer, new ImageWriterOptions.Builder(MetadataVersion.latestProduction()).build());
        return writer.records();
    }

    @Test
    public void testHandleLossOfControllerRegistrations() {
        ClusterImage testImage = new ClusterImage(Map.of(),
            Map.of(1000, new ControllerRegistration.Builder().
                setId(1000).
                setIncarnationId(Uuid.fromString("9ABu6HEgRuS-hjHLgC4cHw")).
                setListeners(Map.of("PLAINTEXT",
                    new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 19092))).
                setSupportedFeatures(Map.of()).build()));
        RecordListWriter writer = new RecordListWriter();
        final AtomicReference<String> lossString = new AtomicReference<>("");
        testImage.write(writer, new ImageWriterOptions.Builder(MetadataVersion.IBP_3_6_IV2).
            setLossHandler(loss -> lossString.compareAndSet("", loss.loss())).
                build());
        assertEquals("controller registration data", lossString.get());
    }

    @Test
    public void testBrokerEpoch() {
        assertEquals(123L, IMAGE1.brokerEpoch(2));
    }

    @Test
    public void testBrokerEpochForNonExistentBroker() {
        assertEquals(-1L, IMAGE1.brokerEpoch(20));
    }

}
