/*
 * Copyright 2017-present Open Networking Foundation
 *
 * Licensed 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 io.atomix.protocols.raft.proxy.impl;

import io.atomix.primitive.PrimitiveType;
import io.atomix.primitive.proxy.PrimitiveProxy;
import io.atomix.primitive.session.SessionId;

import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.function.Consumer;

import static com.google.common.base.Preconditions.checkNotNull;

/**
 * Client state.
 */
public final class RaftProxyState {
  private final String clientId;
  private final SessionId sessionId;
  private final String serviceName;
  private final PrimitiveType primitiveType;
  private final long timeout;
  private volatile PrimitiveProxy.State state = PrimitiveProxy.State.CONNECTED;
  private volatile long commandRequest;
  private volatile long commandResponse;
  private volatile long responseIndex;
  private volatile long eventIndex;
  private final Set<Consumer<PrimitiveProxy.State>> changeListeners = new CopyOnWriteArraySet<>();

  RaftProxyState(String clientId, SessionId sessionId, String serviceName, PrimitiveType primitiveType, long timeout) {
    this.clientId = clientId;
    this.sessionId = sessionId;
    this.serviceName = serviceName;
    this.primitiveType = primitiveType;
    this.timeout = timeout;
    this.responseIndex = sessionId.id();
    this.eventIndex = sessionId.id();
  }

  /**
   * Returns the client identifier.
   *
   * @return The client identifier.
   */
  public String getClientId() {
    return clientId;
  }

  /**
   * Returns the client session ID.
   *
   * @return The client session ID.
   */
  public SessionId getSessionId() {
    return sessionId;
  }

  /**
   * Returns the session name.
   *
   * @return The session name.
   */
  public String getPrimitiveName() {
    return serviceName;
  }

  /**
   * Returns the session type.
   *
   * @return The session type.
   */
  public PrimitiveType getPrimitiveType() {
    return primitiveType;
  }

  /**
   * Returns the session timeout.
   *
   * @return The session timeout.
   */
  public long getSessionTimeout() {
    return timeout;
  }

  /**
   * Returns the session state.
   *
   * @return The session state.
   */
  public PrimitiveProxy.State getState() {
    return state;
  }

  /**
   * Updates the session state.
   *
   * @param state The updates session state.
   */
  public void setState(PrimitiveProxy.State state) {
    if (this.state != state) {
      this.state = state;
      changeListeners.forEach(l -> l.accept(state));
    }
  }

  /**
   * Registers a state change listener on the session manager.
   *
   * @param listener The state change listener callback.
   */
  public void addStateChangeListener(Consumer<PrimitiveProxy.State> listener) {
    changeListeners.add(checkNotNull(listener));
  }

  /**
   * Removes a state change listener.
   *
   * @param listener the listener to remove
   */
  public void removeStateChangeListener(Consumer<PrimitiveProxy.State> listener) {
    changeListeners.remove(checkNotNull(listener));
  }

  /**
   * Sets the last command request sequence number.
   *
   * @param commandRequest The last command request sequence number.
   */
  public void setCommandRequest(long commandRequest) {
    this.commandRequest = commandRequest;
  }

  /**
   * Returns the last command request sequence number for the session.
   *
   * @return The last command request sequence number for the session.
   */
  public long getCommandRequest() {
    return commandRequest;
  }

  /**
   * Returns the next command request sequence number for the session.
   *
   * @return The next command request sequence number for the session.
   */
  public long nextCommandRequest() {
    return ++commandRequest;
  }

  /**
   * Sets the last command sequence number for which a response has been received.
   *
   * @param commandResponse The last command sequence number for which a response has been received.
   */
  public void setCommandResponse(long commandResponse) {
    this.commandResponse = commandResponse;
  }

  /**
   * Returns the last command sequence number for which a response has been received.
   *
   * @return The last command sequence number for which a response has been received.
   */
  public long getCommandResponse() {
    return commandResponse;
  }

  /**
   * Sets the highest index for which a response has been received.
   *
   * @param responseIndex The highest index for which a command or query response has been received.
   */
  public void setResponseIndex(long responseIndex) {
    this.responseIndex = Math.max(this.responseIndex, responseIndex);
  }

  /**
   * Returns the highest index for which a response has been received.
   *
   * @return The highest index for which a command or query response has been received.
   */
  public long getResponseIndex() {
    return responseIndex;
  }

  /**
   * Sets the highest index for which an event has been received in sequence.
   *
   * @param eventIndex The highest index for which an event has been received in sequence.
   */
  public void setEventIndex(long eventIndex) {
    this.eventIndex = eventIndex;
  }

  /**
   * Returns the highest index for which an event has been received in sequence.
   *
   * @return The highest index for which an event has been received in sequence.
   */
  public long getEventIndex() {
    return eventIndex;
  }
}

@@ -1287,16 +1287,55 @@
 NECTED;%0A
+  private volatile Long suspendedTime;%0A
   privat
@@ -3057,54 +3057,445 @@
     
-changeListeners.forEach(l -%3E l.accept(state));
+if (state == PrimitiveProxy.State.SUSPENDED) %7B%0A        if (suspendedTime == null) %7B%0A          suspendedTime = System.currentTimeMillis();%0A        %7D%0A      %7D else %7B%0A        suspendedTime = null;%0A      %7D%0A      changeListeners.forEach(l -%3E l.accept(state));%0A    %7D else if (this.state == PrimitiveProxy.State.SUSPENDED) %7B%0A      if (System.currentTimeMillis() - suspendedTime %3E timeout) %7B%0A        setState(PrimitiveProxy.State.CLOSED);%0A      %7D
 %0A   

Revert "Avoid exporing sessions on the client to ensure leaders can reset session timers after lengthy elections."