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

#include "paimon/format/avro/avro_stats_extractor.h"

#include <vector>

#include "arrow/api.h"
#include "arrow/c/abi.h"
#include "arrow/c/bridge.h"
#include "arrow/ipc/json_simple.h"
#include "gtest/gtest.h"
#include "paimon/common/data/binary_row.h"
#include "paimon/common/utils/date_time_utils.h"
#include "paimon/core/stats/simple_stats.h"
#include "paimon/core/stats/simple_stats_converter.h"
#include "paimon/format/avro/avro_file_format.h"
#include "paimon/format/avro/avro_format_writer.h"
#include "paimon/format/file_format_factory.h"
#include "paimon/fs/file_system.h"
#include "paimon/fs/local/local_file_system.h"
#include "paimon/memory/memory_pool.h"
#include "paimon/status.h"
#include "paimon/testing/utils/testharness.h"

namespace paimon::avro::test {

class AvroStatsExtractorTest : public ::testing::Test {
 public:
    void SetUp() override {}
    void TearDown() override {}

    void WriteAvroFile(const std::string& file_path,
                       const std::shared_ptr<arrow::ChunkedArray>& src_chunk_array,
                       const std::shared_ptr<arrow::Schema>& schema) const {
        ::ArrowSchema c_schema;
        ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok());

        ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileFormat> file_format,
                             FileFormatFactory::Get("avro", options_));
        ASSERT_OK_AND_ASSIGN(auto writer_builder,
                             file_format->CreateWriterBuilder(&c_schema, /*batch_size=*/1024));

        auto fs = std::make_shared<LocalFileSystem>();
        ASSERT_OK_AND_ASSIGN(std::unique_ptr<OutputStream> output_stream,
                             fs->Create(file_path, true));
        ASSERT_OK_AND_ASSIGN(auto writer, writer_builder->Build(std::move(output_stream), "null"));

        for (const auto& array : src_chunk_array->chunks()) {
            ::ArrowArray c_array;
            ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok());
            ASSERT_OK(writer->AddBatch(&c_array));
        }
        ASSERT_OK(writer->Flush());
        ASSERT_OK(writer->Finish());

        ASSERT_OK_AND_ASSIGN(auto file_status, fs->GetFileStatus(file_path));
        ASSERT_GT(file_status->GetLen(), 0);
    }

 private:
    std::map<std::string, std::string> options_ = {{Options::FILE_FORMAT, "avro"},
                                                   {Options::MANIFEST_FORMAT, "avro"}};
};

TEST_F(AvroStatsExtractorTest, TestPrimitiveStatsExtractor) {
    auto timezone = DateTimeUtils::GetLocalTimezoneName();
    arrow::FieldVector fields = {
        arrow::field("f1", arrow::int8()),
        arrow::field("f2", arrow::int16()),
        arrow::field("f3", arrow::int32()),
        arrow::field("f4", arrow::int64()),
        arrow::field("f5", arrow::float32()),
        arrow::field("f6", arrow::float64()),
        arrow::field("f7", arrow::utf8()),
        arrow::field("f8", arrow::binary()),
        arrow::field("f9", arrow::date32()),
        arrow::field("f10", arrow::timestamp(arrow::TimeUnit::NANO)),
        arrow::field("f11", arrow::decimal128(5, 2)),
        arrow::field("f12", arrow::boolean()),
        arrow::field("f13", arrow::timestamp(arrow::TimeUnit::SECOND)),
        arrow::field("f14", arrow::timestamp(arrow::TimeUnit::MILLI)),
        arrow::field("f15", arrow::timestamp(arrow::TimeUnit::MICRO)),
        arrow::field("f16", arrow::timestamp(arrow::TimeUnit::NANO)),
        arrow::field("f17", arrow::timestamp(arrow::TimeUnit::SECOND, timezone)),
        arrow::field("f18", arrow::timestamp(arrow::TimeUnit::MILLI, timezone)),
        arrow::field("f19", arrow::timestamp(arrow::TimeUnit::MICRO, timezone)),
        arrow::field("f20", arrow::timestamp(arrow::TimeUnit::NANO, timezone)),
    };
    auto schema = std::make_shared<arrow::Schema>(fields);
    auto array = std::dynamic_pointer_cast<arrow::StructArray>(
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
        [1, 11, 111, 1111, 1.1, 1.11, "Hello", "你好", 1234, "2033-05-18 03:33:20.0", "1.22", true, "2033-05-18 03:33:20", "2033-05-18 03:33:20.0", "2033-05-18 03:33:20.0", "2033-05-18 03:33:20.0", "2033-05-18 03:33:20", "2033-05-18 03:33:20.0", "2033-05-18 03:33:20.0", "2033-05-18 03:33:20.0"],
        [2, 22, 222, 2222, 2.2, 2.22, "World", "世界", -1234, "1899-01-01 00:59:20.001001001", "2.22", false, "1899-01-01 00:59:20", "1899-01-01 00:59:20", "1899-01-01 00:59:20", "1899-01-01 00:59:20.001001001","1899-01-01 00:59:20", "1899-01-01 00:59:20", "1899-01-01 00:59:20", "1899-01-01 00:59:20.001001001"],
        [null, null, 0, null, null, 0, null, null, null, null, null, null, null, null, null, null, null, null, null, null]
    ])")
            .ValueOrDie());
    auto src_chunk_array = std::make_shared<arrow::ChunkedArray>(arrow::ArrayVector({array}));

    auto dir = paimon::test::UniqueTestDirectory::Create();
    ASSERT_TRUE(dir);
    std::string file_path = dir->Str() + "/test.avro";
    WriteAvroFile(file_path, src_chunk_array, schema);

    AvroFileFormat format(options_);
    ::ArrowSchema arrow_schema;
    ASSERT_TRUE(arrow::ExportSchema(*schema, &arrow_schema).ok());
    ASSERT_OK_AND_ASSIGN(auto extractor, format.CreateStatsExtractor(&arrow_schema));
    auto fs = std::make_shared<LocalFileSystem>();
    ASSERT_OK_AND_ASSIGN(auto stats_with_info,
                         extractor->ExtractWithFileInfo(fs, file_path, GetDefaultPool()));
    const auto& column_stats = stats_with_info.first;
    const auto& file_stats = stats_with_info.second;

    ASSERT_EQ(column_stats.size(), 20);
    for (const auto& stats : column_stats) {
        ASSERT_EQ(stats->ToString(), "min null, max null, null count null");
    }
    ASSERT_EQ(3, file_stats.GetRowCount());
}

