/*
 * Copyright (C) 2026 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/traced/probes/android_aflags/android_aflags_data_source.h"

#include "perfetto/base/build_config.h"
#include "perfetto/base/task_runner.h"
#include "perfetto/ext/base/base64.h"
#include "perfetto/ext/base/subprocess.h"
#include "perfetto/tracing/core/data_source_config.h"
#include "src/base/test/test_task_runner.h"
#include "src/tracing/core/trace_writer_for_testing.h"
#include "test/gtest_and_gmock.h"

#include "protos/perfetto/trace/android/android_aflags.gen.h"
#include "protos/perfetto/trace/android/android_aflags.pbzero.h"
#include "protos/perfetto/trace/trace_packet.gen.h"

// Aflags data source is Android-only.
#if PERFETTO_BUILDFLAG(PERFETTO_OS_ANDROID)
#define ANDROID_ONLY_TEST(x) x
#else
#define ANDROID_ONLY_TEST(x) DISABLED_##x
#endif

namespace perfetto {
namespace {

using ::testing::HasSubstr;

class TestAndroidAflagsDataSource : public AndroidAflagsDataSource {
 public:
  TestAndroidAflagsDataSource(const DataSourceConfig& ds_config,
                              base::TaskRunner* task_runner,
                              TracingSessionID session_id,
                              std::unique_ptr<TraceWriter> writer)
      : AndroidAflagsDataSource(ds_config,
                                task_runner,
                                session_id,
                                std::move(writer)) {}

  using AndroidAflagsDataSource::FinalizeAflagsCapture;

  struct FakeFlag {
    std::string pkg;
    std::string name;
    std::string namespace_val;
    std::string container;
    std::string value;
    std::string staged_value;
    uint32_t permission;
    uint32_t value_picked_from;
    uint32_t storage_backend;
    uint32_t type;
  };

  // Seeds subprocess state as if a poll completed but did not finalize yet.
  void SeedCompletedFakeSubprocess(const std::vector<FakeFlag>& flags) {
    protozero::HeapBuffered<protos::pbzero::AndroidAflags> msg;
    for (const auto& f : flags) {
      auto* flag_proto = msg->add_flags();
      flag_proto->AppendString(1, f.namespace_val);
      flag_proto->AppendString(2, f.name);
      flag_proto->AppendString(3, f.pkg);
      flag_proto->AppendString(4, f.container);
      flag_proto->AppendString(5, f.value);
      flag_proto->AppendString(6, f.staged_value);
      flag_proto->AppendVarInt(7, f.permission);
      flag_proto->AppendVarInt(8, f.value_picked_from);
      flag_proto->AppendVarInt(9, f.storage_backend);
      flag_proto->AppendVarInt(10, f.type);
    }
    std::vector<uint8_t> bytes = msg.SerializeAsArray();
    aflags_output_ = base::Base64Encode(bytes.data(), bytes.size());
    aflags_process_ = std::make_unique<base::Subprocess>(
        std::initializer_list<std::string>{"/bin/true"});
    aflags_process_->Call();
  }

  // Helper to simulate subprocess completion with given flags.
  void FakeSubprocessCompletion(const std::vector<FakeFlag>& flags) {
    SeedCompletedFakeSubprocess(flags);
    FinalizeAflagsCapture(/*error=*/"");
  }

  void FakeSubprocessError(const std::string& output) {
    aflags_output_ = output;
    aflags_process_ = std::make_unique<base::Subprocess>(
        std::initializer_list<std::string>{"/bin/false"});
    aflags_process_->Call();
    FinalizeAflagsCapture(/*error=*/"");
  }

  void FakeInvalidBase64() {
    aflags_output_ = "!!!not-base64!!!";
    aflags_process_ = std::make_unique<base::Subprocess>(
        std::initializer_list<std::string>{"/bin/true"});
    aflags_process_->Call();
    FinalizeAflagsCapture(/*error=*/"");
  }
};

class AndroidAflagsDataSourceTest : public ::testing::Test {
 protected:
  std::unique_ptr<TestAndroidAflagsDataSource> CreateDataSource(
      const DataSourceConfig& ds_config) {
    auto writer =
        std::unique_ptr<TraceWriterForTesting>(new TraceWriterForTesting());
    writer_raw_ = writer.get();
    return std::unique_ptr<TestAndroidAflagsDataSource>(
        new TestAndroidAflagsDataSource(ds_config, &task_runner_, 0,
                                        std::move(writer)));
  }

  base::TestTaskRunner task_runner_;
  TraceWriterForTesting* writer_raw_ = nullptr;
};

