// Copyright 2026 Google LLC
//
// 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
//
//    https://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 "testing/cbor_utils.h"

#include <sys/types.h>

#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/strings/escaping.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/strings/strip.h"
#include "cppbor/cppbor.h"
#include "cppbor/cppbor_parse.h"
#include "nlohmann/json.hpp"

namespace cbor {
namespace {

using Json = ::nlohmann::json;

cppbor::Map B64ToBstr(cppbor::Map* map);
cppbor::Map BstrToB64(cppbor::Map* map);

constexpr absl::string_view kBase64Prefix = "b64'";
constexpr absl::string_view kBase64Suffix = "'";

class NoExceptionParser
    : public nlohmann::detail::json_sax_dom_parser<
          Json, nlohmann::detail::iterator_input_adapter<const char*>> {
 public:
  explicit NoExceptionParser(Json& j)
      : nlohmann::detail::json_sax_dom_parser<
            Json, nlohmann::detail::iterator_input_adapter<const char*>>(
            j, false) {}

  bool parse_error(std::size_t position, absl::string_view last_token,
                   const Json::exception& ex) {
    error_ = ex.what();
    return false;  // Stop processing.
  }

  std::string error() { return error_; }

 private:
  std::string error_;
};

std::string DecodePrefixedBase64(absl::string_view encoded) {
  auto stripped = absl::StripPrefix(encoded, kBase64Prefix);
  if (!absl::EndsWith(stripped, kBase64Suffix)) {
    LOG(FATAL) << "The base64-encoded string " << encoded
               << " does not end with the suffix " << kBase64Suffix;
  }
  stripped = absl::StripSuffix(stripped, kBase64Suffix);
  std::string decoded;
  if (!absl::Base64Unescape(stripped, &decoded)) {
    LOG(FATAL) << "Failed to decode base64: " << stripped;
  }
  return decoded;
}

std::string EncodePrefixedBase64(absl::string_view raw) {
  return absl::StrCat(kBase64Prefix, absl::Base64Escape(raw), kBase64Suffix);
}

std::string EncodeCbor(cppbor::Item* item) {
  std::vector<uint8_t> buf;
  buf.resize(item->encodedSize());
  item->encode(buf.data(), buf.data() + buf.size());
  return std::string(reinterpret_cast<const char*>(buf.data()), buf.size());
}

// Recursively replaces all base64-encoded Tstr with their decoded values as
// Bstr.
cppbor::Array B64ToBstr(cppbor::Array* array) {
  cppbor::Array new_array;
  for (auto& item : *array) {
    if (auto s = item->asTstr();
        s != nullptr && s->value().starts_with(kBase64Prefix)) {
      new_array.add(cppbor::Bstr(DecodePrefixedBase64(s->value())));
    } else if (auto m = item->asMap(); m != nullptr) {
      new_array.add(B64ToBstr(m));
    } else if (auto a = item->asArray(); a != nullptr) {
      new_array.add(B64ToBstr(a));
    } else {
      new_array.add(std::move(item));
    }
  }
  return new_array;
}

// Recursively replaces all base64-encoded Tstr with their decoded values as
// Bstr.
cppbor::Map B64ToBstr(cppbor::Map* map) {
  cppbor::Map new_map;
  for (auto& [key, value] : *map) {
    if (auto v = value->asTstr();
        v != nullptr && v->value().starts_with(kBase64Prefix)) {
      new_map.add(std::move(key),
                  cppbor::Bstr(DecodePrefixedBase64(v->value())));
    } else if (auto m = value->asMap(); m != nullptr) {
      new_map.add(std::move(key), B64ToBstr(m));
    } else if (auto v = value->asArray(); v != nullptr) {
      cppbor::Array new_array = B64ToBstr(v);
      new_map.add(std::move(key), std::move(new_array));
    } else {
      new_map.add(std::move(key), std::move(value));
    }
  }
  return new_map;
}

// Recursively replaces all base64-encoded Tstr with their decoded values as
// Bstr. Supports only arrays and maps.
std::string B64ToBstr(const uint8_t* begin, const uint8_t* end) {
  auto [item, new_position, error] = cppbor::parse(begin, end);
  if (!error.empty()) {
    LOG(FATAL) << "CBOR parsing failed: " << error;
  }
  if (new_position != end) {
    LOG(FATAL) << "Trailing bytes after the parsed CBOR item";
  }

  if (auto map = item->asMap(); map != nullptr) {
    auto new_map = B64ToBstr(map);
    return EncodeCbor(&new_map);
  } else if (auto array = item->asArray(); array != nullptr) {
    auto new_array = B64ToBstr(array);
    return EncodeCbor(&new_array);
  }
  LOG(FATAL) << "Input CBOR item is " << item->type()
             << "; expected array (128) or map (160).";
}

std::unique_ptr<cppbor::Item> UnwrapTag(std::unique_ptr<cppbor::Item> item) {
  while (item->asSemanticTag() != nullptr) {
    if (auto t = item->asTstr(); t != nullptr) return t->clone();
    if (auto u = item->asUint(); u != nullptr) return u->clone();
    if (auto i = item->asInt(); i != nullptr) return i->clone();
    if (auto b = item->asBool(); b != nullptr) return b->clone();
    if (auto f = item->asFloat(); f != nullptr) return f->clone();
    if (auto d = item->asDouble(); d != nullptr) return d->clone();
    if (auto s = item->asSimple(); s != nullptr) return s->clone();
    break;
  }
  return item;
}

// Recursively replaces all Bstr with base64-encoded Tstr.
cppbor::Array BstrToB64(cppbor::Array* array) {
  cppbor::Array new_array;
  for (auto& item_ref : *array) {
    auto item = UnwrapTag(std::move(item_ref));
    if (auto s = item->asBstr(); s != nullptr) {
      absl::string_view sv(reinterpret_cast<const char*>(s->value().data()),
                           s->value().size());
      new_array.add(cppbor::Tstr(EncodePrefixedBase64(sv)));
    } else if (auto m = item->asMap(); m != nullptr) {
      new_array.add(BstrToB64(m));
    } else if (auto a = item->asArray(); a != nullptr) {
      new_array.add(BstrToB64(a));
    } else {
      new_array.add(std::move(item));
    }
  }
  return new_array;
}

// Recursively replaces all Bstr with base64-encoded Tstr.
cppbor::Map BstrToB64(cppbor::Map* map) {
  cppbor::Map new_map;
  for (auto& [key_ref, value_ref] : *map) {
    auto key = UnwrapTag(std::move(key_ref));
    auto value = UnwrapTag(std::move(value_ref));
    if (auto s = value->asBstr(); s != nullptr) {
      absl::string_view sv(reinterpret_cast<const char*>(s->value().data()),
                           s->value().size());
      new_map.add(std::move(key), cppbor::Tstr(EncodePrefixedBase64(sv)));
    } else if (auto m = value->asMap(); m != nullptr) {
      new_map.add(std::move(key), BstrToB64(m));
    } else if (auto v = value->asArray(); v != nullptr) {
      cppbor::Array new_array = BstrToB64(v);
      new_map.add(std::move(key), std::move(new_array));
    } else {
      new_map.add(std::move(key), std::move(value));
    }
  }
  return new_map;
}

// Recursively replaces all Bstr with base64-encoded Tstr. Supports only arrays
// and maps.
std::string BstrToB64(absl::string_view raw) {
  auto [item, new_position, error] =
      cppbor::parse(reinterpret_cast<const uint8_t*>(raw.data()), raw.length());
  if (!error.empty()) {
    LOG(FATAL) << "CBOR parsing failed: " << error;
  }
  if (new_position !=
      reinterpret_cast<const uint8_t*>(raw.data() + raw.length())) {
    LOG(FATAL) << "Trailing bytes after the parsed CBOR item";
  }

  if (auto map = item->asMap(); map != nullptr) {
    auto new_map = BstrToB64(map);
    return EncodeCbor(&new_map);
  } else if (auto array = item->asArray(); array != nullptr) {
    auto new_array = BstrToB64(array);
    return EncodeCbor(&new_array);
  }
  LOG(FATAL) << "Input CBOR item is " << item->type()
             << "; expected array (128) or map (160).";
}

}  // namespace

std::string FromJson(absl::string_view json) {
  Json result;
  NoExceptionParser parser(result);
  if (!Json::sax_parse(json, &parser, Json::input_format_t::json)) {
    LOG(FATAL) << "JSON parse error: " << parser.error();
  }
  std::vector<uint8_t> cbor = Json::to_cbor(result);
  return B64ToBstr(cbor.data(), cbor.data() + cbor.size());
}

std::string ToJson(absl::string_view cbor) {
  std::string encoded = BstrToB64(cbor);
  Json result;
  NoExceptionParser parser(result);
  auto ia = nlohmann::detail::input_adapter(encoded);
  nlohmann::detail::binary_reader<Json, decltype(ia), NoExceptionParser> reader(
      std::move(ia), Json::input_format_t::cbor);
  if (!reader.sax_parse(Json::input_format_t::cbor, &parser, /*strict=*/true,
                        /*tag_handler=*/Json::cbor_tag_handler_t::ignore)) {
    LOG(FATAL) << "CBOR parse error: " << parser.error();
  }
  return result.dump();
}

}  // namespace cbor