/*
 * 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.clients.consumer.internals;

import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData;

import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * This class holds the data that is needed to participate in the Streams rebalance protocol.
 */
public class StreamsRebalanceData {

    public static class TaskId implements Comparable<TaskId> {

        private final String subtopologyId;
        private final int partitionId;

        public TaskId(final String subtopologyId, final int partitionId) {
            this.subtopologyId = Objects.requireNonNull(subtopologyId, "Subtopology ID cannot be null");
            this.partitionId = partitionId;
        }

        public int partitionId() {
            return partitionId;
        }

        public String subtopologyId() {
            return subtopologyId;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            TaskId taskId = (TaskId) o;
            return partitionId == taskId.partitionId && Objects.equals(subtopologyId, taskId.subtopologyId);
        }

        @Override
        public int hashCode() {
            return Objects.hash(subtopologyId, partitionId);
        }

        @Override
        public int compareTo(TaskId taskId) {
            Objects.requireNonNull(taskId, "taskId cannot be null");
            return Comparator.comparing(TaskId::subtopologyId)
                .thenComparingInt(TaskId::partitionId).compare(this, taskId);
        }

        @Override
        public String toString() {
            return subtopologyId + "_" + partitionId;
        }
    }

    public static class HostInfo {

        private final String host;
        private final int port;

        public HostInfo(final String host, final int port) {
            this.host = Objects.requireNonNull(host);
            this.port = port;
        }

        public String host() {
            return host;
        }

        public int port() {
            return port;
        }

        @Override
        public boolean equals(Object o) {
            if (o == null || getClass() != o.getClass()) return false;
            HostInfo hostInfo = (HostInfo) o;
            return port == hostInfo.port && Objects.equals(host, hostInfo.host);
        }

        @Override
        public int hashCode() {
            return Objects.hash(host, port);
        }

        @Override
        public String toString() {
            return "HostInfo{" +
                "host='" + host + '\'' +
                ", port=" + port +
                '}';
        }

    }

    public static class Assignment {

        public static final Assignment EMPTY = new Assignment();

        private final Set<TaskId> activeTasks;

        private final Set<TaskId> standbyTasks;

        private final Set<TaskId> warmupTasks;

        private Assignment() {
            this.activeTasks = Set.of();
            this.standbyTasks = Set.of();
            this.warmupTasks = Set.of();
        }

        public Assignment(final Set<TaskId> activeTasks,
                          final Set<TaskId> standbyTasks,
                          final Set<TaskId> warmupTasks) {
            this.activeTasks = Set.copyOf(Objects.requireNonNull(activeTasks, "Active tasks cannot be null"));
            this.standbyTasks = Set.copyOf(Objects.requireNonNull(standbyTasks, "Standby tasks cannot be null"));
            this.warmupTasks = Set.copyOf(Objects.requireNonNull(warmupTasks, "Warmup tasks cannot be null"));
        }

        public Set<TaskId> activeTasks() {
            return activeTasks;
        }

        public Set<TaskId> standbyTasks() {
            return standbyTasks;
        }

        public Set<TaskId> warmupTasks() {
            return warmupTasks;
        }

        @Override
        public boolean equals(final Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            final Assignment that = (Assignment) o;
            return Objects.equals(activeTasks, that.activeTasks)
                && Objects.equals(standbyTasks, that.standbyTasks)
                && Objects.equals(warmupTasks, that.warmupTasks);
        }

        @Override
        public int hashCode() {
            return Objects.hash(activeTasks, standbyTasks, warmupTasks);
        }

        public Assignment copy() {
            return new Assignment(activeTasks, standbyTasks, warmupTasks);
        }

        @Override
        public String toString() {
            return "Assignment{" +
                "activeTasks=" + activeTasks +
                ", standbyTasks=" + standbyTasks +
                ", warmupTasks=" + warmupTasks +
                '}';
        }
    }

    public static class Subtopology {

        private final Set<String> sourceTopics;
        private final Set<String> repartitionSinkTopics;
        private final Map<String, TopicInfo> stateChangelogTopics;
        private final Map<String, TopicInfo> repartitionSourceTopics;
        private final Collection<Set<String>> copartitionGroups;

        public Subtopology(final Set<String> sourceTopics,
                           final Set<String> repartitionSinkTopics,
                           final Map<String, TopicInfo> repartitionSourceTopics,
                           final Map<String, TopicInfo> stateChangelogTopics,
                           final Collection<Set<String>> copartitionGroups
        ) {
            this.sourceTopics = Set.copyOf(Objects.requireNonNull(sourceTopics, "Subtopology ID cannot be null"));
            this.repartitionSinkTopics =
                Set.copyOf(Objects.requireNonNull(repartitionSinkTopics, "Repartition sink topics cannot be null"));
            this.repartitionSourceTopics =
                Map.copyOf(Objects.requireNonNull(repartitionSourceTopics, "Repartition source topics cannot be null"));
            this.stateChangelogTopics =
                Map.copyOf(Objects.requireNonNull(stateChangelogTopics, "State changelog topics cannot be null"));
            this.copartitionGroups =
                Collections.unmodifiableCollection(Objects.requireNonNull(
                    copartitionGroups,
                    "Co-partition groups cannot be null"
                    )
                );
        }

