

package com.linkedin.kafka.cruisecontrol.model;

import com.linkedin.kafka.cruisecontrol.exception.ModelInputException;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.apache.kafka.common.TopicPartition;


/**
 * A partition consists of replicas with the same topic partition. One replica is the leader, the rest of the replicas
 * are followers.
 */
public class Partition implements Serializable {
  private final TopicPartition _tp;
  private final List<Replica> _replicas;
  private Replica _leader;

  /**
   * Constructor for Partition class.
   *
   * @param tp Topic partition information for the replica in this partition,
   * @throws ModelInputException
   */
  Partition(TopicPartition tp) throws ModelInputException {
    _tp = tp;
    _replicas = new ArrayList<>();
    _leader = null;
  }

  /**
   * @return the topic partition of this partition.
   */
  public TopicPartition topicPartition() {
    return _tp;
  }

  /**
   * Add follower to the partition.
   *
   * @param follower Follower replica.
   * @param index the index the follower should be at.
   *
   * @throws ModelInputException
   */
  void addFollower(Replica follower, int index) throws ModelInputException {
    if (follower.isLeader()) {
      throw new ModelInputException("Inconsistent leadership information. Trying to add follower replica " +
                                        follower + " while it is a leader.");
    }
    if (!follower.topicPartition().equals(_tp)) {
      throw new ModelInputException("Inconsistent topic partition. Trying to add follower replica " + follower +
                                        " to partition " + _tp + ".");
    }
    // Add follower to the list of followers.
    _replicas.add(index, follower);
  }

  /**
   * @return The replica list.
   */
  public List<Replica> replicas() {
    return Collections.unmodifiableList(_replicas);
  }

  /**
   * Get follower replicas.
   */
  public List<Replica> followers() {
    List<Replica> followers = new ArrayList<>();
    _replicas.forEach(r -> {
      if (!r.isLeader()) {
        followers.add(r);
      }
    });
    return followers;
  }

  /**
   * Get the leader replica.
   */
  public Replica leader() {
    return _leader;
  }

  /**
   * Get replica with the given broker id in this partition.
   *
   * @param brokerId Broker id of the requested replica in this partition.
   * @return Replica with the given broker id in this partition.
   * @throws ModelInputException
   */
  Replica replica(long brokerId)
      throws ModelInputException {
    for (Replica replica : _replicas) {
      if (replica.broker().id() == brokerId) {
        return replica;
      }
    }

    throw new ModelInputException("Requested replica " + brokerId + " is not a replica of partition " + _tp);
  }

  /**
   * Get the set of brokers that followers reside in.
   */
  public List<Broker> followerBrokers() {
    List<Broker> followerBrokers = new ArrayList<>();
    _replicas.forEach(r -> {
      if (!r.isLeader()) {
        followerBrokers.add(r.broker());
      }
    });
    return followerBrokers;
  }

  /**
   * Given two follower indices in the replica list, swap their positions.
   *
   * @param index1 The index of the first follower to be swapped.
   * @param index2 The index of the second follower to be swapped
   */
  public void swapFollowerPositions(int index1, int index2) {
    Replica follower1 = _replicas.get(index1);
    Replica follower2 = _replicas.get(index2);

    if (follower1.isLeader() || follower2.isLeader()) {
      throw new IllegalArgumentException(String.format("%s is not a follower.",
                                                       follower1.isLeader() ? follower1 : follower2));
    }
    _replicas.set(index2, follower1);
    _replicas.set(index1, follower2);
  }

  /**
   * Given two replica indices in the replica list, swap their positions.
   *
   * @param index1 The index of the first replica to be swapped.
   * @param index2 The index of the second replica to be swapped
   */
  public void swapReplicaPositions(int index1, int index2) {
    Replica replica1 = _replicas.get(index1);
    Replica replica2 = _replicas.get(index2);

    _replicas.set(index2, replica1);
    _replicas.set(index1, replica2);
  }

  /**
   * Get the set of brokers that contain replicas of the partition.
   */
  public Set<Broker> partitionBrokers() {
    Set<Broker> partitionBrokers = new HashSet<>();
    _replicas.forEach(r -> partitionBrokers.add(r.broker()));
    return partitionBrokers;
  }

  /**
   * @return the set of racks that contains replicas of the partition.
   */
  public Set<Rack> partitionRacks() {
    Set<Rack> partitionRacks = new HashSet<>();
    _replicas.forEach(r -> partitionRacks.add(r.broker().rack()));
    return partitionRacks;
  }

  /**
   * Set the leader to the value specified by the leader parameter.
   *
   * @param leader Leader replica of partition.
   * @param index the index the leader replica should be at.
   * @throws ModelInputException
   */
  void addLeader(Replica leader, int index) throws ModelInputException {
    if (_leader != null) {
      throw new ModelInputException(String.format("Partition %s already has a leader replica %s. Cannot "
                                                      + "add a new leader replica %s",
                                                  _tp, _leader, leader));
    }
    if (!leader.isLeader()) {
      throw new ModelInputException("Inconsistent leadership information. Trying to set " + leader.broker() +
                                    " as the leader for partition " + _tp + " while the replica is not marked as a leader.");
    }
    _leader = leader;
    _replicas.add(index, leader);
  }

  /**
   * Relocate leadership.
   * Move the leader replica to the head of the replica list.
   *
   * @param prospectiveLeader Prospective leader.
   */
  void relocateLeadership(Replica prospectiveLeader) throws ModelInputException {
    int leaderPos = _replicas.indexOf(prospectiveLeader);
    swapReplicaPositions(0, leaderPos);
    _leader = prospectiveLeader;
  }

  /**
   * Clear the leader to null and clear followers.
   */
  public void clear() {
    _replicas.clear();
    _leader = null;
  }

  /**
   * Get string representation of Partition in XML format.
   */
  @Override
  public String toString() {
    StringBuilder partition = new StringBuilder().append(String.format("<Partition>%n<Leader>%s</Leader>%n", _leader));

    for (Replica replica : _replicas) {
      if (!replica.isLeader()) {
        partition.append(String.format("<Follower>%s</Follower>%n", replica));
      }
    }
    return partition.append("</Partition>%n").toString();
  }
}