/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

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 "tensorflow/js/ops/ts_op_gen.h"

#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/op_gen_lib.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"

namespace tensorflow {
namespace {

void ExpectContainsStr(absl::string_view s, absl::string_view expected) {
  EXPECT_TRUE(absl::StrContains(s, expected))
      << "'" << s << "' does not contain '" << expected << "'";
}

void ExpectDoesNotContainStr(absl::string_view s, absl::string_view expected) {
  EXPECT_FALSE(absl::StrContains(s, expected))
      << "'" << s << "' does not contain '" << expected << "'";
}

constexpr char kBaseOpDef[] = R"(
op {
  name: "Foo"
  input_arg {
    name: "images"
    type_attr: "T"
    number_attr: "N"
    description: "Images to process."
  }
  input_arg {
    name: "dim"
    description: "Description for dim."
    type: DT_FLOAT
  }
  output_arg {
    name: "output"
    description: "Description for output."
    type: DT_FLOAT
  }
  attr {
    name: "T"
    type: "type"
    description: "Type for images"
    allowed_values {
      list {
        type: DT_UINT8
        type: DT_INT8
      }
    }
    default_value {
      i: 1
    }
  }
  attr {
    name: "N"
    type: "int"
    has_minimum: true
    minimum: 1
  }
  summary: "Summary for op Foo."
  description: "Description for op Foo."
}
)";

// Generate TypeScript code
void GenerateTsOpFileText(const std::string& op_def_str,
                          const std::string& api_def_str,
                          std::string* ts_file_text) {
  Env* env = Env::Default();
  OpList op_defs;
  protobuf::TextFormat::ParseFromString(
      op_def_str.empty() ? kBaseOpDef : op_def_str, &op_defs);
  ApiDefMap api_def_map(op_defs);

  if (!api_def_str.empty()) {
    TF_ASSERT_OK(api_def_map.LoadApiDef(api_def_str));
  }

  const std::string& tmpdir = testing::TmpDir();
  const auto ts_file_path = io::JoinPath(tmpdir, "test.ts");

  WriteTSOps(op_defs, api_def_map, ts_file_path);
  TF_ASSERT_OK(ReadFileToString(env, ts_file_path, ts_file_text));
}

TEST(TsOpGenTest, TestImports) {
  std::string ts_file_text;
  GenerateTsOpFileText("", "", &ts_file_text);

  const std::string expected = R"(
import * as tfc from '@tensorflow/tfjs-core';
import {createTensorsTypeOpAttr, nodeBackend} from './op_utils';
)";
  ExpectContainsStr(ts_file_text, expected);
}

TEST(TsOpGenTest, InputSingleAndList) {
  const std::string api_def = R"pb(
    op { graph_op_name: "Foo" arg_order: "dim" arg_order: "images" }
  )pb";

  std::string ts_file_text;
  GenerateTsOpFileText("", api_def, &ts_file_text);

  const std::string expected = R"(
export function Foo(dim: tfc.Tensor, images: tfc.Tensor[]): tfc.Tensor {
)";
  ExpectContainsStr(ts_file_text, expected);
}

TEST(TsOpGenTest, TestVisibility) {
  const std::string api_def = R"(
op {
  graph_op_name: "Foo"
  visibility: HIDDEN
}
)";

  std::string ts_file_text;
  GenerateTsOpFileText("", api_def, &ts_file_text);

  const std::string expected = R"(
export function Foo(images: tfc.Tensor[], dim: tfc.Tensor): tfc.Tensor {
)";
  ExpectDoesNotContainStr(ts_file_text, expected);
}

TEST(TsOpGenTest, SkipDeprecated) {
  const std::string op_def = R"(
op {
  name: "DeprecatedFoo"
  input_arg {
    name: "input"
    type_attr: "T"
    description: "Description for input."
  }
  output_arg {
    name: "output"
    description: "Description for output."
    type: DT_FLOAT
  }
  attr {
    name: "T"
    type: "type"
    description: "Type for input"
    allowed_values {
      list {
        type: DT_FLOAT 
      }
    }
  }
  deprecation {
    explanation: "Deprecated."
  }
}
)";

  std::string ts_file_text;
  GenerateTsOpFileText(op_def, "", &ts_file_text);

  ExpectDoesNotContainStr(ts_file_text, "DeprecatedFoo");
}

TEST(TsOpGenTest, MultiOutput) {
  const std::string op_def = R"(
op {
  name: "MultiOutputFoo"
  input_arg {
    name: "input"
    description: "Description for input."
    type_attr: "T"
  }
  output_arg {
    name: "output1"
    description: "Description for output 1."
    type: DT_FLOAT
  }
  output_arg {
    name: "output2"
    description: "Description for output 2."
    type: DT_FLOAT
  }
  attr {
    name: "T"
    type: "type"
    description: "Type for input"
    allowed_values {
      list {
        type: DT_FLOAT 
      }
    }
  }
  summary: "Summary for op MultiOutputFoo."
  description: "Description for op MultiOutputFoo."
}
)";

  std::string ts_file_text;
  GenerateTsOpFileText(op_def, "", &ts_file_text);

  const std::string expected = R"(
export function MultiOutputFoo(input: tfc.Tensor): tfc.Tensor[] {
)";
  ExpectContainsStr(ts_file_text, expected);
}

TEST(TsOpGenTest, OpAttrs) {
  std::string ts_file_text;
  GenerateTsOpFileText("", "", &ts_file_text);

  const std::string expectedFooAttrs = R"(
  const opAttrs = [
    createTensorsTypeOpAttr('T', images),
    {name: 'N', type: nodeBackend().binding.TF_ATTR_INT, value: images.length}
  ];
)";

  ExpectContainsStr(ts_file_text, expectedFooAttrs);
}

}  // namespace
}  // namespace tensorflow
