package com.didiglobal.thriftmock.server;

import org.apache.thrift.AsyncProcessFunction;
import org.apache.thrift.TApplicationException;
import org.apache.thrift.TBase;
import org.apache.thrift.TException;
import org.apache.thrift.TSerializable;
import org.apache.thrift.async.AsyncMethodCallback;
import org.apache.thrift.protocol.TMessageType;
import org.apache.thrift.server.AbstractNonblockingServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.function.Function;

public class AsyncProcessFunctionMock extends AsyncProcessFunction implements Delay{

  private static final Logger LOG = LoggerFactory.getLogger(AsyncProcessFunctionMock.class);

  private String methodName;
  private int delay;
  private TBase args;
  private Function<TBase, TBase> mockResultFunction;

  public AsyncProcessFunctionMock(String methodName, TBase result) {
    this(methodName, result, 0);
  }

  public AsyncProcessFunctionMock(String methodName, TBase result, int delay) {
    this(methodName, result, a -> result, delay);
  }

  public AsyncProcessFunctionMock(String methodName,
                                  TBase args,
                                  Function<TBase, TBase> mockResultFunction,
                                  int delay) {
    super(methodName);
    this.methodName = methodName;
    this.delay = delay;
    this.args = args;
    this.mockResultFunction = mockResultFunction;
  }

  @Override
  protected boolean isOneway() {
    return false;
  }

  @Override
  public void start(Object iface, TBase args, AsyncMethodCallback resultHandler) throws TException {
    delay(delay);
    resultHandler.onComplete(new MockResult(methodName, mockResultFunction.apply(args)));
  }

  @Override
  public TBase getEmptyArgsInstance() {
    return args;
  }

  @Override
  public AsyncMethodCallback getResultHandler(AbstractNonblockingServer.AsyncFrameBuffer fb,
                                              int seqid) {
    return new AsyncMethodCallback() {
      @Override
      public void onComplete(Object response) {
        try {
          sendResponse(fb, (TSerializable) response, TMessageType.REPLY, seqid);
          return;
        } catch (Exception e) {
          LOG.error("Exception writing to internal frame buffer", e);
        }
        fb.close();
      }

      @Override
      public void onError(Exception exception) {
        byte msgType;
        TBase msg;
        msgType = TMessageType.EXCEPTION;
        msg = (TBase)new TApplicationException(TApplicationException.INTERNAL_ERROR,
                                               exception.getMessage());
        try {
          sendResponse(fb,msg,msgType,seqid);
          return;
        } catch (Exception ex) {
          LOG.error("Exception writing to internal frame buffer", ex);
        }
        fb.close();
      }
    };
  }
}