// Copyright 2025 The ODML Authors.
//
// 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 "runtime/components/embedding_lookup/embedding_lookup_text.h"

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <filesystem>  // NOLINT: Required for path manipulation.
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"  // from @com_google_absl
#include "absl/types/span.h"  // from @com_google_absl
#include "litert/cc/litert_element_type.h"  // from @litert
#include "litert/cc/litert_environment.h"  // from @litert
#include "litert/cc/litert_expected.h"  // from @litert
#include "litert/cc/litert_layout.h"  // from @litert
#include "litert/cc/litert_macros.h"  // from @litert
#include "litert/cc/litert_model.h"  // from @litert
#include "litert/cc/litert_ranked_tensor_type.h"  // from @litert
#include "litert/cc/litert_tensor_buffer.h"  // from @litert
#include "litert/cc/litert_tensor_buffer_types.h"  // from @litert
#include "litert/test/matchers.h"  // from @litert

#ifndef EXPECT_OK
#define EXPECT_OK(status) EXPECT_TRUE((status).ok())
#endif

namespace litert::lm {

constexpr char kTestdataDir[] =
    "litert_lm/runtime/components/testdata/";

class EmbeddingLookupTextTest : public testing::Test {
 protected:
  absl::Status CreateModelFromFile() {
    // The model is a dummy embedding model that takes a token as input and
    // returns an embedding vector. The embedding table has the shape
    // [10, 1, 4, 32].
    // The values in the table are generated by the following formula:
    // for idx_0 in range(dim[0]):
    //   for idx_1 in range(dim[1]):
    //     for idx_2 in range(dim[2]):
    //       for idx_3 in range(dim[3]):
    //         table[idx_0, idx_1, idx_2, idx_3] =
    //                              idx_0*10000 + idx_1*1000 + idx_2*100 + idx_3
    auto model_path = std::filesystem::path(::testing::SrcDir()) /
                      kTestdataDir / "dummy_embedding_cpu_model.tflite";
    LITERT_ASSIGN_OR_RETURN(model_, Model::CreateFromFile(model_path.string()));
    return absl::OkStatus();
  }

  std::unique_ptr<EmbeddingLookupText> GetEmbeddingLookupText(
      std::optional<std::string> signature_key = std::nullopt) {
    if (!CreateModelFromFile().ok()) {
      return nullptr;
    }
    if (!model_.has_value()) {
      return nullptr;
    }
    auto status = EmbeddingLookupText::Create(&*model_, signature_key, &*env_);
    if (!status.ok()) {
      return nullptr;
    }
    return std::move(status.value());
  }

  Expected<TensorBuffer> GetTensorBuffer(Dimensions& dimensions) {
    size_t buffer_size = sizeof(float);
    for (auto dim : dimensions) {
      buffer_size *= dim;
    }
    Layout layout(dimensions);
    RankedTensorType ranked_tensor_type(ElementType::Float32,
                                        std::move(layout));

    return TensorBuffer::CreateManaged(*env_,
                                       ::litert::TensorBufferType::kHostMemory,
                                       ranked_tensor_type, buffer_size);
  }

  Expected<Environment> env_ = Environment::Create({});
  std::optional<Model> model_;
};

TEST_F(EmbeddingLookupTextTest, LookupDecodeVector) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  std::vector<float> output_vector(4 * 32);

  int32_t token = 1;
  EXPECT_OK(embedding->LookupDecode(token, output_vector));

  size_t offset = 0;
  // Dimensions 0 and 1 both have size 1.
  for (int idx2 = 0; idx2 < 4; ++idx2) {
    for (int idx3 = 0; idx3 < 32; ++idx3) {
      // Dimensions 0 and 1 both have size 1 so offset and expected value can
      // ignore them.
      float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
      EXPECT_NEAR(output_vector[offset++], expected_value, 1e-5);
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupDecodeVectorBadOutputVector) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  std::vector<float> output_vector(4 * 32 + 1);

  int32_t token = 1;
  EXPECT_THAT(
      embedding->LookupDecode(token, output_vector),
      testing::status::StatusIs(
          absl::StatusCode::kInvalidArgument,
          testing::HasSubstr(
              "The text embedding lookup output vector must be of size")));
}

TEST_F(EmbeddingLookupTextTest, LookupDecodeVectorNegativeToken) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  std::vector<float> output_vector(4 * 32);

  int32_t token = -1;
  EXPECT_OK(embedding->LookupDecode(token, output_vector));

