// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/optimization_guide/core/delivery/prediction_model_override.h"

#include <algorithm>

#include "base/base64.h"
#include "base/command_line.h"
#include "base/strings/string_util.h"
#include "base/strings/stringprintf.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace optimization_guide {

namespace {

#if BUILDFLAG(IS_WIN)
const char kTestAbsoluteFilePath[] = "C:\\absolute\\file\\path";
const char kOtherAbsoluteFilePath[] = "C:\\other\\absolute\\file\\path";
const char kTestRelativeFilePath[] = "relative\\file\\path";
#else
const char kTestAbsoluteFilePath[] = "/absolutefilepath";
const char kOtherAbsoluteFilePath[] = "/other/abs/file/path";
const char kTestRelativeFilePath[] = "relativefilepath";
#endif

}  // namespace

TEST(PredictionModelOverridesTest, NotSet) {
  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(0u, overrides.size());
}

TEST(PredictionModelOverridesTest, EmptyInput) {
  base::CommandLine::ForCurrentProcess()->AppendSwitch(
      kModelOverrideSwitch);
  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(0u, overrides.size());
}

TEST(PredictionModelOverridesTest, BadInput) {
  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch, "whatever");
  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(0u, overrides.size());
}

TEST(PredictionModelOverridesTest, InvalidOptimizationTarget) {
  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch,
      "notanoptimizationtarget:" + std::string(kTestAbsoluteFilePath));
  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(0u, overrides.size());
}

TEST(PredictionModelOverridesTest, RelativeFilePath) {
  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch, "OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD:" +
                                    std::string(kTestRelativeFilePath));
  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(0u, overrides.size());
}

TEST(PredictionModelOverridesTest, RelativeFilePathWithMetadata) {
  proto::Any metadata;
  metadata.set_type_url("sometypeurl");
  std::string encoded_metadata;
  metadata.SerializeToString(&encoded_metadata);
  encoded_metadata = base::Base64Encode(encoded_metadata);

  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch,
      base::StringPrintf("OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD:%s:%s",
                         kTestRelativeFilePath, encoded_metadata));
  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(0u, overrides.size());
}

TEST(PredictionModelOverridesTest, OneFilePath) {
  proto::Any metadata;
  metadata.set_type_url("sometypeurl");
  std::string encoded_metadata;
  metadata.SerializeToString(&encoded_metadata);
  encoded_metadata = base::Base64Encode(encoded_metadata);
#if BUILDFLAG(IS_WIN)
  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch,
      base::StringPrintf("OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD|%s|%s",
                         kTestAbsoluteFilePath, encoded_metadata));
#else
  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch,
      base::StringPrintf("OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD:%s:%s",
                         kTestAbsoluteFilePath, encoded_metadata));
#endif

  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(1u, overrides.size());
  auto* entry = overrides.Get(proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
  ASSERT_TRUE(entry);
  EXPECT_EQ(entry->path().MaybeAsASCII(), kTestAbsoluteFilePath);
  EXPECT_EQ(entry->metadata()->type_url(), "sometypeurl");
  ASSERT_FALSE(overrides.Get(proto::OPTIMIZATION_TARGET_PAGE_TOPICS));
}

TEST(PredictionModelOverridesTest, MultipleFilePath) {
  proto::Any metadata;
  metadata.set_type_url("sometypeurl");
  std::string encoded_metadata;
  metadata.SerializeToString(&encoded_metadata);
  encoded_metadata = base::Base64Encode(encoded_metadata);
#if BUILDFLAG(IS_WIN)
  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch,
      base::StringPrintf("OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD|%s,"
                         "OPTIMIZATION_TARGET_PAGE_TOPICS|%s|%s",
                         kTestAbsoluteFilePath, kOtherAbsoluteFilePath,
                         encoded_metadata));
#else
  base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
      kModelOverrideSwitch,
      base::StringPrintf("OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD:%s,"
                         "OPTIMIZATION_TARGET_PAGE_TOPICS:%s:%s",
                         kTestAbsoluteFilePath, kOtherAbsoluteFilePath,
                         encoded_metadata));
#endif

  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  EXPECT_EQ(2u, overrides.size());
  {
    auto* entry = overrides.Get(proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD);
    ASSERT_TRUE(entry);
    EXPECT_EQ(entry->path().MaybeAsASCII(), kTestAbsoluteFilePath);
    EXPECT_FALSE(entry->metadata());
  }
  {
    auto* entry = overrides.Get(proto::OPTIMIZATION_TARGET_PAGE_TOPICS);
    ASSERT_TRUE(entry);
    EXPECT_EQ(entry->path().MaybeAsASCII(), kOtherAbsoluteFilePath);
    EXPECT_EQ(entry->metadata()->type_url(), "sometypeurl");
  }
}

}  // namespace optimization_guide