package water.parser.avro;

import org.apache.avro.Schema;
import org.apache.avro.file.DataFileReader;
import org.apache.avro.file.DataFileStream;
import org.apache.avro.file.SeekableByteArrayInput;
import org.apache.avro.file.SeekableInput;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.DatumReader;
import org.apache.avro.util.Utf8;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.List;

import water.H2O;
import water.Job;
import water.Key;

import water.fvec.Vec;
import water.parser.BufferedString;
import water.parser.ParseReader;
import water.parser.ParseSetup;
import water.parser.ParseWriter;
import water.parser.Parser;

import water.util.ArrayUtils;
import water.util.Log;

import static water.parser.avro.AvroUtil.*;

/**
 * AVRO parser for H2O distributed parsing subsystem.
 */
public class AvroParser extends Parser {

  /** Avro header */
  private final byte[] header;

  AvroParser(ParseSetup setup, Key<Job> jobKey) {
    super(setup, jobKey);
    this.header = ((AvroParser.AvroParseSetup) setup).header;
  }

  @Override
  protected final ParseWriter parseChunk(int cidx, ParseReader din, ParseWriter dout) {
    // We will read GenericRecord and load them based on schema
    final DatumReader<GenericRecord> datumReader = new GenericDatumReader<>();
    final H2OSeekableInputAdaptor sbai = new H2OSeekableInputAdaptor(cidx, din);
    DataFileReader<GenericRecord> dataFileReader = null;
    int cnt = 0;
    try {
      // Reconstruct Avro header
      DataFileStream.Header
          fakeHeader = new DataFileReader<>(new SeekableByteArrayInput(this.header), datumReader).getHeader();
      dataFileReader = DataFileReader.openReader(sbai, datumReader, fakeHeader, true);
      Schema schema = dataFileReader.getSchema();
      GenericRecord gr = new GenericData.Record(schema);
      Schema.Field[] flatSchema = flatSchema(schema);
      long sync = dataFileReader.previousSync();
      if (sbai.chunkCnt == 0) { // Find data in first chunk
        while (dataFileReader.hasNext() && dataFileReader.previousSync() == sync) {
          gr = dataFileReader.next(gr);
          // Write values to the output
          // FIXME: what if user change input names, or ignore an input column?
          write2frame(gr, _setup.getColumnNames(), flatSchema, _setup.getColumnTypes(), dout);
          cnt++;
        }
      } // else first chunk does not contain synchronization block, so give up and let another reader to use it
    } catch (IOException e) {
      throw new IllegalStateException("Failed to read AVRO.", e);
    }

    Log.trace(String.format("Avro: ChunkIdx: %d read %d records, start at %d off, block count: %d, block size: %d", cidx, cnt, din.getChunkDataStart(cidx), dataFileReader.getBlockCount(), dataFileReader.getBlockSize()));

    return dout;
  }


  /** A simple adaptor for Avro Seekable Input.
   *
   * It implements lazy loading of chunks from ParseReader and track how many chunks
   * were loaded.
   *
   * Warning: This is not designed to be accessed by multiple threads!
   */
  private static class H2OSeekableInputAdaptor implements SeekableInput {

    private final ParseReader din;
    private final int startCidx;

    protected int pos;
    protected int mark;

    private byte[] data;
    // Additional chunks loaded
    protected int chunkCnt;

    public H2OSeekableInputAdaptor(int cidx, ParseReader din) {
      this.din = din;
      this.startCidx = cidx;
      this.data = din.getChunkData(cidx);
      this.chunkCnt = 0;

      this.mark = din.getChunkDataStart(cidx) > 0 ? din.getChunkDataStart(cidx) : 0;
      this.pos = mark;
    }

    @Override
    public void seek(long p) throws IOException {
      this.reset();
      this.skip(p);
    }

    @Override
    public long tell() throws IOException {
      return this.pos;
    }

    @Override
    public long length() throws IOException {
      return -1;
    }

    @Override
    public int read(byte[] b, int off, int len) throws IOException {
      if (b == null) {
        throw new NullPointerException();
      } else if (off < 0 || len < 0 || len > b.length - off) {
        throw new IndexOutOfBoundsException();
      }

      needData(len);

      if (pos >= count()) {
        return -1;
      }

      int avail = count() - pos;
      if (len > avail) {
        len = avail;
      }
      if (len <= 0) {
        return 0;
      }

      // FIXME drop read data
      System.arraycopy(data, pos, b, off, len);
      pos += len;
      return len;
    }

