/*
 * Copyright (C) 2024 The Android Open Source Project
 *
 * Licensed 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 "src/trace_processor/importers/proto/proto_trace_reader.h"

#include <cstdint>
#include <cstring>
#include <memory>
#include <utility>
#include <vector>

#include "perfetto/base/status.h"
#include "perfetto/protozero/scattered_heap_buffer.h"
#include "perfetto/trace_processor/trace_blob.h"
#include "perfetto/trace_processor/trace_blob_view.h"
#include "protos/perfetto/common/builtin_clock.pbzero.h"
#include "protos/perfetto/trace/trace.pbzero.h"
#include "src/trace_processor/importers/common/clock_tracker.h"
#include "src/trace_processor/importers/common/global_args_tracker.h"
#include "src/trace_processor/importers/common/global_metadata_tracker.h"
#include "src/trace_processor/importers/common/import_logs_tracker.h"
#include "src/trace_processor/importers/common/machine_tracker.h"
#include "src/trace_processor/importers/common/metadata_tracker.h"
#include "src/trace_processor/importers/proto/additional_modules.h"
#include "src/trace_processor/sorter/trace_sorter.h"
#include "src/trace_processor/storage/trace_storage.h"
#include "src/trace_processor/types/trace_processor_context.h"
#include "src/trace_processor/types/trace_processor_context_ptr.h"
#include "src/trace_processor/util/clock_synchronizer.h"
#include "src/trace_processor/util/descriptors.h"
#include "test/gtest_and_gmock.h"

#include "protos/perfetto/trace/clock_snapshot.pbzero.h"
#include "protos/perfetto/trace/remote_clock_sync.pbzero.h"

namespace perfetto::trace_processor {
namespace {

constexpr auto REALTIME = protos::pbzero::BUILTIN_CLOCK_REALTIME;
constexpr auto BOOTTIME = protos::pbzero::BUILTIN_CLOCK_BOOTTIME;

class ProtoTraceReaderTest : public ::testing::Test {
 public:
  ProtoTraceReaderTest() {
    host_context_.storage = std::make_unique<TraceStorage>();
    host_context_.trace_state =
        TraceProcessorContextPtr<TraceProcessorContext::TraceState>::MakeRoot(
            TraceProcessorContext::TraceState{TraceId{1}});
    host_context_.forked_context_state = TraceProcessorContextPtr<
        TraceProcessorContext::ForkedContextState>::MakeRoot();
    host_context_.machine_tracker =
        std::make_unique<MachineTracker>(&host_context_, kDefaultMachineId);
    host_context_.global_args_tracker =
        std::make_unique<GlobalArgsTracker>(host_context_.storage.get());
    host_context_.import_logs_tracker =
        std::make_unique<ImportLogsTracker>(&host_context_, TraceId(1));
    host_context_.trace_time_state = std::make_unique<TraceTimeState>(
        ClockId::Machine(protos::pbzero::BUILTIN_CLOCK_BOOTTIME));
    primary_sync_ = std::make_unique<ClockSynchronizer>(
        host_context_.trace_time_state.get(),
        std::make_unique<ClockSynchronizerListenerImpl>(&host_context_));
    host_context_.clock_tracker = std::make_unique<ClockTracker>(
        &host_context_,
        std::make_unique<ClockSynchronizerListenerImpl>(&host_context_),
        primary_sync_.get(), true);
    host_context_.global_metadata_tracker =
        std::make_unique<GlobalMetadataTracker>(host_context_.storage.get());
    host_context_.metadata_tracker =
        std::make_unique<MetadataTracker>(&host_context_);
    host_context_.sorter = std::make_unique<TraceSorter>(
        &host_context_, TraceSorter::SortingMode::kDefault);
    host_context_.descriptor_pool_ = std::make_unique<DescriptorPool>();
    host_context_.register_additional_proto_modules =
        &RegisterAdditionalModules;
    proto_trace_reader_ = std::make_unique<ProtoTraceReader>(&host_context_);
  }

  base::Status Tokenize() {
    trace_->Finalize();
    std::vector<uint8_t> trace_bytes = trace_.SerializeAsArray();
    std::unique_ptr<uint8_t[]> raw_trace(new uint8_t[trace_bytes.size()]);
    memcpy(raw_trace.get(), trace_bytes.data(), trace_bytes.size());
    auto status = proto_trace_reader_->Parse(TraceBlobView(
        TraceBlob::TakeOwnership(std::move(raw_trace), trace_bytes.size())));

    trace_.Reset();
    return status;
  }

 protected:
  protozero::HeapBuffered<protos::pbzero::Trace> trace_;
  TraceProcessorContext host_context_;
  std::unique_ptr<ClockSynchronizer> primary_sync_;
  std::unique_ptr<ProtoTraceReader> proto_trace_reader_;
};

TEST_F(ProtoTraceReaderTest, RemoteClockSync_Valid) {
  auto* machine_context =
      host_context_.ForkContextForMachineInCurrentTrace(0x1001);

  auto* packet = trace_->add_packet();
  packet->set_machine_id(0x1001);
  auto* remote_clock_sync = packet->set_remote_clock_sync();
  auto* synced_clocks = remote_clock_sync->add_synced_clocks();
  auto* client_clocks = synced_clocks->set_client_clocks();

  // First synced clock snapshots on both sides.
  auto* clock = client_clocks->add_clocks();
  clock->set_clock_id(BOOTTIME);
  clock->set_timestamp(10000);

  auto* host_clocks = synced_clocks->set_host_clocks();
  clock = host_clocks->add_clocks();
  clock->set_clock_id(BOOTTIME);
  clock->set_timestamp(120000);

  // Second synced clock snapshots on both sides.
  synced_clocks = remote_clock_sync->add_synced_clocks();

  client_clocks = synced_clocks->set_client_clocks();
  clock = client_clocks->add_clocks();
  clock->set_clock_id(BOOTTIME);
  clock->set_timestamp(25000);

  host_clocks = synced_clocks->set_host_clocks();
  clock = host_clocks->add_clocks();
  clock->set_clock_id(BOOTTIME);
  clock->set_timestamp(135000);

  ASSERT_TRUE(Tokenize().ok());
  ASSERT_EQ(1u,
            machine_context->clock_tracker->remote_clock_offsets_for_testing()
                .size());
}

TEST_F(ProtoTraceReaderTest, RemoteClockSync_Incomplete) {
  auto* machine_context =
      host_context_.ForkContextForMachineInCurrentTrace(0x1001);

  auto* packet = trace_->add_packet();
  packet->set_machine_id(0x1001);
  auto* remote_clock_sync = packet->set_remote_clock_sync();
  auto* synced_clocks = remote_clock_sync->add_synced_clocks();
  auto* client_clocks = synced_clocks->set_client_clocks();

  // First synced clock snapshots on both sides.
  auto* clock = client_clocks->add_clocks();
  clock->set_clock_id(BOOTTIME);
  clock->set_timestamp(10000);

  auto* host_clocks = synced_clocks->set_host_clocks();
  clock = host_clocks->add_clocks();
  clock->set_clock_id(BOOTTIME);
  clock->set_timestamp(120000);

  // Second synced clock snapshots on both sides.
  synced_clocks = remote_clock_sync->add_synced_clocks();

  client_clocks = synced_clocks->set_client_clocks();
  clock = client_clocks->add_clocks();
  clock->set_clock_id(BOOTTIME);
  clock->set_timestamp(25000);

  // Missing the second host CLOCK_BOOTTIME making it below the minimum
  // requirement for using the remote_clock_sync for calculating clock offset.

  ASSERT_TRUE(Tokenize().ok());
  // No valid clock offset.
  ASSERT_EQ(0u,
            machine_context->clock_tracker->remote_clock_offsets_for_testing()
                .size());
}

TEST_F(ProtoTraceReaderTest, CalculateClockOffset) {
  std::vector<ProtoTraceReader::SyncClockSnapshots> sync_clock_snapshots;
  ProtoTraceReader::SyncClockSnapshots snapshots;
  snapshots[BOOTTIME] = {120000, 10000};
  snapshots[REALTIME] = {135000, 25000};
  sync_clock_snapshots.push_back(std::move(snapshots));

  snapshots[BOOTTIME] = {140000, 20000};
  snapshots[REALTIME] = {150000, 35000};
  sync_clock_snapshots.push_back(std::move(snapshots));

  auto clock_offsets = perfetto::trace_processor::ProtoTraceReader::
      CalculateClockOffsetsForTesting(sync_clock_snapshots);
  ASSERT_EQ(2u, clock_offsets.size());
  // Client 10000      20000
  // Host     120000     140000
  // Estimated offsets: (10000 + 20000)/2 - 120000 = -105000,
  //                    20000 - (120000 + 140000) / 2 = -110000.
  // Average = -107500.
  ASSERT_EQ(-107500, clock_offsets[ClockId::Machine(BOOTTIME)]);
  // Client 25000      35000
  // Host     135000     150000
  // Estimated offsets: (25000 + 35000)/2 - 135000 = -105000,
  //                    35000 - (135000 + 150000) / 2 = -107500.
  // Average = -106250.
  ASSERT_EQ(-106250, clock_offsets[ClockId::Machine(REALTIME)]);
}

TEST_F(ProtoTraceReaderTest, CalculateClockOffset_AboveThreshold) {
  std::vector<ProtoTraceReader::SyncClockSnapshots> sync_clock_snapshots;
  ProtoTraceReader::SyncClockSnapshots snapshots;
  snapshots[BOOTTIME] = {120000, 10000};
  snapshots[REALTIME] = {135000, 25000};
  sync_clock_snapshots.push_back(std::move(snapshots));

  // 30 sec interval: the 2 clock snapshots will be considered 2 different
  // rounds of clock synchronization IPC exchange and won't be used.
  auto interval = 30ull * 1000 * 1000 * 1000;
  snapshots[BOOTTIME] = {120000 + interval, 10000 + interval};
  snapshots[REALTIME] = {135000 + interval, 25000 + interval};
  sync_clock_snapshots.push_back(std::move(snapshots));

  auto clock_offsets = perfetto::trace_processor::ProtoTraceReader::
      CalculateClockOffsetsForTesting(sync_clock_snapshots);
  ASSERT_EQ(0u, clock_offsets.size());
}

TEST_F(ProtoTraceReaderTest, CalculateClockOffset_MultiRounds) {
  std::vector<ProtoTraceReader::SyncClockSnapshots> sync_clock_snapshots;
  ProtoTraceReader::SyncClockSnapshots snapshots;
  // This emits clock offsets -105000, -110000.
  snapshots[BOOTTIME] = {120000, 10000};
  sync_clock_snapshots.push_back(std::move(snapshots));
  snapshots[BOOTTIME] = {140000, 20000};
  sync_clock_snapshots.push_back(std::move(snapshots));

  // The interval works as a delimiter of IPC exchange.
  auto interval = 30ull * 1000 * 1000 * 1000;

  // This emits clock offsets: (30000 + 45000) / 2 - 160000 = -122500,
  //                           45000 - (160000 + 170000) / 2 = -120000.
  snapshots[BOOTTIME] = {160000 + interval, 30000 + interval};
  sync_clock_snapshots.push_back(std::move(snapshots));
  snapshots[BOOTTIME] = {170000 + interval, 45000 + interval};
  sync_clock_snapshots.push_back(std::move(snapshots));

  auto clock_offsets = perfetto::trace_processor::ProtoTraceReader::
      CalculateClockOffsetsForTesting(sync_clock_snapshots);
  ASSERT_EQ(1u, clock_offsets.size());
  // Average(-105000, -110000, -122500, -120000) = -114375.
  ASSERT_EQ(-114375, clock_offsets[ClockId::Machine(BOOTTIME)]);
}

}  // namespace
}  // namespace perfetto::trace_processor
