// 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 <rapidjson/allocators.h>

#include <rapidjson/document.h>

#include <rapidjson/encodings.h>

#include <rapidjson/rapidjson.h>

#include <simdjson/common_defs.h>

#include <simdjson/simdjson.h> // IWYU pragma: keep

#include <memory>

#include <string>

#include <string_view>

#include <unordered_map>

#include <unordered_set>

#include <vector>

#include "common/status.h"

#include "core/custom_allocator.h"

#include "core/string_ref.h"

#include "core/types.h"

#include "exprs/json_functions.h"

#include "format/line_reader.h"

#include "format/table/table_format_reader.h"

#include "io/file_factory.h"

#include "io/fs/file_reader_writer_fwd.h"

#include "runtime/runtime_profile.h"

#include "util/decompressor.h"

namespace simdjson::fallback::ondemand {

class object;

} // namespace simdjson::fallback::ondemand

namespace doris {

class SlotDescriptor;

class RuntimeState;

class TFileRangeDesc;

class TFileScanRangeParams;

namespace io {

class FileSystem;

struct IOContext;

} // namespace io

struct ScannerCounter;

class Block;

class IColumn;

class NewPlainTextLineReader;

namespace json_reader_detail {

Status append_null_for_malformed_json(Block& block);

void truncate_block_to_rows(Block& block, size_t num_rows);

void pop_back_last_inserted_value(Block& block, size_t column_index);

} // namespace json_reader_detail

/// JSON-specific initialization context.

/// Extends ReaderInitContext with default value context (unique to JSON reader).

struct JsonInitContext final : public ReaderInitContext {

    const std::unordered_map<std::string, VExprContextSPtr>* col_default_value_ctx = nullptr;

    bool is_load = false;

};

class NewJsonReader : public TableFormatReader {

    ENABLE_FACTORY_CREATOR(NewJsonReader);

public:

    NewJsonReader(RuntimeState* state, RuntimeProfile* profile, ScannerCounter* counter,

                  const TFileScanRangeParams& params, const TFileRangeDesc& range,

                  const std::vector<SlotDescriptor*>& file_slot_descs, bool* scanner_eof,

                  size_t batch_size, io::IOContext* io_ctx,

                  std::shared_ptr<io::IOContext> io_ctx_holder = nullptr);

    NewJsonReader(RuntimeProfile* profile, const TFileScanRangeParams& params,

                  const TFileRangeDesc& range, const std::vector<SlotDescriptor*>& file_slot_descs,

                  size_t batch_size, io::IOContext* io_ctx,

                  std::shared_ptr<io::IOContext> io_ctx_holder = nullptr);

    ~NewJsonReader() override;

    Status init_reader(

            const std::unordered_map<std::string, VExprContextSPtr>& col_default_value_ctx,

            bool is_load);

    Status _do_get_next_block(Block* block, size_t* read_rows, bool* eof) override;

    Status _get_columns_impl(std::unordered_map<std::string, DataTypePtr>* name_to_type) override;

    // Row-based readers control throughput via row count, not byte budget.

    // The FileScanner's AdaptiveBlockSizePredictor converts the byte budget

    // into a predicted row count and calls set_batch_size() with it.

    void set_batch_size(size_t batch_size) override;

    size_t get_batch_size() const override { return _batch_size; }

    Status init_schema_reader() override;

    Status get_parsed_schema(std::vector<std::string>* col_names,

                             std::vector<DataTypePtr>* col_types) override;

protected:

    // ---- Unified init_reader(ReaderInitContext*) overrides ----

    Status _open_file_reader(ReaderInitContext* ctx) override;

    Status _do_init_reader(ReaderInitContext* ctx) override;

    void _collect_profile_before_close() override;

private:

    Status _get_range_params();

    void _init_system_properties();

    void _init_file_description();

    Status _open_file_reader(bool need_schema);

    Status _open_line_reader();

    Status _parse_jsonpath_and_json_root();

    Status _read_json_column(RuntimeState* state, Block& block,

                             const std::vector<SlotDescriptor*>& slot_descs, bool* is_empty_row,

                             bool* eof);

    Status _read_one_message(DorisUniqueBufferPtr<uint8_t>* file_buf, size_t* read_size);

    // StreamLoadPipe::read_one_message only reads a portion of the data when stream loading with a chunked transfer HTTP request.

    // Need to read all the data before performing JSON parsing.

    Status _read_one_message_from_pipe(DorisUniqueBufferPtr<uint8_t>* file_buf, size_t* read_size);

    // simdjson, replace none simdjson function if it is ready

    Status _simdjson_init_reader();

