// 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.

#pragma once

#include <memory>
#include <random>
#include <sstream>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

#include "arrow/array.h"
#include "arrow/array/builder_binary.h"
#include "arrow/io/memory.h"
#include "arrow/json/converter.h"
#include "arrow/json/options.h"
#include "arrow/json/parser.h"
#include "arrow/json/rapidjson_defs.h"
#include "arrow/testing/gtest_util.h"
#include "arrow/type.h"
#include "arrow/util/checked_cast.h"
#include "arrow/visit_type_inline.h"

#include "rapidjson/document.h"
#include "rapidjson/prettywriter.h"
#include "rapidjson/reader.h"
#include "rapidjson/writer.h"

namespace arrow {

using internal::checked_cast;

namespace json {

namespace rj = arrow::rapidjson;

using rj::StringBuffer;
using std::string_view;
using Writer = rj::Writer<StringBuffer>;

struct GenerateOptions {
  // Probability of a field being written
  double field_probability = 1.0;
  // Probability of a value being null
  double null_probability = 0.2;
  // Whether to randomize the order of written fields
  bool randomize_field_order = false;

  static constexpr GenerateOptions Defaults() { return GenerateOptions{}; }
};

inline static Status OK(bool ok) { return ok ? Status::OK() : Status::Invalid(""); }

template <typename Engine>
inline static Status Generate(
    const std::shared_ptr<DataType>& type, Engine& e, Writer* writer,
    const GenerateOptions& options = GenerateOptions::Defaults());

template <typename Engine>
inline static Status Generate(
    const std::vector<std::shared_ptr<Field>>& fields, Engine& e, Writer* writer,
    const GenerateOptions& options = GenerateOptions::Defaults());

template <typename Engine>
inline static Status Generate(
    const std::shared_ptr<Schema>& schm, Engine& e, Writer* writer,
    const GenerateOptions& options = GenerateOptions::Defaults()) {
  return Generate(schm->fields(), e, writer, options);
}

template <typename Engine>
struct GenerateImpl {
  Status Visit(const NullType&) { return OK(writer.Null()); }

  Status Visit(const BooleanType&) {
    return OK(writer.Bool(std::uniform_int_distribution<uint16_t>{}(e)&1));
  }

  template <typename T>
  enable_if_physical_unsigned_integer<T, Status> Visit(const T&) {
    auto val = std::uniform_int_distribution<>{}(e);
    return OK(writer.Uint64(static_cast<typename T::c_type>(val)));
  }

  template <typename T>
  enable_if_physical_signed_integer<T, Status> Visit(const T&) {
    auto val = std::uniform_int_distribution<>{}(e);
    return OK(writer.Int64(static_cast<typename T::c_type>(val)));
  }

  template <typename T>
  enable_if_physical_floating_point<T, Status> Visit(const T&) {
    auto val = std::normal_distribution<typename T::c_type>{0, 1 << 10}(e);
    return OK(writer.Double(val));
  }

  Status GenerateAscii(const DataType&) {
    auto size = std::poisson_distribution<>{4}(e);
    std::uniform_int_distribution<uint16_t> gen_char(32, 126);  // FIXME generate UTF8
    std::string s(size, '\0');
    for (char& ch : s) ch = static_cast<char>(gen_char(e));
    return OK(writer.String(s.c_str()));
  }

  template <typename T>
  enable_if_base_binary<T, Status> Visit(const T& t) {
    return GenerateAscii(t);
  }

  Status Visit(const BinaryViewType& t) { return GenerateAscii(t); }

  template <typename T>
  enable_if_list_like<T, Status> Visit(const T& t) {
    auto size = std::poisson_distribution<>{4}(e);
    writer.StartArray();
    for (int i = 0; i < size; ++i) {
      RETURN_NOT_OK(Generate(t.value_type(), e, &writer, options));
    }
    return OK(writer.EndArray(size));
  }

  Status Visit(const ListViewType& t) { return NotImplemented(t); }

  Status Visit(const LargeListViewType& t) { return NotImplemented(t); }

  Status Visit(const StructType& t) { return Generate(t.fields(), e, &writer, options); }

  Status Visit(const DayTimeIntervalType& t) { return NotImplemented(t); }

  Status Visit(const MonthDayNanoIntervalType& t) { return NotImplemented(t); }

  Status Visit(const DictionaryType& t) { return NotImplemented(t); }

  Status Visit(const ExtensionType& t) { return NotImplemented(t); }

  Status Visit(const Decimal128Type& t) { return NotImplemented(t); }

  Status Visit(const FixedSizeBinaryType& t) { return NotImplemented(t); }

  Status Visit(const UnionType& t) { return NotImplemented(t); }

  Status Visit(const RunEndEncodedType& t) { return NotImplemented(t); }

  Status NotImplemented(const DataType& t) {
    return Status::NotImplemented("random generation of arrays of type ", t);
  }

