/*
 * 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.ratis.examples.filestore;

import org.apache.ratis.conf.RaftProperties;
import org.apache.ratis.proto.ExamplesProtos;
import org.apache.ratis.proto.ExamplesProtos.DeleteReplyProto;
import org.apache.ratis.proto.ExamplesProtos.DeleteRequestProto;
import org.apache.ratis.proto.ExamplesProtos.FileStoreRequestProto;
import org.apache.ratis.proto.ExamplesProtos.ReadRequestProto;
import org.apache.ratis.proto.ExamplesProtos.StreamWriteRequestProto;
import org.apache.ratis.proto.ExamplesProtos.WriteRequestHeaderProto;
import org.apache.ratis.proto.ExamplesProtos.WriteRequestProto;
import org.apache.ratis.proto.RaftProtos;
import org.apache.ratis.proto.RaftProtos.LogEntryProto;
import org.apache.ratis.protocol.Message;
import org.apache.ratis.protocol.RaftClientRequest;
import org.apache.ratis.protocol.RaftGroupId;
import org.apache.ratis.protocol.RaftPeerId;
import org.apache.ratis.server.RaftServer;
import org.apache.ratis.server.storage.RaftStorage;
import org.apache.ratis.statemachine.StateMachineStorage;
import org.apache.ratis.statemachine.TransactionContext;
import org.apache.ratis.statemachine.impl.BaseStateMachine;
import org.apache.ratis.statemachine.impl.SimpleStateMachineStorage;
import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
import org.apache.ratis.thirdparty.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.ratis.util.FileUtils;

import java.io.IOException;
import java.nio.channels.WritableByteChannel;
import java.nio.file.Path;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;

public class FileStoreStateMachine extends BaseStateMachine {
  static class SimulateFailure {
    static final String PATH_TO_FAIL = "path-to-fail";
    private static final AtomicReference<RaftPeerId> CHOSEN = new AtomicReference<>();

    static boolean chooseServer(RaftPeerId serverId) {
      if (SimulateFailure.CHOSEN.compareAndSet(null, serverId)) {
        LOG.info("Server {} is chosen", serverId);
        return true;
      } else {
        return serverId.equals(SimulateFailure.CHOSEN.get());
      }
    }
  }

  private final SimpleStateMachineStorage storage = new SimpleStateMachineStorage();

  private final FileStore files;

  public FileStoreStateMachine(RaftProperties properties) {
    this.files = new FileStore(this::getId, properties);
  }

  @Override
  public void initialize(RaftServer server, RaftGroupId groupId, RaftStorage raftStorage)
      throws IOException {
    super.initialize(server, groupId, raftStorage);
    this.storage.init(raftStorage);
    for (Path path : files.getRoots()) {
      FileUtils.createDirectories(path);
    }
  }

  @Override
  public StateMachineStorage getStateMachineStorage() {
    return storage;
  }

  @Override
  public void close() {
    files.close();
    setLastAppliedTermIndex(null);
  }

  @Override
  public CompletableFuture<Message> query(Message request) {
    final ReadRequestProto proto;
    try {
      proto = ReadRequestProto.parseFrom(request.getContent());
    } catch (InvalidProtocolBufferException e) {
      return FileStoreCommon.completeExceptionally("Failed to parse " + request, e);
    }

    final String path = proto.getPath().toStringUtf8();
    return (proto.getIsWatch()? files.watch(path)
        : files.read(path, proto.getOffset(), proto.getLength(), true))
        .thenApply(reply -> Message.valueOf(reply.toByteString()));
  }

  @Override
  public long transferTo(Message request, WritableByteChannel stream) throws IOException {
    try {
      final ReadRequestProto proto = ReadRequestProto.parseFrom(request.getContent());
      if (proto.getIsWatch()) {
        throw new IOException("Watch is not supported for streaming read: " + proto);
      }
      final long length = proto.getLength();
      files.streamRead(proto.getPath().toStringUtf8(), proto.getOffset(), length, stream);
      return length;
    } catch (IOException e) {
      LOG.error(getId() + ": Failed streaming read for " + request, e);
      try {
        stream.close();
      } catch (IOException ignored) {
        // ignore
      }
      throw e;
    }
  }

  @Override
  public TransactionContext startTransaction(RaftClientRequest request) throws IOException {
    final ByteString content = request.getMessage().getContent();
    final FileStoreRequestProto proto = FileStoreRequestProto.parseFrom(content);
    final TransactionContext.Builder b = TransactionContext.newBuilder()
        .setStateMachine(this)
        .setClientRequest(request);
    if (proto.getRequestCase() == FileStoreRequestProto.RequestCase.WRITE) {
      final WriteRequestProto write = proto.getWrite();
      final FileStoreRequestProto newProto = FileStoreRequestProto.newBuilder()
          .setWriteHeader(write.getHeader()).build();
      b.setLogData(newProto.toByteString()).setStateMachineData(write.getData())
       .setStateMachineContext(newProto);
    } else {
      b.setLogData(content)
       .setStateMachineContext(proto);
    }
    return b.build();
  }

  @Override
  public TransactionContext startTransaction(LogEntryProto entry, RaftProtos.RaftPeerRole role) {
    return TransactionContext.newBuilder()
        .setStateMachine(this)
        .setLogEntry(entry)
        .setServerRole(role)
        .setStateMachineContext(getProto(entry))
        .build();
  }

  void simulateFailure(String path) throws Exception {
    final RaftServer.Division division = getServer().get().getDivision(getGroupId());
    if (path.equals(SimulateFailure.PATH_TO_FAIL)
        && division.getInfo().isFollower()
        && SimulateFailure.chooseServer(division.getId())) {
      throw new IOException(getId() + ": Simulated failure for path " + path);
    }
  }

  @Override
  public CompletableFuture<Integer> write(LogEntryProto entry, TransactionContext context) {
    final FileStoreRequestProto proto = getProto(context, entry);
    if (proto.getRequestCase() != FileStoreRequestProto.RequestCase.WRITEHEADER) {
      return null;
    }

    final WriteRequestHeaderProto h = proto.getWriteHeader();
    try {
      simulateFailure(h.getPath().toStringUtf8());
    } catch (Exception e) {
      return FileStoreCommon.completeExceptionally(entry.getIndex(), getId() + ": Failed simulateFailure", e);
    }

    final CompletableFuture<Integer> f = files.write(entry.getIndex(),
        h.getPath().toStringUtf8(), h.getClose(),  h.getSync(), h.getOffset(),
        entry.getStateMachineLogEntry().getStateMachineEntry().getStateMachineData());
    // sync only if closing the file
    return h.getClose()? f: null;
  }

  static FileStoreRequestProto getProto(TransactionContext context, LogEntryProto entry) {
    if (context != null) {
      final FileStoreRequestProto proto = (FileStoreRequestProto) context.getStateMachineContext();
      if (proto != null) {
        return proto;
      }
    }
    return getProto(entry);
  }

  static FileStoreRequestProto getProto(LogEntryProto entry) {
    try {
      return FileStoreRequestProto.parseFrom(entry.getStateMachineLogEntry().getLogData());
    } catch (InvalidProtocolBufferException e) {
      throw new IllegalArgumentException("Failed to parse data, entry=" + entry, e);
    }
  }

  @Override
  public CompletableFuture<ByteString> read(LogEntryProto entry, TransactionContext context) {
    final FileStoreRequestProto proto = getProto(context, entry);
    if (proto.getRequestCase() != FileStoreRequestProto.RequestCase.WRITEHEADER) {
      return null;
    }

    final WriteRequestHeaderProto h = proto.getWriteHeader();
    CompletableFuture<ExamplesProtos.ReadReplyProto> reply =
        files.read(h.getPath().toStringUtf8(), h.getOffset(), h.getLength(), false);

    return reply.thenApply(ExamplesProtos.ReadReplyProto::getData);
  }

  static class LocalStream implements DataStream {
    private final DataChannel dataChannel;

    LocalStream(DataChannel dataChannel) {
      this.dataChannel = dataChannel;
    }

    @Override
    public DataChannel getDataChannel() {
      return dataChannel;
    }

    @Override
    public CompletableFuture<?> cleanUp() {
      return CompletableFuture.supplyAsync(() -> {
        try {
          dataChannel.close();
          return true;
        } catch (IOException e) {
          return FileStoreCommon.completeExceptionally("Failed to close data channel", e);
        }
      });
    }
  }

  @Override
  public CompletableFuture<DataStream> stream(RaftClientRequest request) {
    final ByteString reqByteString = request.getMessage().getContent();
    final FileStoreRequestProto proto;
    try {
      proto = FileStoreRequestProto.parseFrom(reqByteString);
    } catch (InvalidProtocolBufferException e) {
      return FileStoreCommon.completeExceptionally(
          "Failed to parse stream header", e);
    }
    return files.createDataChannel(proto.getStream().getPath().toStringUtf8())
        .thenApply(LocalStream::new);
  }

  @Override
  public CompletableFuture<?> link(DataStream stream, LogEntryProto entry) {
    LOG.info("linking {}", stream);
    return files.streamLink(stream);
  }

  @Override
  public CompletableFuture<Message> applyTransaction(TransactionContext trx) {
    final LogEntryProto entry = trx.getLogEntry();

    final long index = entry.getIndex();
    updateLastAppliedTermIndex(entry.getTerm(), index);

    final FileStoreRequestProto request = getProto(trx, entry);

    switch(request.getRequestCase()) {
      case DELETE:
        return delete(index, request.getDelete());
      case WRITEHEADER:
        return writeCommit(index, request.getWriteHeader(),
            entry.getStateMachineLogEntry().getStateMachineEntry().getStateMachineData().size());
      case STREAM:
        return streamCommit(request.getStream());
      case WRITE:
        // WRITE should not happen here since
        // startTransaction converts WRITE requests to WRITEHEADER requests.
      default:
        LOG.error(getId() + ": Unexpected request case " + request.getRequestCase());
        return FileStoreCommon.completeExceptionally(index,
            "Unexpected request case " + request.getRequestCase());
    }
  }

  private CompletableFuture<Message> writeCommit(
      long index, WriteRequestHeaderProto header, int size) {
    final String path = header.getPath().toStringUtf8();
    return files.submitCommit(index, path, header.getClose(), header.getOffset(), size)
        .thenApply(reply -> Message.valueOf(reply.toByteString()));
  }

  private CompletableFuture<Message> streamCommit(StreamWriteRequestProto stream) {
    final String path = stream.getPath().toStringUtf8();
    final long size = stream.getLength();
    return files.streamCommit(path, size).thenApply(reply -> Message.valueOf(reply.toByteString()));
  }

  private CompletableFuture<Message> delete(long index, DeleteRequestProto request) {
    final String path = request.getPath().toStringUtf8();
    return files.delete(index, path).thenApply(resolved ->
        Message.valueOf(DeleteReplyProto.newBuilder().setResolvedPath(
            FileStoreCommon.toByteString(resolved)).build().toByteString(),
            () -> "Message:" + resolved));
  }
}