TEST_F(AndroidAflagsDataSourceTest, ANDROID_ONLY_TEST(EmitAflags)) {
  DataSourceConfig ds_config;
  ds_config.set_name("android.aflags");

  auto ds = CreateDataSource(ds_config);

  std::vector<TestAndroidAflagsDataSource::FakeFlag> flags;
  flags.push_back({
      "com.android.settings",
      "my_flag",
      "settings_ns",
      "system",
      "enabled",
      "disabled",
      static_cast<uint32_t>(
          protos::pbzero::AndroidAflags::FLAG_PERMISSION_READ_WRITE),
      static_cast<uint32_t>(
          protos::pbzero::AndroidAflags::VALUE_PICKED_FROM_LOCAL),
      static_cast<uint32_t>(
          protos::pbzero::AndroidAflags::FLAG_STORAGE_BACKEND_ACONFIGD),
      static_cast<uint32_t>(protos::pbzero::AndroidAflags::FLAG_TYPE_BOOLEAN),
  });

  ds->FakeSubprocessCompletion(flags);

  protos::gen::TracePacket packet = writer_raw_->GetOnlyTracePacket();
  ASSERT_TRUE(packet.has_android_aflags());

  const auto& aflags = packet.android_aflags();
  ASSERT_EQ(aflags.flags().size(), 1u);
  EXPECT_EQ(aflags.flags()[0].pkg(), "com.android.settings");
  EXPECT_EQ(aflags.flags()[0].name(), "my_flag");
  EXPECT_EQ(aflags.flags()[0].flag_namespace(), "settings_ns");
  EXPECT_EQ(aflags.flags()[0].container(), "system");
  EXPECT_EQ(aflags.flags()[0].value(), "enabled");
  EXPECT_EQ(aflags.flags()[0].staged_value(), "disabled");
  EXPECT_EQ(aflags.flags()[0].permission(),
            protos::gen::AndroidAflags::FLAG_PERMISSION_READ_WRITE);
  EXPECT_EQ(aflags.flags()[0].value_picked_from(),
            protos::gen::AndroidAflags::VALUE_PICKED_FROM_LOCAL);
  EXPECT_EQ(aflags.flags()[0].storage_backend(),
            protos::gen::AndroidAflags::FLAG_STORAGE_BACKEND_ACONFIGD);
  EXPECT_EQ(aflags.flags()[0].type(),
            protos::gen::AndroidAflags::FLAG_TYPE_BOOLEAN);
}

TEST_F(AndroidAflagsDataSourceTest,
       ANDROID_ONLY_TEST(SubprocessErrorEmitsErrorPacket)) {
  DataSourceConfig ds_config;
  ds_config.set_name("android.aflags");

  auto ds = CreateDataSource(ds_config);

  ds->FakeSubprocessError("aflags error message");

  protos::gen::TracePacket packet = writer_raw_->GetOnlyTracePacket();
  ASSERT_TRUE(packet.has_android_aflags());
  const auto& aflags = packet.android_aflags();
  EXPECT_TRUE(aflags.has_error());
  EXPECT_THAT(aflags.error(),
              testing::AllOf(HasSubstr("aflags failed"),
                             HasSubstr("aflags error message")));
}

TEST_F(AndroidAflagsDataSourceTest,
       ANDROID_ONLY_TEST(InvalidBase64EmitsErrorPacket)) {
  DataSourceConfig ds_config;
  ds_config.set_name("android.aflags");

  auto ds = CreateDataSource(ds_config);

  ds->FakeInvalidBase64();

  protos::gen::TracePacket packet = writer_raw_->GetOnlyTracePacket();
  ASSERT_TRUE(packet.has_android_aflags());
  const auto& aflags = packet.android_aflags();
  EXPECT_TRUE(aflags.has_error());
  EXPECT_THAT(aflags.error(), HasSubstr("Failed to decode aflags output"));
}

TEST_F(AndroidAflagsDataSourceTest,
       ANDROID_ONLY_TEST(FlushCompletesImmediatelyWhenIdle)) {
  DataSourceConfig ds_config;
  ds_config.set_name("android.aflags");
  auto ds = CreateDataSource(ds_config);

  bool callback_fired = false;
  ds->Flush(/*flush_request_id=*/42,
            [&callback_fired] { callback_fired = true; });

  EXPECT_TRUE(callback_fired);
}

TEST_F(AndroidAflagsDataSourceTest,
       ANDROID_ONLY_TEST(FlushDefersUntilSubprocessCompletes)) {
  DataSourceConfig ds_config;
  ds_config.set_name("android.aflags");
  auto ds = CreateDataSource(ds_config);

  std::vector<TestAndroidAflagsDataSource::FakeFlag> flags;
  flags.push_back({
      "com.android.test",
      "a_flag",
      "ns",
      "system",
      "enabled",
      "disabled",
      static_cast<uint32_t>(
          protos::pbzero::AndroidAflags::FLAG_PERMISSION_READ_ONLY),
      static_cast<uint32_t>(
          protos::pbzero::AndroidAflags::VALUE_PICKED_FROM_DEFAULT),
      static_cast<uint32_t>(
          protos::pbzero::AndroidAflags::FLAG_STORAGE_BACKEND_ACONFIGD),
      static_cast<uint32_t>(protos::pbzero::AndroidAflags::FLAG_TYPE_BOOLEAN),
  });
  ds->SeedCompletedFakeSubprocess(flags);

  bool callback_fired = false;
  ds->Flush(/*flush_request_id=*/1,
            [&callback_fired] { callback_fired = true; });
  EXPECT_FALSE(callback_fired);

  ds->FinalizeAflagsCapture(/*error=*/"");

  EXPECT_TRUE(callback_fired);
  protos::gen::TracePacket packet = writer_raw_->GetOnlyTracePacket();
  ASSERT_TRUE(packet.has_android_aflags());
  EXPECT_EQ(packet.android_aflags().flags().size(), 1u);
}

}  // namespace
}  // namespace perfetto