    @Override
    public void close() throws IOException {
      data = null;
    }

    public void reset() {
      pos = 0;
    }

    public long skip(long n) {
      long remain = 0;
      while ((remain = count() - pos) < n && loadNextData()) ;
      if (n < remain) {
        remain = n < 0 ? 0 : n;
      }
      pos += remain;
      return remain;
    }

    private int count() {
      return data.length;
    }

    private boolean needData(int len) {
      boolean loaded = false;
      while ((count() - pos) < len && (loaded = loadNextData())) ;
      return loaded;
    }

    private boolean loadNextData() {
      // FIXME: just replace data
      byte[] nextChunk = this.din.getChunkData(this.startCidx + chunkCnt + 1);
      if (nextChunk != null && nextChunk.length > 0) {
        this.data = ArrayUtils.append(this.data, nextChunk);
        this.chunkCnt++;
        Log.trace(String.format("Avro stream wrapper - loading another chunk: StartChunkIdx: %d, LoadedChunkCnt: %d",  startCidx, chunkCnt));
        return true;
      } else {
        return false;
      }
    }
  }

  /**
   * The main method transforming Avro record into a row in H2O frame.
   *
   * @param gr  Avro generic record
   * @param columnNames Column names prepared by parser setup
   * @param inSchema  Flattenized Avro schema which corresponds to passed column names
   * @param columnTypes  Target H2O types
   * @param dout  Parser writer
   */
  private static void write2frame(GenericRecord gr, String[] columnNames, Schema.Field[] inSchema, byte[] columnTypes, ParseWriter dout) {
    assert inSchema.length == columnTypes.length : "AVRO field flatenized schema has to match to parser setup";
    BufferedString bs = new BufferedString();
    for (int cIdx = 0; cIdx < columnNames.length; cIdx++) {
      int inputFieldIdx = inSchema[cIdx].pos();
      Schema.Type inputType = toPrimitiveType(inSchema[cIdx].schema());
      byte targetType = columnTypes[cIdx]; // FIXME: support target conversions
      Object value = gr.get(inputFieldIdx);
      if (value == null) {
        dout.addInvalidCol(cIdx);
      } else {
        switch (inputType) {
          case BOOLEAN:
            dout.addNumCol(cIdx, ((Boolean) value) ? 1 : 0);
            break;
          case INT:
            dout.addNumCol(cIdx, ((Integer) value), 0);
            break;
          case LONG:
            dout.addNumCol(cIdx, ((Long) value), 0);
            break;
          case FLOAT:
            dout.addNumCol(cIdx, (Float) value);
            break;
          case DOUBLE:
            dout.addNumCol(cIdx, (Double) value);
            break;
          case ENUM:
            // Note: this code expects ordering of categoricals provided by Avro remain same
            // as in H2O!!!
            GenericData.EnumSymbol es = (GenericData.EnumSymbol) value;
            dout.addNumCol(cIdx, es.getSchema().getEnumOrdinal(es.toString()));
            break;
          case BYTES:
            dout.addStrCol(cIdx, bs.set(((ByteBuffer) value).array()));
            break;
          case STRING:
            dout.addStrCol(cIdx, bs.set(((Utf8) value).getBytes()));
            break;
          case NULL:
            dout.addInvalidCol(cIdx);
            break;
        }
      }
    }
  }

  public static class AvroParseSetup extends ParseSetup {
    final byte[] header;
    final long blockSize;

    public AvroParseSetup(int ncols,
                          String[] columnNames,
                          byte[] ctypes,
                          String[][] domains,
                          String[][] naStrings,
                          String[][] data,
                          byte[] header,
                          long blockSize) {
      super(AvroParserProvider.AVRO_INFO, (byte) '|', true, HAS_HEADER , ncols, columnNames, ctypes, domains, naStrings, data);
      this.header = header;
      this.blockSize = blockSize;
      this.setChunkSize((int) blockSize);
    }

    public AvroParseSetup(ParseSetup ps, byte[] header, long blockSize, String[][] domains) {
      super(ps);
      this.header = header;
      this.blockSize = blockSize;
      this.setDomains(domains);
      this.setChunkSize((int) blockSize);
    }

    @Override
    protected Parser parser(Key jobKey) {
      return new AvroParser(this, jobKey);
    }
  }