        public Set<String> sourceTopics() {
            return sourceTopics;
        }

        public Set<String> repartitionSinkTopics() {
            return repartitionSinkTopics;
        }

        public Map<String, TopicInfo> stateChangelogTopics() {
            return stateChangelogTopics;
        }

        public Map<String, TopicInfo> repartitionSourceTopics() {
            return repartitionSourceTopics;
        }

        public Collection<Set<String>> copartitionGroups() {
            return copartitionGroups;
        }

        @Override
        public String toString() {
            return "Subtopology{" +
                "sourceTopics=" + sourceTopics +
                ", repartitionSinkTopics=" + repartitionSinkTopics +
                ", stateChangelogTopics=" + stateChangelogTopics +
                ", repartitionSourceTopics=" + repartitionSourceTopics +
                ", copartitionGroups=" + copartitionGroups +
                '}';
        }
    }

    public static class TopicInfo {

        private final Optional<Integer> numPartitions;
        private final Optional<Short> replicationFactor;
        private final Map<String, String> topicConfigs;

        public TopicInfo(final Optional<Integer> numPartitions,
                         final Optional<Short> replicationFactor,
                         final Map<String, String> topicConfigs) {
            this.numPartitions = Objects.requireNonNull(numPartitions, "Number of partitions cannot be null");
            this.replicationFactor = Objects.requireNonNull(replicationFactor, "Replication factor cannot be null");
            this.topicConfigs =
                Map.copyOf(Objects.requireNonNull(topicConfigs, "Additional topic configs cannot be null"));
        }

        public Optional<Integer> numPartitions() {
            return numPartitions;
        }

        public Optional<Short> replicationFactor() {
            return replicationFactor;
        }

        public Map<String, String> topicConfigs() {
            return topicConfigs;
        }

        @Override
        public String toString() {
            return "TopicInfo{" +
                "numPartitions=" + numPartitions +
                ", replicationFactor=" + replicationFactor +
                ", topicConfigs=" + topicConfigs +
                '}';
        }
    }

    private final UUID processId;

    private final Optional<HostInfo> endpoint;

    private final Map<String, String> clientTags;

    private final Map<String, Subtopology> subtopologies;

    private final AtomicReference<Assignment> reconciledAssignment = new AtomicReference<>(Assignment.EMPTY);

    private final AtomicReference<Map<HostInfo, List<TopicPartition>>> partitionsByHost = new AtomicReference<>(Collections.emptyMap());

    private final AtomicBoolean shutdownRequested = new AtomicBoolean(false);

    private final AtomicReference<List<StreamsGroupHeartbeatResponseData.Status>> statuses = new AtomicReference<>(List.of());

    public StreamsRebalanceData(final UUID processId,
                                final Optional<HostInfo> endpoint,
                                final Map<String, Subtopology> subtopologies,
                                final Map<String, String> clientTags) {
        this.processId = Objects.requireNonNull(processId, "Process ID cannot be null");
        this.endpoint = Objects.requireNonNull(endpoint, "Endpoint cannot be null");
        this.subtopologies = Map.copyOf(Objects.requireNonNull(subtopologies, "Subtopologies cannot be null"));
        this.clientTags = Map.copyOf(Objects.requireNonNull(clientTags, "Client tags cannot be null"));
    }

    public UUID processId() {
        return processId;
    }

    public Optional<HostInfo> endpoint() {
        return endpoint;
    }

    public Map<String, String> clientTags() {
        return clientTags;
    }

    public Map<String, Subtopology> subtopologies() {
        return subtopologies;
    }

    public int topologyEpoch() {
        return 0;
    }

    public void setReconciledAssignment(final Assignment assignment) {
        reconciledAssignment.set(assignment);
    }

    public Assignment reconciledAssignment() {
        return reconciledAssignment.get();
    }

    public void setPartitionsByHost(final Map<HostInfo, List<TopicPartition>> partitionsByHost) {
        this.partitionsByHost.set(partitionsByHost);
    }

    public Map<HostInfo, List<TopicPartition>> partitionsByHost() {
        return partitionsByHost.get();
    }

    /** For the current stream thread to request a shutdown of all Streams clients belonging to the same application. */
    public void requestShutdown() {
        shutdownRequested.set(true);
    }

    /** True if the current stream thread requested a shutdown of all Streams clients belonging to the same application. */
    public boolean shutdownRequested() {
        return shutdownRequested.get();
    }

    /** Updated whenever the status of the streams group is updated. */
    public void setStatuses(final List<StreamsGroupHeartbeatResponseData.Status> s) {
        statuses.set(s);
    }

    /** For communicating the current status of the group to the stream thread */
    public List<StreamsGroupHeartbeatResponseData.Status> statuses() {
        return statuses.get();
    }

}