    Status _simdjson_parse_json(size_t* size, bool* is_empty_row, bool* eof,

                                simdjson::error_code* error);

    Status _get_json_value(size_t* size, bool* eof, simdjson::error_code* error,

                           bool* is_empty_row);

    Status _judge_empty_row(size_t size, bool eof, bool* is_empty_row);

    Status _handle_simdjson_error(simdjson::simdjson_error& error, Block& block, size_t num_rows,

                                  bool* eof);

    Status _simdjson_handle_simple_json(RuntimeState* state, Block& block,

                                        const std::vector<SlotDescriptor*>& slot_descs,

                                        bool* is_empty_row, bool* eof);

    Status _simdjson_handle_simple_json_write_columns(

            Block& block, const std::vector<SlotDescriptor*>& slot_descs, bool* is_empty_row,

            bool* eof);

    Status _simdjson_handle_flat_array_complex_json(RuntimeState* state, Block& block,

                                                    const std::vector<SlotDescriptor*>& slot_descs,

                                                    bool* is_empty_row, bool* eof);

    Status _simdjson_handle_flat_array_complex_json_write_columns(

            Block& block, const std::vector<SlotDescriptor*>& slot_descs, bool* is_empty_row,

            bool* eof);

    Status _simdjson_handle_nested_complex_json(RuntimeState* state, Block& block,

                                                const std::vector<SlotDescriptor*>& slot_descs,

                                                bool* is_empty_row, bool* eof);

    Status _simdjson_set_column_value(simdjson::ondemand::object* value, Block& block,

                                      const std::vector<SlotDescriptor*>& slot_descs, bool* valid);

    template <bool use_string_cache>

    Status _simdjson_write_data_to_column(simdjson::ondemand::value& value,

                                          const DataTypePtr& type_desc, IColumn* column_ptr,

                                          const std::string& column_name, DataTypeSerDeSPtr serde,

                                          bool* valid);

    Status _simdjson_write_columns_by_jsonpath(simdjson::ondemand::object* value,

                                               const std::vector<SlotDescriptor*>& slot_descs,

                                               Block& block, bool* valid);

    Status _append_error_msg(simdjson::ondemand::object* obj, std::string error_msg,

                             std::string col_name, bool* valid);

    size_t _column_index(const StringRef& name, size_t key_index);

    Status (NewJsonReader::*_vhandle_json_callback)(RuntimeState* state, Block& block,

                                                    const std::vector<SlotDescriptor*>& slot_descs,

                                                    bool* is_empty_row, bool* eof);

    Status _get_column_default_value(

            const std::vector<SlotDescriptor*>& slot_descs,

            const std::unordered_map<std::string, VExprContextSPtr>& col_default_value_ctx);

    Status _fill_missing_column(SlotDescriptor* slot_desc, DataTypeSerDeSPtr serde,

                                IColumn* column_ptr, bool* valid);

    // fe will add skip_bitmap_col to _file_slot_descs iff the target olap table has skip_bitmap_col

    // and the current load is a flexible partial update

    // flexible partial update can not be used when user specify jsonpaths, so we just fill the skip bitmap

    // in `_simdjson_handle_simple_json` and `_vhandle_simple_json` (which will be used when jsonpaths is not specified)

    bool _should_process_skip_bitmap_col() const { return skip_bitmap_col_idx != -1; }

    void _append_empty_skip_bitmap_value(Block& block, size_t cur_row_count);

    void _set_skip_bitmap_mark(SlotDescriptor* slot_desc, IColumn* column_ptr, Block& block,

                               size_t cur_row_count, bool* valid);

    RuntimeState* _state = nullptr;

    RuntimeProfile* _profile = nullptr;

    ScannerCounter* _counter = nullptr;

    const TFileScanRangeParams& _params;

    const TFileRangeDesc& _range;

    io::FileSystemProperties _system_properties;

    io::FileDescription _file_description;

    const std::vector<SlotDescriptor*>& _file_slot_descs;

    io::FileReaderSPtr _file_reader;

    std::unique_ptr<NewPlainTextLineReader> _line_reader;

    bool _reader_eof;

    std::unique_ptr<Decompressor> _decompressor;

    TFileCompressType::type _file_compress_type;

    // When we fetch range doesn't start from 0 will always skip the first line

    bool _skip_first_line;

    std::string _line_delimiter;

    size_t _line_delimiter_length;

    uint32_t _next_row;

    size_t _total_rows;

    std::string _jsonpaths;

    std::string _json_root;

    bool _read_json_by_line;

    bool _strip_outer_array;

    bool _num_as_string;

    bool _fuzzy_parse;

