/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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.
 */
#pragma once

#include "velox/expression/fuzzer/ArgTypesGenerator.h"

namespace facebook::velox::fuzzer {

/// An argument type generator for decimal function. A map keyed on the pair of
/// precision and scale could be initialized with all possible input types.
/// Argument types are generated by looking up the map with the precision and
/// scale of return type, and randomly selecting valid input types. Derived
/// classes should call 'initialize' from the constructor and specify the number
/// of decimal arguments. They should also implement toReturnType with matching
/// number of pairs of precision and scale.
class DecimalArgTypesGeneratorBase : public ArgTypesGenerator {
 public:
  std::vector<TypePtr> generateArgs(
      const exec::FunctionSignature& signature,
      const TypePtr& returnType,
      FuzzerGenerator& rng) override;

 protected:
  // Computes result type for all possible pairs of decimal input types. Stores
  // the results in 'inputs_' map keyed by the precision and scale of return
  // type.
  // @param numArgs the number of decimal argument types. It only supports
  // initialization with one or two argument types.
  virtual void initialize(uint32_t numArgs);

  // Given precisions and scales of the inputs, returns precision and scale of
  // the result. Returns std::nullopt if a valid return type cannot be generated
  // with inputs. Used when the return type is generated with one pair of input
  // precision and scale.
  virtual std::optional<std::pair<int, int>> toReturnType(int p, int s) {
    VELOX_UNREACHABLE();
  }

  // Used when the return type is generated with two pairs of input precision
  // and scale.
  virtual std::optional<std::pair<int, int>>
  toReturnType(int p1, int s1, int p2, int s2) {
    VELOX_UNREACHABLE();
  }

 private:
  // Returns randomly selected pair of input types that produce the specified
  // result type.
  std::vector<TypePtr> findInputs(
      const TypePtr& returnType,
      FuzzerGenerator& rng) const;

  // Maps from the precision and scale of return type to corresponding input
  // types.
  std::unordered_map<std::pair<int, int>, std::vector<std::vector<TypePtr>>>
      inputs_;
};

} // namespace facebook::velox::fuzzer