  size_t offset = 0;
  // Dimensions 0 and 1 both have size 1.
  for (int idx2 = 0; idx2 < 4; ++idx2) {
    for (int idx3 = 0; idx3 < 32; ++idx3) {
      // Dimensions 0 and 1 both have size 1 so offset and expected value can
      // ignore them.
      float expected_value = 10000.0 * 0 + 100.0 * idx2 + idx3;
      EXPECT_NEAR(output_vector[offset++], expected_value, 1e-5);
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupDecode) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 1, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  int32_t token = 1;
  EXPECT_OK(embedding->LookupDecode(token, &output_tensor));

  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  // Dimensions 0 and 1 both have size 1.
  for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
    for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
      // Dimensions 0 and 1 both have size 1 so offset and expected value can
      // ignore them.
      size_t offset = idx2 * dimensions[3] + idx3;
      float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
      EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupDecodeBadOutputTensorDimNum) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 1, 4});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  EXPECT_THAT(embedding->LookupDecode(1, &output_tensor),
              testing::status::StatusIs(
                  absl::StatusCode::kInvalidArgument,
                  testing::HasSubstr("The output tensor from the Embedding "
                                     "model must be have the same "
                                     "number of dimensions")));
}

TEST_F(EmbeddingLookupTextTest, LookupDecodeBadOutputTensorDimSize) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 1, 4, 256});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  EXPECT_THAT(embedding->LookupDecode(1, &output_tensor),
              testing::status::StatusIs(
                  absl::StatusCode::kInvalidArgument,
                  testing::HasSubstr("The output tensor from the Embedding "
                                     "model must be have the same "
                                     "dimensions")));
}

TEST_F(EmbeddingLookupTextTest, LookupDecodeNullOutputTensor) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  EXPECT_THAT(embedding->LookupDecode(1, nullptr),
              testing::status::StatusIs(
                  absl::StatusCode::kInvalidArgument,
                  testing::HasSubstr("Decode output tensor buffer is null")));
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillVector) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  std::vector<float> output_vector(4 * 32);

  int32_t token = 1;
  EXPECT_OK(embedding->LookupPrefill(token, output_vector));

  size_t offset = 0;
  // Dimensions 0 and 1 both have size 1.
  for (int idx2 = 0; idx2 < 4; ++idx2) {
    for (int idx3 = 0; idx3 < 32; ++idx3) {
      // Dimensions 0 and 1 both have size 1 so offset and expected value can
      // ignore them.
      float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
      EXPECT_NEAR(output_vector[offset++], expected_value, 1e-5);
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillVectorBadOutputVector) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  std::vector<float> output_vector(4 * 32 + 1);

  int32_t token = 1;
  EXPECT_THAT(
      embedding->LookupPrefill(token, output_vector),
      testing::status::StatusIs(
          absl::StatusCode::kInvalidArgument,
          testing::HasSubstr(
              "The text embedding lookup output vector must be of size")));
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillVectorNegativeToken) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  std::vector<float> output_vector(4 * 32);

  int32_t token = -1;
  EXPECT_OK(embedding->LookupPrefill(token, output_vector));

  size_t offset = 0;
  // Dimensions 0 and 1 both have size 1.
  for (int idx2 = 0; idx2 < 4; ++idx2) {
    for (int idx3 = 0; idx3 < 32; ++idx3) {
      // Dimensions 0 and 1 both have size 1 so offset and expected value can
      // ignore them.
      float expected_value = 10000.0 * 0 + 100.0 * idx2 + idx3;
      EXPECT_NEAR(output_vector[offset++], expected_value, 1e-5);
    }
  }
}

TEST_F(EmbeddingLookupTextTest, GetFloatsPerToken) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);
  EXPECT_EQ(embedding->GetFloatsPerToken(), 4 * 32);
}