    std::vector<std::vector<JsonPath>> _parsed_jsonpaths;

    std::vector<JsonPath> _parsed_json_root;

    bool _parsed_from_json_root = false; // to avoid parsing json root multiple times

    char _value_buffer[4 * 1024 * 1024]; // 4MB

    char _parse_buffer[512 * 1024];      // 512KB

    using Document = rapidjson::GenericDocument<rapidjson::UTF8<>, rapidjson::MemoryPoolAllocator<>,

                                                rapidjson::MemoryPoolAllocator<>>;

    rapidjson::MemoryPoolAllocator<> _value_allocator;

    rapidjson::MemoryPoolAllocator<> _parse_allocator;

    Document _origin_json_doc;   // origin json document object from parsed json string

    rapidjson::Value* _json_doc; // _json_doc equals _final_json_doc iff not set `json_root`

    std::unordered_map<std::string, int> _name_map;

    bool* _scanner_eof = nullptr;

    size_t _current_offset;

    io::IOContext* _io_ctx = nullptr;

    std::shared_ptr<io::IOContext> _io_ctx_holder;

    RuntimeProfile::Counter* _read_timer = nullptr;

    // ======SIMD JSON======

    // name mapping

    /// Hash table match `field name -> position in the block`. NOTE You can use perfect hash map.

    using NameMap = phmap::flat_hash_map<StringRef, size_t, StringRefHash>;

    NameMap _slot_desc_index;

    /// Cached search results for previous row (keyed as index in JSON object) - used as a hint.

    std::vector<NameMap::iterator> _prev_positions;

    /// Set of columns which already met in row. Exception is thrown if there are more than one column with the same name.

    std::vector<UInt8> _seen_columns;

    // simdjson

    DorisUniqueBufferPtr<uint8_t> _json_str_ptr;

    const uint8_t* _json_str = nullptr;

    static constexpr size_t _init_buffer_size = 1024 * 1024 * 8;

    size_t _padded_size = _init_buffer_size + simdjson::SIMDJSON_PADDING;

    size_t _original_doc_size = 0;

    std::string _simdjson_ondemand_padding_buffer;

    std::string _simdjson_ondemand_unscape_padding_buffer;

    // char _simdjson_ondemand_padding_buffer[_padded_size];

    simdjson::ondemand::document _original_json_doc;

    simdjson::ondemand::value _json_value;

    // for strip outer array

    // array_iter pointed to _array

    simdjson::ondemand::array_iterator _array_iter;

    simdjson::ondemand::array _array;

    std::unique_ptr<simdjson::ondemand::parser> _ondemand_json_parser;

    // column to default value string map

    std::unordered_map<std::string, std::string> _col_default_value_map;

    // From document of simdjson:

    // ```

    //   Important: a value should be consumed once. Calling get_string() twice on the same value is an error.

    // ```

    // We should cache the string_views to avoid multiple get_string() calls.

    struct StringViewHash {

        size_t operator()(const std::string_view& str) const {

            return std::hash<int64_t>()(reinterpret_cast<int64_t>(str.data()));

        }

    };

    struct StringViewEqual {

        bool operator()(const std::string_view& lhs, const std::string_view& rhs) const {

            return lhs.data() == rhs.data() && lhs.size() == rhs.size();

        }

    };

    std::unordered_map<std::string_view, std::string_view, StringViewHash, StringViewEqual>

            _cached_string_values;

    int32_t skip_bitmap_col_idx {-1};

    //Used to indicate whether it is a stream load. When loading, only data will be inserted into columnString.

    //If an illegal value is encountered during the load process, `_append_error_msg` should be called

    //instead of directly returning `Status::DataQualityError`

    bool _is_load = true;

    // In hive : create table xxx ROW FORMAT SERDE 'org.apache.hive.hcatalog.data.JsonSerDe';

    // Hive will not allow you to create columns with the same name but different case, including field names inside

    // structs, and will automatically convert uppercase names in create sql to lowercase.However, when Hive loads data

    // to table, the column names in the data may be uppercase,and there may be multiple columns with

    // the same name but different capitalization.We refer to the behavior of hive, convert all column names

    // in the data to lowercase,and use the last one as the insertion value

    bool _is_hive_table = false;

    // hive : org.openx.data.jsonserde.JsonSerDe, `ignore.malformed.json` prop.

    // If the variable is true, `null` will be inserted for llegal json format instead of return error.

    bool _openx_json_ignore_malformed = false;

    DataTypeSerDeSPtrs _serdes;

    DataTypeSerDe::FormatOptions _serde_options;

    // Adaptive batch size set by FileScanner.

    size_t _batch_size;

};

} // namespace doris