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

#ifndef COMPONENTS_OPTIMIZATION_GUIDE_CORE_DELIVERY_PREDICTION_MODEL_OVERRIDE_H_
#define COMPONENTS_OPTIMIZATION_GUIDE_CORE_DELIVERY_PREDICTION_MODEL_OVERRIDE_H_

#include <cstddef>
#include <initializer_list>
#include <map>
#include <optional>
#include <string>

#include "base/command_line.h"
#include "base/files/file_path.h"
#include "base/functional/callback.h"
#include "components/optimization_guide/proto/models.pb.h"
#include "mojo/public/cpp/bindings/pending_remote.h"

namespace base {
class FilePath;
}  // namespace base

namespace unzip::mojom {
class Unzipper;
}  // namespace unzip::mojom

namespace unzip {
// TODO: crbug.com/421262905 - Avoid duplicating this alias.
using UnzipperFactory =
    base::RepeatingCallback<mojo::PendingRemote<mojom::Unzipper>()>;
}  // namespace unzip

namespace optimization_guide {

// Disables the fetching of models and overrides the file path and metadata to
// be used for the session to use what's passed via command-line instead of what
// is already stored.
//
// We expect that the string be a comma-separated string of model overrides with
// each model override be: OPTIMIZATION_TARGET_STRING:file_path or
// OPTIMIZATION_TARGET_STRING:file_path:base64_encoded_any_proto_model_metadata.
//
// It is possible this only works on Desktop since file paths are less easily
// accessible on Android, but may work.
inline constexpr char kModelOverrideSwitch[] =
    "optimization-guide-model-override";

// Return the separator used in the model override switch.
std::string ModelOverrideSeparator();

// Holds command-line overrides for prediction models.
class PredictionModelOverrides {
 public:
  struct Entry {
   public:
    using BuiltCallback =
        base::OnceCallback<void(std::unique_ptr<proto::PredictionModel>)>;
    Entry(proto::OptimizationTarget target,
          base::FilePath path,
          std::optional<proto::Any> metadata = std::nullopt);
    Entry(Entry&&);
    ~Entry();

    // Attempts to parse the Entry into a |proto::PredictionModel|.
    // Returns the result (or nullptr if there was an error) via `callback`.
    // In the event of an error, check LOG(ERROR).
    void BuildModel(const base::FilePath& base_model_dir,
                    unzip::UnzipperFactory unzipper_factory,
                    BuiltCallback callback) const;

    proto::OptimizationTarget target() const { return target_; }
    base::FilePath path() const { return path_; }
    std::optional<proto::Any> metadata() const { return metadata_; }

   private:
    proto::OptimizationTarget target_;
    // An absolute path to the model file.
    base::FilePath path_;
    // The provided metadata for the model.
    std::optional<proto::Any> metadata_ = std::nullopt;
  };

  ~PredictionModelOverrides();
  PredictionModelOverrides(PredictionModelOverrides&&);

  // Get the override entry for the given optimization target. Returns nullptr
  // if there is no override for the given optimization target.
  const Entry* Get(proto::OptimizationTarget optimization_target) const;

  size_t size() const { return overrides_.size(); }

  static PredictionModelOverrides ParseFromCommandLine(
      base::CommandLine* command_line);

 private:
  explicit PredictionModelOverrides(std::vector<Entry> overrides);

  std::map<proto::OptimizationTarget, Entry> overrides_;
};

}  // namespace optimization_guide

#endif  // COMPONENTS_OPTIMIZATION_GUIDE_CORE_DELIVERY_PREDICTION_MODEL_OVERRIDE_H_