  Engine& e;
  rj::Writer<rj::StringBuffer>& writer;
  const GenerateOptions& options;
};

template <typename Engine>
inline static Status Generate(const std::shared_ptr<DataType>& type, Engine& e,
                              Writer* writer, const GenerateOptions& options) {
  if (std::bernoulli_distribution(options.null_probability)(e)) {
    writer->Null();
    return Status::OK();
  }
  GenerateImpl<Engine> visitor = {e, *writer, options};
  return VisitTypeInline(*type, &visitor);
}

template <typename Engine>
inline static Status Generate(const std::vector<std::shared_ptr<Field>>& fields,
                              Engine& e, Writer* writer, const GenerateOptions& options) {
  RETURN_NOT_OK(OK(writer->StartObject()));

  int num_fields = 0;
  auto write_field = [&](const Field& f) {
    ++num_fields;
    writer->Key(f.name().c_str());
    return Generate(f.type(), e, writer, options);
  };

  std::bernoulli_distribution bool_dist(options.field_probability);
  if (options.randomize_field_order) {
    std::vector<size_t> indices;
    indices.reserve(static_cast<size_t>(fields.size() * options.field_probability));
    for (size_t i = 0; i < fields.size(); ++i) {
      if (bool_dist(e)) {
        indices.push_back(i);
      }
    }
    std::shuffle(indices.begin(), indices.end(), e);
    for (auto i : indices) {
      RETURN_NOT_OK(write_field(*fields[i]));
    }
  } else {
    for (const auto& f : fields) {
      if (bool_dist(e)) {
        RETURN_NOT_OK(write_field(*f));
      }
    }
  }

  return OK(writer->EndObject(num_fields));
}

inline static Status MakeStream(string_view src_str,
                                std::shared_ptr<io::InputStream>* out) {
  auto src = std::make_shared<Buffer>(src_str);
  *out = std::make_shared<io::BufferReader>(src);
  return Status::OK();
}

// scalar values (numbers and strings) are parsed into a
// dictionary<index:int32, value:string>. This can be decoded for ease of comparison
inline static Status DecodeStringDictionary(const DictionaryArray& dict_array,
                                            std::shared_ptr<Array>* decoded) {
  const StringArray& dict = checked_cast<const StringArray&>(*dict_array.dictionary());
  const Int32Array& indices = checked_cast<const Int32Array&>(*dict_array.indices());
  StringBuilder builder;
  RETURN_NOT_OK(builder.Resize(indices.length()));
  for (int64_t i = 0; i < indices.length(); ++i) {
    if (indices.IsNull(i)) {
      builder.UnsafeAppendNull();
      continue;
    }
    auto value = dict.GetView(indices.GetView(i));
    RETURN_NOT_OK(builder.ReserveData(value.size()));
    builder.UnsafeAppend(value);
  }
  return builder.Finish(decoded);
}

inline static Status ParseFromString(ParseOptions options, string_view src_str,
                                     std::shared_ptr<Array>* parsed) {
  auto src = std::make_shared<Buffer>(src_str);
  std::unique_ptr<BlockParser> parser;
  RETURN_NOT_OK(BlockParser::Make(options, &parser));
  RETURN_NOT_OK(parser->Parse(src));
  return parser->Finish(parsed);
}

inline static Status ParseFromString(ParseOptions options, string_view src_str,
                                     std::shared_ptr<StructArray>* parsed) {
  std::shared_ptr<Array> parsed_non_struct;
  RETURN_NOT_OK(ParseFromString(options, src_str, &parsed_non_struct));
  *parsed = internal::checked_pointer_cast<StructArray>(parsed_non_struct);
  return Status::OK();
}

static inline std::string PrettyPrint(string_view one_line) {
  rj::Document document;

  // Must pass size to avoid ASAN issues.
  document.Parse(one_line.data(), one_line.size());
  rj::StringBuffer sb;
  rj::PrettyWriter<rj::StringBuffer> writer(sb);
  document.Accept(writer);
  return sb.GetString();
}

template <typename T>
std::string RowsOfOneColumn(std::string_view name, std::initializer_list<T> values,
                            decltype(std::to_string(*values.begin()))* = nullptr) {
  std::stringstream ss;
  for (auto value : values) {
    ss << R"({")" << name << R"(":)" << std::to_string(value) << "}\n";
  }
  return ss.str();
}

inline std::string RowsOfOneColumn(std::string_view name,
                                   std::initializer_list<std::string> values) {
  std::stringstream ss;
  for (auto value : values) {
    ss << R"({")" << name << R"(":)" << value << "}\n";
  }
  return ss.str();
}

inline static std::string scalars_only_src() {
  return R"(
    { "hello": 3.5, "world": false, "yo": "thing" }
    { "hello": 3.25, "world": null }
    { "hello": 3.125, "world": null, "yo": "\u5fcd" }
    { "hello": 0.0, "world": true, "yo": null }
  )";
}

inline static std::string nested_src() {
  return R"(
    { "hello": 3.5, "world": false, "yo": "thing", "arr": [1, 2, 3], "nuf": {} }
    { "hello": 3.25, "world": null, "arr": [2], "nuf": null }
    { "hello": 3.125, "world": null, "yo": "\u5fcd", "arr": [], "nuf": { "ps": 78 } }
    { "hello": 0.0, "world": true, "yo": null, "arr": null, "nuf": { "ps": 90 } }
  )";
}

inline static std::string null_src() {
  return R"(
    { "plain": null, "list1": [], "list2": [], "struct": { "plain": null } }
    { "plain": null, "list1": [], "list2": [null], "struct": {} }
  )";
}

inline static std::string unquoted_decimal_src() {
  return R"(
    { "price": 30.04, "cost":30.001 }
    { "price": 1.23, "cost":1.229 }
  )";
}

inline static std::string mixed_decimal_src() {
  return R"(
    { "price": 30.04, "cost": 30.001 }
    { "price": "1.23", "cost": "1.229" }
  )";
}

}  // namespace json
}  // namespace arrow