/*
 * 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
 *
 *     https://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.avro.util;

import org.apache.avro.LogicalType;
import org.apache.avro.LogicalTypes;
import org.apache.avro.Schema;
import org.apache.avro.SchemaParser;
import org.apache.avro.file.CodecFactory;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.generic.GenericArray;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericDatumWriter;

import java.io.File;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;

/** Generates schema data as Java objects with random values. */
public class RandomData implements Iterable<Object> {
  public static final String USE_DEFAULT = "use-default";
  private final GenericData genericData;

  private static final int MILLIS_IN_DAY = (int) Duration.ofDays(1).toMillis();

  private final Schema root;
  private final long seed;
  private final int count;
  private final boolean utf8ForString;

  public RandomData(Schema schema, int count) {
    this(schema, count, false);
  }

  public RandomData(Schema schema, int count, long seed) {
    this(schema, count, seed, false);
  }

  public RandomData(Schema schema, int count, boolean utf8ForString) {
    this(schema, count, System.currentTimeMillis(), utf8ForString);
  }

  public RandomData(Schema schema, int count, long seed, boolean utf8ForString) {
    this(GenericData.get(), schema, count, seed, utf8ForString);
  }

  public RandomData(GenericData genericData, Schema schema, int count) {
    this(genericData, schema, count, false);
  }

  public RandomData(GenericData genericData, Schema schema, int count, long seed) {
    this(genericData, schema, count, seed, false);
  }

  public RandomData(GenericData genericData, Schema schema, int count, boolean utf8ForString) {
    this(genericData, schema, count, System.currentTimeMillis(), utf8ForString);
  }

  public RandomData(GenericData genericData, Schema schema, int count, long seed, boolean utf8ForString) {
    this.genericData = genericData;
    this.root = schema;
    this.seed = seed;
    this.count = count;
    this.utf8ForString = utf8ForString;
  }

  @Override
  public Iterator<Object> iterator() {
    return new Iterator<Object>() {
      private int n;
      private final Random random = new Random(seed);

      @Override
      public boolean hasNext() {
        return n < count;
      }

      @Override
      public Object next() {
        n++;
        return generate(root, random, 0);
      }

      @Override
      public void remove() {
        throw new UnsupportedOperationException();
      }
    };
  }

  @SuppressWarnings(value = "unchecked")
  private Object generate(Schema schema, Random random, int d) {
    switch (schema.getType()) {
    case RECORD:
      Object record = genericData.newRecord(null, schema);
      for (Schema.Field field : schema.getFields()) {
        Object value = (field.getObjectProp(USE_DEFAULT) == null) ? generate(field.schema(), random, d + 1)
            : GenericData.get().getDefaultValue(field);
        genericData.setField(record, field.name(), field.pos(), value);
      }
      return record;
    case ENUM:
      List<String> symbols = schema.getEnumSymbols();
      return genericData.createEnum(symbols.get(random.nextInt(symbols.size())), schema);
    case ARRAY:
      int length = Math.max(0, (random.nextInt(5) + 2) - d);
      GenericArray<Object> array = (GenericArray<Object>) genericData.newArray(null, length, schema);
      for (int i = 0; i < length; i++)
        array.add(generate(schema.getElementType(), random, d + 1));
      return array;
    case MAP:
      length = Math.max(0, (random.nextInt(5) + 2) - d);
      Map<Object, Object> map = (Map<Object, Object>) genericData.newMap(null, length);
      for (int i = 0; i < length; i++) {
        map.put(randomString(random, 40), generate(schema.getValueType(), random, d + 1));
      }
      return map;
    case UNION:
      List<Schema> types = schema.getTypes();
      return generate(types.get(random.nextInt(types.size())), random, d);
    case FIXED:
      byte[] bytes = new byte[schema.getFixedSize()];
      random.nextBytes(bytes);
      return genericData.createFixed(null, bytes, schema);
    case STRING:
      return randomString(random, 40);
    case BYTES:
      return randomBytes(random, 40);
    case INT:
      return this.randomInt(random, schema.getLogicalType());
    case LONG:
      return this.randomLong(random, schema.getLogicalType());
    case FLOAT:
      return random.nextFloat();
    case DOUBLE:
      return random.nextDouble();
    case BOOLEAN:
      return random.nextBoolean();
    case NULL:
      return null;
    default:
      throw new RuntimeException("Unknown type: " + schema);
    }
  }

  private static final Charset UTF8 = StandardCharsets.UTF_8;

  private int randomInt(Random random, LogicalType type) {
    if (type instanceof LogicalTypes.TimeMillis) {
      return random.nextInt(RandomData.MILLIS_IN_DAY - 1);
    }
    // LogicalTypes.Date LocalDate.MAX.toEpochDay() > Integer.MAX;
    return random.nextInt();
  }

  private long randomLong(Random random, LogicalType type) {
    if (type instanceof LogicalTypes.TimeMicros) {
      return ThreadLocalRandom.current().nextLong(RandomData.MILLIS_IN_DAY * 1000L);
    }
    // For LogicalTypes.TimestampMillis, every long would be OK,
    // Instant.MAX.toEpochMilli() failed and would be > Long.MAX_VALUE.
    return random.nextLong();
  }

  private Object randomString(Random random, int maxLength) {
    int length = random.nextInt(maxLength);
    byte[] bytes = new byte[length];
    for (int i = 0; i < length; i++) {
      bytes[i] = (byte) ('a' + random.nextInt('z' - 'a'));
    }
    return utf8ForString ? new Utf8(bytes) : new String(bytes, UTF8);
  }

  private static ByteBuffer randomBytes(Random rand, int maxLength) {
    ByteBuffer bytes = ByteBuffer.allocate(rand.nextInt(maxLength));
    bytes.limit(bytes.capacity());
    rand.nextBytes(bytes.array());
    return bytes;
  }

  public static void main(String[] args) throws Exception {
    if (args.length < 3 || args.length > 4) {
      System.out.println("Usage: RandomData <schemafile> <outputfile> <count> [codec]");
      System.exit(-1);
    }
    Schema sch = new SchemaParser().parse(new File(args[0])).mainSchema();
    try (DataFileWriter<Object> writer = new DataFileWriter<>(new GenericDatumWriter<>())) {
      writer.setCodec(CodecFactory.fromString(args.length >= 4 ? args[3] : "null"));
      writer.setMeta("user_metadata", "someByteArray".getBytes(StandardCharsets.UTF_8));
      File file = new File(args[1]);
      Files.createDirectories(Paths.get(file.getParent()));
      writer.create(sch, file);

      for (Object datum : new RandomData(sch, Integer.parseInt(args[2]))) {
        writer.append(datum);
      }
    }
  }
}