TEST_F(EmbeddingLookupTextTest, LookupPrefill) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 2, 3};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_OK(embedding->LookupPrefill(tokens_span, &output_tensor, 0));

  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  for (int idx0 = 0; idx0 < tokens.size(); ++idx0) {
    int token = tokens[idx0];
    for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
      for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
        // Since dimension 1 is of size 1, the offset and expected value can
        // ignore it.
        size_t offset =
            idx0 * dimensions[2] * dimensions[3] + idx2 * dimensions[3] + idx3;
        float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
        EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
      }
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillDecendingTokens) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {3, 2, 1};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_OK(embedding->LookupPrefill(tokens_span, &output_tensor, 0));

  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  for (int idx0 = 0; idx0 < tokens.size(); ++idx0) {
    int token = tokens[idx0];
    for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
      for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
        // Since dimension 1 is of size 1, the offset and expected value can
        // ignore it.
        size_t offset =
            idx0 * dimensions[2] * dimensions[3] + idx2 * dimensions[3] + idx3;
        float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
        EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
      }
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillRepeatedToken) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 1, 1};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_OK(embedding->LookupPrefill(tokens_span, &output_tensor, 0));

  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  for (int idx0 = 0; idx0 < tokens.size(); ++idx0) {
    int token = tokens[idx0];
    for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
      for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
        // Since dimension 1 is of size 1, the offset and expected value can
        // ignore it.
        size_t offset =
            idx0 * dimensions[2] * dimensions[3] + idx2 * dimensions[3] + idx3;
        float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
        EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
      }
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillBadOutputTensorDimNum) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 2, 3};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_THAT(embedding->LookupPrefill(tokens_span, &output_tensor, 0),
              testing::status::StatusIs(
                  absl::StatusCode::kInvalidArgument,
                  testing::HasSubstr("The output tensor from the Embedding "
                                     "model must be have the same "
                                     "number of dimensions")));
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillBadOutputTensorDimSize) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4, 256});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 2, 3};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_THAT(embedding->LookupPrefill(tokens_span, &output_tensor, 0),
              testing::status::StatusIs(
                  absl::StatusCode::kInvalidArgument,
                  testing::HasSubstr("The output tensor from the Embedding "
                                     "model must be have the same "
                                     "dimensions")));
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillBadOutputTensorDim0) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({3, 1, 4, 256});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 2, 3};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_THAT(
      embedding->LookupPrefill(tokens_span, &output_tensor, 0),
      testing::status::StatusIs(
          absl::StatusCode::kUnimplemented,
          testing::HasSubstr("The output tensor to fill from the Embedding "
                             "model must be have the 0th dimension as 1.")));
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillBadOutputTensorDim1) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 1, 4, 256});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 2, 3};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_THAT(
      embedding->LookupPrefill(tokens_span, &output_tensor, 0),
      testing::status::StatusIs(
          absl::StatusCode::kInvalidArgument,
          testing::HasSubstr("The output tensor to fill from the Embedding "
                             "model must have a 1st dimension that is at least "
                             "the same size as the number of tokens")));
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillLargerOutputTensor) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 4, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 2, 3};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_OK(embedding->LookupPrefill(tokens_span, &output_tensor, 0));

  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  for (int idx0 = 0; idx0 < tokens.size() + 1; ++idx0) {
    int token = 0;
    if (idx0 < tokens.size()) {
      token = tokens[idx0];
    }
    for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
      for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
        // Since dimension 1 is of size 1, the offset and expected value can
        // ignore it.
        size_t offset =
            idx0 * dimensions[2] * dimensions[3] + idx2 * dimensions[3] + idx3;
        float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
        EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
      }
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillNullOutputTensor) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  std::vector<int> tokens = {1, 2, 3};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_THAT(embedding->LookupPrefill(tokens_span, nullptr, 0),
              testing::status::StatusIs(
                  absl::StatusCode::kInvalidArgument,
                  testing::HasSubstr("Prefill output tensor buffer is null")));
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillNegativeToken) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  float filler_value = 9999.0;
  {
    auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
        output_tensor, ::litert::TensorBuffer::LockMode::kWrite);
    auto output_tensor_ptr =
        reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

    LITERT_ASSERT_OK_AND_ASSIGN(size_t output_tensor_size,
                                output_tensor.Size());
    for (int i = 0; i < output_tensor_size / sizeof(float); ++i) {
      output_tensor_ptr[i] = filler_value;
    }
  }

  std::vector<int> tokens = {1, -1, 2};
  absl::Span<const int> tokens_span(tokens);
  EXPECT_OK(embedding->LookupPrefill(tokens_span, &output_tensor, 0));

  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  for (int idx0 = 0; idx0 < tokens.size(); ++idx0) {
    int token = tokens[idx0];
    for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
      for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
        float expected_value;
        if (token < 0) {
          // If the token is negative, the expected value is the value of token
          // 0.
          expected_value = 10000.0 * 0 + 100.0 * idx2 + idx3;
        } else {
          // Since dimension 1 is of size 1, the offset and expected value can
          // ignore it.
          expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
        }
        size_t offset =
            idx0 * dimensions[2] * dimensions[3] + idx2 * dimensions[3] + idx3;
        EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
      }
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillWithOffset) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  {
    auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
        output_tensor, ::litert::TensorBuffer::LockMode::kWrite);
    auto output_tensor_ptr =
        reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

    LITERT_ASSERT_OK_AND_ASSIGN(size_t output_tensor_size,
                                output_tensor.Size());
    memset(output_tensor_ptr, 99, output_tensor_size);
  }

  std::vector<int> tokens = {1, 2};
  absl::Span<const int> tokens_span(tokens);
  const size_t token_offset = 1;
  const size_t float_offset = token_offset * 4 * 32;
  const size_t byte_offset = float_offset * sizeof(float);

  EXPECT_OK(embedding->LookupPrefill(tokens_span, &output_tensor, byte_offset));

  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  for (int i = 0; i < byte_offset; ++i) {
    EXPECT_EQ(reinterpret_cast<uint8_t*>(output_tensor_ptr)[i], 99);
  }

  for (int idx0 = 0; idx0 < tokens.size(); ++idx0) {
    int token = tokens[idx0];
    for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
      for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
        // Since dimension 1 is of size 1, the offset and expected value can
        // ignore it.
        size_t offset = float_offset + idx0 * dimensions[2] * dimensions[3] +
                        idx2 * dimensions[3] + idx3;
        float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
        EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
      }
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillWithOffsetAndDefaultEmbedding) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 3, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  {
    auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
        output_tensor, ::litert::TensorBuffer::LockMode::kWrite);
    auto output_tensor_ptr =
        reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

    LITERT_ASSERT_OK_AND_ASSIGN(size_t output_tensor_size,
                                output_tensor.Size());
    memset(output_tensor_ptr, 99, output_tensor_size);
  }

  std::vector<int> tokens = {1};
  absl::Span<const int> tokens_span(tokens);
  const size_t token_offset = 1;
  const size_t float_offset = token_offset * 4 * 32;
  const size_t byte_offset = float_offset * sizeof(float);

  EXPECT_OK(embedding->LookupPrefill(tokens_span, &output_tensor, byte_offset));

  // Check that the first token is not overwritten.
  auto output_tensor_lock_and_addr = ::litert::TensorBufferScopedLock::Create(
      output_tensor, ::litert::TensorBuffer::LockMode::kRead);
  auto output_tensor_ptr =
      reinterpret_cast<float*>(output_tensor_lock_and_addr->second);

  for (int i = 0; i < byte_offset; ++i) {
    EXPECT_EQ(reinterpret_cast<uint8_t*>(output_tensor_ptr)[i], 99);
  }

  for (int idx0 = 0; idx0 < tokens.size(); ++idx0) {
    int token = tokens[idx0];
    for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
      for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
        // Since dimension 1 is of size 1, the offset and expected value can
        // ignore it.
        size_t offset = float_offset + idx0 * dimensions[2] * dimensions[3] +
                        idx2 * dimensions[3] + idx3;
        float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
        EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
      }
    }
  }

  // Check that the last token is filled with the default embedding.
  int token = 0;
  for (int idx2 = 0; idx2 < dimensions[2]; ++idx2) {
    for (int idx3 = 0; idx3 < dimensions[3]; ++idx3) {
      // Since dimension 1 is of size 1, the offset and expected value can
      // ignore it.
      size_t offset = float_offset * 2 + idx2 * dimensions[3] + idx3;
      float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
      EXPECT_NEAR(output_tensor_ptr[offset], expected_value, 1e-5);
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupPrefillWithBadOffset) {
  std::unique_ptr<EmbeddingLookupText> embedding = GetEmbeddingLookupText();
  EXPECT_NE(embedding, nullptr);

  Dimensions dimensions({1, 2, 4, 32});
  LITERT_ASSERT_OK_AND_ASSIGN(TensorBuffer output_tensor,
                              GetTensorBuffer(dimensions));

  std::vector<int> tokens = {1, 2};
  absl::Span<const int> tokens_span(tokens);
  const size_t token_offset = 1;
  const size_t float_offset = token_offset * 4 * 32;
  const size_t byte_offset = float_offset * sizeof(float);

  EXPECT_THAT(
      embedding->LookupPrefill(tokens_span, &output_tensor, byte_offset),
      testing::status::StatusIs(
          absl::StatusCode::kInvalidArgument,
          testing::HasSubstr(
              "The byte offset and the total number of bytes to be written "
              "must not exceed the size of the output tensor")));
}

TEST_F(EmbeddingLookupTextTest, LookupDecodeVectorSpecifySignatureKey) {
  std::unique_ptr<EmbeddingLookupText> embedding =
      GetEmbeddingLookupText("serving_default");
  EXPECT_NE(embedding, nullptr);

  std::vector<float> output_vector(4 * 32);

  int32_t token = 1;
  EXPECT_OK(embedding->LookupDecode(token, output_vector));

  size_t offset = 0;
  // Dimensions 0 and 1 both have size 1.
  for (int idx2 = 0; idx2 < 4; ++idx2) {
    for (int idx3 = 0; idx3 < 32; ++idx3) {
      // Dimensions 0 and 1 both have size 1 so offset and expected value can
      // ignore them.
      float expected_value = 10000.0 * token + 100.0 * idx2 + idx3;
      EXPECT_NEAR(output_vector[offset++], expected_value, 1e-5);
    }
  }
}

TEST_F(EmbeddingLookupTextTest, LookupDecodeVectorSpecifySignatureKeyNotFound) {
  EXPECT_TRUE(CreateModelFromFile().ok());
  EXPECT_TRUE(model_.has_value());
  auto status = EmbeddingLookupText::Create(&*model_, "not_found");
  EXPECT_TRUE(!status.ok());
}

}  // namespace litert::lm