/**
 * 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.apex.malhar.flume.operator;

import java.io.IOException;
import java.io.Serializable;
import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ArrayBlockingQueue;

import javax.validation.constraints.Min;
import javax.validation.constraints.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.apex.malhar.flume.discovery.Discovery;
import org.apache.apex.malhar.flume.discovery.ZKAssistedDiscovery;
import org.apache.apex.malhar.flume.sink.Server;

import org.apache.flume.Event;

import com.datatorrent.api.Context;
import com.datatorrent.api.Context.OperatorContext;
import com.datatorrent.api.DefaultOutputPort;
import com.datatorrent.api.DefaultPartition;
import com.datatorrent.api.InputOperator;
import com.datatorrent.api.Operator;
import com.datatorrent.api.Partitioner;
import com.datatorrent.api.Stats.OperatorStats;
import com.datatorrent.api.StreamCodec;
import com.datatorrent.netlet.AbstractLengthPrependerClient;
import com.datatorrent.netlet.DefaultEventLoop;
import com.datatorrent.netlet.util.Slice;

import static java.lang.Thread.sleep;

/**
 * <p>
 * Abstract AbstractFlumeInputOperator class.</p>
 *
 * @param <T> Type of the output payload.
 * @since 0.9.2
 */
public abstract class AbstractFlumeInputOperator<T>
    implements InputOperator, Operator.ActivationListener<OperatorContext>, Operator.IdleTimeHandler,
    Operator.CheckpointListener, Partitioner<AbstractFlumeInputOperator<T>>
{
  public final transient DefaultOutputPort<T> output = new DefaultOutputPort<T>();
  public final transient DefaultOutputPort<Slice> drop = new DefaultOutputPort<Slice>();
  @NotNull
  private String[] connectionSpecs;
  @NotNull
  private StreamCodec<Event> codec;
  private final ArrayList<RecoveryAddress> recoveryAddresses;
  @SuppressWarnings("FieldMayBeFinal") // it's not final because that mucks with the serialization somehow
  private transient ArrayBlockingQueue<Slice> handoverBuffer;
  private transient int idleCounter;
  private transient int eventCounter;
  private transient DefaultEventLoop eventloop;
  private transient volatile boolean connected;
  private transient OperatorContext context;
  private transient Client client;
  private transient long windowId;
  private transient byte[] address;
  @Min(0)
  private long maxEventsPerSecond;
  //This is calculated from maxEventsPerSecond, App window count and streaming window size
  private transient long maxEventsPerWindow;

  public AbstractFlumeInputOperator()
  {
    handoverBuffer = new ArrayBlockingQueue<Slice>(1024 * 5);
    connectionSpecs = new String[0];
    recoveryAddresses = new ArrayList<RecoveryAddress>();
    maxEventsPerSecond = Long.MAX_VALUE;
  }

  @Override
  public void setup(OperatorContext context)
  {
    long windowDurationMillis = context.getValue(OperatorContext.APPLICATION_WINDOW_COUNT) *
        context.getValue(Context.DAGContext.STREAMING_WINDOW_SIZE_MILLIS);
    maxEventsPerWindow = (long)(windowDurationMillis / 1000.0 * maxEventsPerSecond);
    logger.debug("max-events per-second {} per-window {}", maxEventsPerSecond, maxEventsPerWindow);

    try {
      eventloop = new DefaultEventLoop("EventLoop-" + context.getId());
      eventloop.start();
    } catch (IOException ex) {
      throw new RuntimeException(ex);
    }
  }

  @Override
  @SuppressWarnings({"unchecked"})
  public void activate(OperatorContext ctx)
  {
    if (connectionSpecs.length == 0) {
      logger.info("Discovered zero FlumeSink");
    } else if (connectionSpecs.length == 1) {
      for (String connectAddresse: connectionSpecs) {
        logger.debug("Connection spec is {}", connectAddresse);
        String[] parts = connectAddresse.split(":");
        eventloop.connect(new InetSocketAddress(parts[1], Integer.parseInt(parts[2])), client = new Client(parts[0]));
      }
    } else {
      throw new IllegalArgumentException(
          String.format("A physical %s operator cannot connect to more than 1 addresses!",
              this.getClass().getSimpleName()));
    }

    context = ctx;
  }

  @Override
  public void beginWindow(long windowId)
  {
    this.windowId = windowId;
    idleCounter = 0;
    eventCounter = 0;
  }

  @Override
  public void emitTuples()
  {
    int i = handoverBuffer.size();
    if (i > 0 && eventCounter < maxEventsPerWindow) {

      while (--i > 0 && eventCounter < maxEventsPerWindow - 1) {
        final Slice slice = handoverBuffer.poll();
        slice.offset += 8;
        slice.length -= 8;
        T convert = convert((Event)codec.fromByteArray(slice));
        if (convert == null) {
          drop.emit(slice);
        } else {
          output.emit(convert);
        }
        eventCounter++;
      }

      final Slice slice = handoverBuffer.poll();
      slice.offset += 8;
      slice.length -= 8;
      T convert = convert((Event)codec.fromByteArray(slice));
      if (convert == null) {
        drop.emit(slice);
      } else {
        output.emit(convert);
      }
      eventCounter++;

      address = Arrays.copyOfRange(slice.buffer, slice.offset - 8, slice.offset);
    }
  }

  @Override
  public void endWindow()
  {
    if (connected) {
      byte[] array = new byte[Server.Request.FIXED_SIZE];

      array[0] = Server.Command.WINDOWED.getOrdinal();
      Server.writeInt(array, 1, eventCounter);
      Server.writeInt(array, 5, idleCounter);
      Server.writeLong(array, Server.Request.TIME_OFFSET, System.currentTimeMillis());

      logger.debug("wrote {} with eventCounter = {} and idleCounter = {}", Server.Command.WINDOWED, eventCounter, idleCounter);
      client.write(array);
    }

    if (address != null) {
      RecoveryAddress rAddress = new RecoveryAddress();
      rAddress.address = address;
      address = null;
      rAddress.windowId = windowId;
      recoveryAddresses.add(rAddress);
    }
  }

  @Override
  public void deactivate()
  {
    if (connected) {
      eventloop.disconnect(client);
    }
    context = null;
  }

  @Override
  public void teardown()
  {
    eventloop.stop();
    eventloop = null;
  }

  @Override
  public void handleIdleTime()
  {
    idleCounter++;
    try {
      sleep(context.getValue(OperatorContext.SPIN_MILLIS));
    } catch (InterruptedException ex) {
      throw new RuntimeException(ex);
    }
  }

  public abstract T convert(Event event);

  /**
   * @return the connectAddress
   */
  public String[] getConnectAddresses()
  {
    return connectionSpecs.clone();
  }

  /**
   * @param specs - sinkid:host:port specification of all the sinks.
   */
  public void setConnectAddresses(String[] specs)
  {
    this.connectionSpecs = specs.clone();
  }

  /**
   * @return the codec
   */
  public StreamCodec<Event> getCodec()
  {
    return codec;
  }

  /**
   * @param codec the codec to set
   */
  public void setCodec(StreamCodec<Event> codec)
  {
    this.codec = codec;
  }

  private static class RecoveryAddress implements Serializable
  {
    long windowId;
    byte[] address;

    @Override
    public String toString()
    {
      return "RecoveryAddress{" + "windowId=" + windowId + ", address=" + Arrays.toString(address) + '}';
    }

    @Override
    public boolean equals(Object o)
    {
      if (this == o) {
        return true;
      }
      if (!(o instanceof RecoveryAddress)) {
        return false;
      }

      RecoveryAddress that = (RecoveryAddress)o;

      if (windowId != that.windowId) {
        return false;
      }
      return Arrays.equals(address, that.address);
    }

    @Override
    public int hashCode()
    {
      int result = (int)(windowId ^ (windowId >>> 32));
      result = 31 * result + (address != null ? Arrays.hashCode(address) : 0);
      return result;
    }

    private static final long serialVersionUID = 201312021432L;
  }

  @Override
  public void checkpointed(long windowId)
  {
    /* dont do anything */
  }

  @Override
  public void committed(long windowId)
  {
    if (!connected) {
      return;
    }

    synchronized (recoveryAddresses) {
      byte[] addr = null;

      Iterator<RecoveryAddress> iterator = recoveryAddresses.iterator();
      while (iterator.hasNext()) {
        RecoveryAddress ra = iterator.next();
        if (ra.windowId > windowId) {
          break;
        }

        iterator.remove();
        if (ra.address != null) {
          addr = ra.address;
        }
      }

      if (addr != null) {
        /*
         * Make sure that we store the last valid address processed
         */
        if (recoveryAddresses.isEmpty()) {
          RecoveryAddress ra = new RecoveryAddress();
          ra.address = addr;
          recoveryAddresses.add(ra);
        }

        int arraySize = 1/* for the type of the message */
            + 8 /* for the location to commit */
            + 8 /* for storing the current time stamp*/;
        byte[] array = new byte[arraySize];

        array[0] = Server.Command.COMMITTED.getOrdinal();
        System.arraycopy(addr, 0, array, 1, 8);
        Server.writeLong(array, Server.Request.TIME_OFFSET, System.currentTimeMillis());
        logger.debug("wrote {} with recoveryOffset = {}", Server.Command.COMMITTED, Arrays.toString(addr));
        client.write(array);
      }
    }
  }

  @Override
  public Collection<Partition<AbstractFlumeInputOperator<T>>> definePartitions(
      Collection<Partition<AbstractFlumeInputOperator<T>>> partitions, PartitioningContext context)
  {
    Collection<Discovery.Service<byte[]>> discovered = discoveredFlumeSinks.get();
    if (discovered == null) {
      return partitions;
    }

    HashMap<String, ArrayList<RecoveryAddress>> allRecoveryAddresses = abandonedRecoveryAddresses.get();
    ArrayList<String> allConnectAddresses = new ArrayList<String>(partitions.size());
    for (Partition<AbstractFlumeInputOperator<T>> partition: partitions) {
      String[] lAddresses = partition.getPartitionedInstance().connectionSpecs;
      allConnectAddresses.addAll(Arrays.asList(lAddresses));
      for (int i = lAddresses.length; i-- > 0;) {
        String[] parts = lAddresses[i].split(":", 2);
        allRecoveryAddresses.put(parts[0], partition.getPartitionedInstance().recoveryAddresses);
      }
    }

    HashMap<String, String> connections = new HashMap<String, String>(discovered.size());
    for (Discovery.Service<byte[]> service: discovered) {
      String previousSpec = connections.get(service.getId());
      String newspec = service.getId() + ':' + service.getHost() + ':' + service.getPort();
      if (previousSpec == null) {
        connections.put(service.getId(), newspec);
      } else {
        boolean found = false;
        for (ConnectionStatus cs: partitionedInstanceStatus.get().values()) {
          if (previousSpec.equals(cs.spec) && !cs.connected) {
            connections.put(service.getId(), newspec);
            found = true;
            break;
          }
        }

        if (!found) {
          logger.warn("2 sinks found with the same id: {} and {}... Ignoring previous.", previousSpec, newspec);
          connections.put(service.getId(), newspec);
        }
      }
    }

    for (int i = allConnectAddresses.size(); i-- > 0;) {
      String[] parts = allConnectAddresses.get(i).split(":");
      String connection = connections.remove(parts[0]);
      if (connection == null) {
        allConnectAddresses.remove(i);
      } else {
        allConnectAddresses.set(i, connection);
      }
    }

    allConnectAddresses.addAll(connections.values());

    partitions.clear();
    try {
      if (allConnectAddresses.isEmpty()) {
        /* return at least one of them; otherwise stram becomes grumpy */
        @SuppressWarnings("unchecked")
        AbstractFlumeInputOperator<T> operator = getClass().newInstance();
        operator.setCodec(codec);
        operator.setMaxEventsPerSecond(maxEventsPerSecond);
        for (ArrayList<RecoveryAddress> lRecoveryAddresses: allRecoveryAddresses.values()) {
          operator.recoveryAddresses.addAll(lRecoveryAddresses);
        }
        operator.connectionSpecs = new String[allConnectAddresses.size()];
        for (int i = connectionSpecs.length; i-- > 0;) {
          connectionSpecs[i] = allConnectAddresses.get(i);
        }

        partitions.add(new DefaultPartition<AbstractFlumeInputOperator<T>>(operator));
      } else {
        long maxEventsPerSecondPerOperator = maxEventsPerSecond / allConnectAddresses.size();
        for (int i = allConnectAddresses.size(); i-- > 0;) {
          @SuppressWarnings("unchecked")
          AbstractFlumeInputOperator<T> operator = getClass().newInstance();
          operator.setCodec(codec);
          operator.setMaxEventsPerSecond(maxEventsPerSecondPerOperator);
          String connectAddress = allConnectAddresses.get(i);
          operator.connectionSpecs = new String[] {connectAddress};

          String[] parts = connectAddress.split(":", 2);
          ArrayList<RecoveryAddress> remove = allRecoveryAddresses.remove(parts[0]);
          if (remove != null) {
            operator.recoveryAddresses.addAll(remove);
          }

          partitions.add(new DefaultPartition<AbstractFlumeInputOperator<T>>(operator));
        }
      }
    } catch (IllegalAccessException ex) {
      throw new RuntimeException(ex);
    } catch (InstantiationException ex) {
      throw new RuntimeException(ex);
    }

    logger.debug("Requesting partitions: {}", partitions);
    return partitions;
  }

  @Override
  public void partitioned(Map<Integer, Partition<AbstractFlumeInputOperator<T>>> partitions)
  {
    logger.debug("Partitioned Map: {}", partitions);
    HashMap<Integer, ConnectionStatus> map = partitionedInstanceStatus.get();
    map.clear();
    for (Entry<Integer, Partition<AbstractFlumeInputOperator<T>>> entry: partitions.entrySet()) {
      if (map.containsKey(entry.getKey())) {
        // what can be done here?
      } else {
        map.put(entry.getKey(), null);
      }
    }
  }

  @Override
  public String toString()
  {
    return "AbstractFlumeInputOperator{" + "connected=" + connected + ", connectionSpecs=" +
        (connectionSpecs.length == 0 ? "empty" : connectionSpecs[0]) + ", recoveryAddresses=" + recoveryAddresses + '}';
  }

  class Client extends AbstractLengthPrependerClient
  {
    private final String id;

    Client(String id)
    {
      this.id = id;
    }

    @Override
    public void onMessage(byte[] buffer, int offset, int size)
    {
      try {
        handoverBuffer.put(new Slice(buffer, offset, size));
      } catch (InterruptedException ex) {
        handleException(ex, eventloop);
      }
    }

    @Override
    public void connected()
    {
      super.connected();

      byte[] address;
      synchronized (recoveryAddresses) {
        if (recoveryAddresses.size() > 0) {
          address = recoveryAddresses.get(recoveryAddresses.size() - 1).address;
        } else {
          address = new byte[8];
        }
      }

      int len = 1 /* for the message type SEEK */
          + 8 /* for the address */
          + 8 /* for storing the current time stamp*/;

      byte[] array = new byte[len];
      array[0] = Server.Command.SEEK.getOrdinal();
      System.arraycopy(address, 0, array, 1, 8);
      Server.writeLong(array, 9, System.currentTimeMillis());
      write(array);

      connected = true;
      ConnectionStatus connectionStatus = new ConnectionStatus();
      connectionStatus.connected = true;
      connectionStatus.spec = connectionSpecs[0];
      OperatorContext ctx = context;
      synchronized (ctx) {
        logger.debug("{} Submitting ConnectionStatus = {}", AbstractFlumeInputOperator.this, connectionStatus);
        context.setCounters(connectionStatus);
      }
    }

    @Override
    public void disconnected()
    {
      connected = false;
      ConnectionStatus connectionStatus = new ConnectionStatus();
      connectionStatus.connected = false;
      connectionStatus.spec = connectionSpecs[0];
      OperatorContext ctx = context;
      synchronized (ctx) {
        logger.debug("{} Submitting ConnectionStatus = {}", AbstractFlumeInputOperator.this, connectionStatus);
        context.setCounters(connectionStatus);
      }
      super.disconnected();
    }

  }

  public static class ZKStatsListner extends ZKAssistedDiscovery implements com.datatorrent.api.StatsListener,
      Serializable
  {
    /*
     * In the current design, one input operator is able to connect
     * to only one flume adapter. Sometime in future, we should support
     * any number of input operators connecting to any number of flume
     * sinks and vice a versa.
     *
     * Until that happens the following map should be sufficient to
     * keep track of which input operator is connected to which flume sink.
     */
    long intervalMillis;
    private final Response response;
    private transient long nextMillis;

    public ZKStatsListner()
    {
      intervalMillis = 60 * 1000L;
      response = new Response();
    }

    @Override
    public Response processStats(BatchedOperatorStats stats)
    {
      final HashMap<Integer, ConnectionStatus> map = partitionedInstanceStatus.get();
      response.repartitionRequired = false;

      Object lastStat = null;
      List<OperatorStats> lastWindowedStats = stats.getLastWindowedStats();
      for (OperatorStats os: lastWindowedStats) {
        if (os.counters != null) {
          lastStat = os.counters;
          logger.debug("Received custom stats = {}", lastStat);
        }
      }

      if (lastStat instanceof ConnectionStatus) {
        ConnectionStatus cs = (ConnectionStatus)lastStat;
        map.put(stats.getOperatorId(), cs);
        if (!cs.connected) {
          logger.debug("setting repatitioned = true because of lastStat = {}", lastStat);
          response.repartitionRequired = true;
        }
      }

      if (System.currentTimeMillis() >= nextMillis) {
        logger.debug("nextMillis = {}", nextMillis);
        try {
          super.setup(null);
          Collection<Discovery.Service<byte[]>> addresses;
          try {
            addresses = discover();
          } finally {
            super.teardown();
          }
          AbstractFlumeInputOperator.discoveredFlumeSinks.set(addresses);
          logger.debug("\ncurrent map = {}\ndiscovered sinks = {}", map, addresses);
          switch (addresses.size()) {
            case 0:
              response.repartitionRequired = map.size() != 1;
              break;

            default:
              if (addresses.size() == map.size()) {
                for (ConnectionStatus value: map.values()) {
                  if (value == null || !value.connected) {
                    response.repartitionRequired = true;
                    break;
                  }
                }
              } else {
                response.repartitionRequired = true;
              }
              break;
          }
        } catch (Error er) {
          throw er;
        } catch (Throwable cause) {
          logger.warn("Unable to discover services, using values from last successful discovery", cause);
        } finally {
          nextMillis = System.currentTimeMillis() + intervalMillis;
          logger.debug("Proposed NextMillis = {}", nextMillis);
        }
      }

      return response;
    }

    /**
     * @return the intervalMillis
     */
    public long getIntervalMillis()
    {
      return intervalMillis;
    }

    /**
     * @param intervalMillis the intervalMillis to set
     */
    public void setIntervalMillis(long intervalMillis)
    {
      this.intervalMillis = intervalMillis;
    }

    private static final long serialVersionUID = 201312241646L;
  }

  public static class ConnectionStatus implements Serializable
  {
    int id;
    String spec;
    boolean connected;

    @Override
    public int hashCode()
    {
      return spec.hashCode();
    }

    @Override
    public boolean equals(Object obj)
    {
      if (obj == null) {
        return false;
      }
      if (getClass() != obj.getClass()) {
        return false;
      }
      final ConnectionStatus other = (ConnectionStatus)obj;
      return spec == null ? other.spec == null : spec.equals(other.spec);
    }

    @Override
    public String toString()
    {
      return "ConnectionStatus{" + "id=" + id + ", spec=" + spec + ", connected=" + connected + '}';
    }

    private static final long serialVersionUID = 201312261615L;
  }

  private static final transient ThreadLocal<HashMap<Integer, ConnectionStatus>> partitionedInstanceStatus =
      new ThreadLocal<HashMap<Integer, ConnectionStatus>>()
    {
      @Override
      protected HashMap<Integer, ConnectionStatus> initialValue()
      {
        return new HashMap<Integer, ConnectionStatus>();
      }

    };
  /**
   * When a sink goes away and a replacement sink is not found, we stash the recovery addresses associated
   * with the sink in a hope that the new sink may show up in near future.
   */
  private static final transient ThreadLocal<HashMap<String, ArrayList<RecoveryAddress>>> abandonedRecoveryAddresses =
      new ThreadLocal<HashMap<String, ArrayList<RecoveryAddress>>>()
  {
    @Override
    protected HashMap<String, ArrayList<RecoveryAddress>> initialValue()
    {
      return new HashMap<String, ArrayList<RecoveryAddress>>();
    }

  };
  protected static final transient ThreadLocal<Collection<Discovery.Service<byte[]>>> discoveredFlumeSinks =
      new ThreadLocal<Collection<Discovery.Service<byte[]>>>();

  @Override
  public boolean equals(Object o)
  {
    if (this == o) {
      return true;
    }
    if (!(o instanceof AbstractFlumeInputOperator)) {
      return false;
    }

    AbstractFlumeInputOperator<?> that = (AbstractFlumeInputOperator<?>)o;

    if (!Arrays.equals(connectionSpecs, that.connectionSpecs)) {
      return false;
    }
    return recoveryAddresses.equals(that.recoveryAddresses);

  }

  @Override
  public int hashCode()
  {
    int result = connectionSpecs != null ? Arrays.hashCode(connectionSpecs) : 0;
    result = 31 * result + (recoveryAddresses.hashCode());
    return result;
  }

  public void setMaxEventsPerSecond(long maxEventsPerSecond)
  {
    this.maxEventsPerSecond = maxEventsPerSecond;
  }

  public long getMaxEventsPerSecond()
  {
    return maxEventsPerSecond;
  }

  private static final Logger logger = LoggerFactory.getLogger(AbstractFlumeInputOperator.class);
}