  public static ParseSetup guessSetup(byte[] bits) {
    try {
      return runOnPreview(bits, new AvroPreviewProcessor<ParseSetup>() {
        @Override
        public ParseSetup process(byte[] header, GenericRecord gr, long blockCount,
                                  long blockSize) {
          return deriveParseSetup(header, gr, blockCount, blockSize);
        }
      });
    } catch (IOException e) {
      throw new RuntimeException("Avro format was not recognized", e);
    }
  }

  static AvroInfo extractAvroInfo(byte[] bits, final ParseSetup requiredSetup) throws IOException {
    return runOnPreview(bits, new AvroPreviewProcessor<AvroInfo>() {
      @Override
      public AvroInfo process(byte[] header, GenericRecord gr, long blockCount,
                              long blockSize) {
        Schema recordSchema = gr.getSchema();
        List<Schema.Field> fields = recordSchema.getFields();
        int supportedFieldCnt = 0 ;
        for (Schema.Field f : fields) if (isSupportedSchema(f.schema())) supportedFieldCnt++;
        assert supportedFieldCnt == requiredSetup.getColumnNames().length : "User-driven changes are not not supported in Avro format";
        String[][] domains = new String[supportedFieldCnt][];
        int i = 0;
        for (Schema.Field f : fields) {
          Schema schema = f.schema();
          if (isSupportedSchema(schema)) {
            byte type = schemaToColumnType(schema);
            if (type == Vec.T_CAT) {
              domains[i] = getDomain(schema);
            }
            i++;
          }
        }
        return new AvroInfo(header, blockCount, blockSize, domains);
      }
    });
  }

  static <T> T runOnPreview(byte[] bits, AvroPreviewProcessor<T> processor) throws IOException {
    DatumReader<GenericRecord> datumReader = new GenericDatumReader<GenericRecord>();
    SeekableByteArrayInput sbai = new SeekableByteArrayInput(bits);
    DataFileReader<GenericRecord> dataFileReader = null;

    try {
      dataFileReader = new DataFileReader<>(sbai, datumReader);
      int headerLen = (int) dataFileReader.previousSync();
      byte[] header = Arrays.copyOf(bits, headerLen);
      if (dataFileReader.hasNext()) {
        GenericRecord gr = dataFileReader.next();
        return processor.process(header, gr, dataFileReader.getBlockCount(), dataFileReader.getBlockSize());
      } else {
        throw new RuntimeException("Empty Avro file - cannot run preview! ");
      }
    } finally {
      try { if (dataFileReader!=null) dataFileReader.close(); } catch (IOException safeToIgnore) {}
    }
  }

  private static ParseSetup deriveParseSetup(byte[] header, GenericRecord gr,
                                             long blockCount, long blockSize) {
    // Expect flat structure
    Schema recordSchema = gr.getSchema();
    List<Schema.Field> fields = recordSchema.getFields();
    int supportedFieldCnt = 0 ;
    for (Schema.Field f : fields) if (isSupportedSchema(f.schema())) supportedFieldCnt++;
    String[] names = new String[supportedFieldCnt];
    byte[] types = new byte[supportedFieldCnt];
    String[][] domains = new String[supportedFieldCnt][];
    String[] dataPreview = new String[supportedFieldCnt];
    int i = 0;
    for (Schema.Field f : fields) {
      Schema schema = f.schema();
      if (isSupportedSchema(schema)) {
        names[i] = f.name();
        types[i] = schemaToColumnType(schema);
        if (types[i] == Vec.T_CAT) {
          domains[i] = getDomain(schema);
        }
        dataPreview[i] = gr.get(f.name()) != null ? gr.get(f.name()).toString() : "null";
        i++;
      } else {
        Log.warn("Skipping field: " + f.name() + " because of unsupported type: " + schema.getType() + " schema: " + schema);
      }
    }

    AvroParseSetup ps = new AvroParseSetup(
        supportedFieldCnt,
        names,
        types,
        domains,
        null,
        new String[][] { dataPreview },
        header,
        blockSize
    );
    return ps;
  }

  /** Helper to represent Avro header
   * and size of 1st block of data.
   */
  static class AvroInfo {

    public AvroInfo(byte[] header, long firstBlockCount, long firstBlockSize, String[][] domains) {
      this.header = header;
      this.firstBlockCount = firstBlockCount;
      this.firstBlockSize = firstBlockSize;
      this.domains = domains;
    }

    byte[] header;
    long firstBlockCount;
    long firstBlockSize;
    String[][] domains;
  }

  private interface AvroPreviewProcessor<R> {
    R process(byte[] header, GenericRecord gr, long blockCount, long blockSize);
  }

}