TEST_F(AvroStatsExtractorTest, TestNestedType) {
    arrow::FieldVector fields = {
        arrow::field("f0", arrow::list(arrow::float32())),
        arrow::field("f1", arrow::struct_({arrow::field("sub_f0", arrow::boolean()),
                                           arrow::field("sub_f1", arrow::int64())}))};
    auto schema = arrow::schema(fields);
    auto array = std::dynamic_pointer_cast<arrow::StructArray>(
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
        [null, [true, 2]],
        [[0.1, 0.3], [true, 1]],
        [[1.1, 1.2], null]
    ])")
            .ValueOrDie());
    auto src_chunk_array = std::make_shared<arrow::ChunkedArray>(arrow::ArrayVector({array}));

    auto dir = paimon::test::UniqueTestDirectory::Create();
    ASSERT_TRUE(dir);
    std::string file_path = dir->Str() + "/test.avro";
    WriteAvroFile(file_path, src_chunk_array, schema);

    AvroStatsExtractor extractor(options_);
    auto fs = std::make_shared<LocalFileSystem>();
    ASSERT_OK_AND_ASSIGN(auto results, extractor.Extract(fs, file_path, GetDefaultPool()));

    ASSERT_EQ(results.size(), 2);
    for (const auto& stats : results) {
        ASSERT_EQ(stats->ToString(), "min null, max null, null count null");
    }
}

TEST_F(AvroStatsExtractorTest, TestNullForAllType) {
    auto timezone = DateTimeUtils::GetLocalTimezoneName();
    arrow::FieldVector fields = {
        arrow::field("f0", arrow::boolean()),
        arrow::field("f1", arrow::int8()),
        arrow::field("f2", arrow::int16()),
        arrow::field("f3", arrow::int32()),
        arrow::field("f4", arrow::int64()),
        arrow::field("f5", arrow::float32()),
        arrow::field("f6", arrow::float64()),
        arrow::field("f7", arrow::utf8()),
        arrow::field("f8", arrow::binary()),
        arrow::field("f9", arrow::list(arrow::struct_({arrow::field("key", arrow::int8()),
                                                       arrow::field("value", arrow::int16())}))),
        arrow::field("f10", arrow::list(arrow::float32())),
        arrow::field("f11", arrow::struct_({arrow::field("f0", arrow::boolean()),
                                            arrow::field("f1", arrow::int64())})),
        arrow::field("f12", arrow::timestamp(arrow::TimeUnit::NANO)),
        arrow::field("f13", arrow::date32()),
        arrow::field("f14", arrow::decimal128(2, 2)),
        arrow::field("f15", arrow::decimal128(30, 2)),
        arrow::field("f16", arrow::timestamp(arrow::TimeUnit::SECOND)),
        arrow::field("f17", arrow::timestamp(arrow::TimeUnit::MILLI)),
        arrow::field("f18", arrow::timestamp(arrow::TimeUnit::MICRO)),
        arrow::field("f19", arrow::timestamp(arrow::TimeUnit::NANO)),
        arrow::field("f20", arrow::timestamp(arrow::TimeUnit::SECOND, timezone)),
        arrow::field("f21", arrow::timestamp(arrow::TimeUnit::MILLI, timezone)),
        arrow::field("f22", arrow::timestamp(arrow::TimeUnit::MICRO, timezone)),
        arrow::field("f23", arrow::timestamp(arrow::TimeUnit::NANO, timezone)),
    };
    auto schema = std::make_shared<arrow::Schema>(fields);
    auto src_array = std::dynamic_pointer_cast<arrow::StructArray>(
        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([
        [null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null]
    ])")
            .ValueOrDie());
    auto src_chunk_array = std::make_shared<arrow::ChunkedArray>(arrow::ArrayVector({src_array}));

    auto dir = paimon::test::UniqueTestDirectory::Create();
    ASSERT_TRUE(dir);
    std::string file_path = dir->Str() + "/test.avro";
    WriteAvroFile(file_path, src_chunk_array, schema);

    AvroStatsExtractor extractor(options_);
    auto fs = std::make_shared<LocalFileSystem>();
    ASSERT_OK_AND_ASSIGN(auto column_stats, extractor.Extract(fs, file_path, GetDefaultPool()));

    ASSERT_OK_AND_ASSIGN(auto stats,
                         SimpleStatsConverter::ToBinary(column_stats, GetDefaultPool().get()));
    ASSERT_EQ(stats.min_values_.HashCode(), 0xf890741a);
    ASSERT_EQ(stats.max_values_.HashCode(), 0xf890741a);
}

}  // namespace paimon::avro::test
