/** @file

  A brief file description

  @section license License

  Licensed to the Apache Software Foundation (ASF) under one
  or more contributor license agreements.  See the NOTICE file
  distributed with this work for additional information
  regarding copyright ownership.  The ASF licenses this file
  to you 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 <tsutil/Regex.h>
#include <tsutil/Assert.h>

#define PCRE2_CODE_UNIT_WIDTH 8
#include <pcre2.h>

#include <array>
#include <vector>
#include <mutex>

static_assert(RE_CASE_INSENSITIVE == PCRE2_CASELESS, "Update RE_CASE_INSERSITIVE for current PCRE2 version.");
static_assert(RE_UNANCHORED == PCRE2_MULTILINE, "Update RE_MULTILINE for current PCRE2 version.");
static_assert(RE_ANCHORED == PCRE2_ANCHORED, "Update RE_ANCHORED for current PCRE2 version.");
static_assert(RE_NOTEMPTY == PCRE2_NOTEMPTY, "Update RE_NOTEMPTY for current PCRE2 version.");

static_assert(RE_ERROR_NOMATCH == PCRE2_ERROR_NOMATCH, "Update RE_ERROR_NOMATCH for current PCRE2 version.");
static_assert(RE_ERROR_NULL == PCRE2_ERROR_NULL, "Update RE_ERROR_NULL for current PCRE2 version.");

// PCRE2 10.30 added PCRE2_ENDANCHORED. Older PCRE2 (e.g., CentOS 7 ships 10.23) lacks it.
// On modern PCRE2 we pass RE_ENDANCHORED through natively (zero overhead); on old PCRE2 the
// bit is not a valid pcre2_compile option and would be rejected with PCRE2_ERROR_BADOPTION,
// so we transparently rewrite the pattern to "(?:pattern)\z" and strip the bit. See
// Regex::compile() for the rewrite. This preserves alternation-with-backtracking semantics
// (unlike a post-match length check, which stops at the first successful alternative).
#ifdef PCRE2_ENDANCHORED
static constexpr bool ATS_PCRE2_HAS_ENDANCHORED = true;
static_assert(RE_ENDANCHORED == PCRE2_ENDANCHORED, "Update RE_ENDANCHORED for current PCRE2 version.");
static_assert((RE_FULL_MATCH & PCRE2_ENDANCHORED) == 0, "RE_FULL_MATCH bit collides with PCRE2_ENDANCHORED");
#else
static constexpr bool ATS_PCRE2_HAS_ENDANCHORED = false;
#endif

// RE_FULL_MATCH is an ATS-only flag; it must not collide with any PCRE2 compile or match flag.
// We strip it before forwarding to pcre2_match, but a collision would cause spurious behavior
// if someone OR'd it into a flag word that's also passed elsewhere.
static_assert((RE_FULL_MATCH & PCRE2_ANCHORED) == 0, "RE_FULL_MATCH bit collides with PCRE2_ANCHORED");
static_assert((RE_FULL_MATCH & PCRE2_NO_UTF_CHECK) == 0, "RE_FULL_MATCH bit collides with PCRE2_NO_UTF_CHECK");
static_assert((RE_FULL_MATCH & PCRE2_CASELESS) == 0, "RE_FULL_MATCH bit collides with PCRE2_CASELESS");
static_assert((RE_FULL_MATCH & PCRE2_MULTILINE) == 0, "RE_FULL_MATCH bit collides with PCRE2_MULTILINE");
static_assert((RE_FULL_MATCH & PCRE2_NOTEMPTY) == 0, "RE_FULL_MATCH bit collides with PCRE2_NOTEMPTY");

//----------------------------------------------------------------------------
namespace
{
void *
my_malloc(size_t size, void * /*caller*/)
{
  void *ptr = malloc(size);
  return ptr;
}

void
my_free(void *ptr, void * /*caller*/)
{
  free(ptr);
}

//----------------------------------------------------------------------------
class RegexContext
{
public:
  static RegexContext *
  get_instance()
  {
    thread_local RegexContext ctx;
    return &ctx;
  }
  ~RegexContext()
  {
    if (_general_context != nullptr) {
      pcre2_general_context_free(_general_context);
    }
    if (_compile_context != nullptr) {
      pcre2_compile_context_free(_compile_context);
    }
    if (_match_context != nullptr) {
      pcre2_match_context_free(_match_context);
    }
    if (_jit_stack != nullptr) {
      pcre2_jit_stack_free(_jit_stack);
    }
  }
  pcre2_general_context *
  get_general_context()
  {
    return _general_context;
  }
  pcre2_compile_context *
  get_compile_context()
  {
    return _compile_context;
  }
  pcre2_match_context *
  get_match_context()
  {
    return _match_context;
  }

private:
  RegexContext()
  {
    _general_context = pcre2_general_context_create(my_malloc, my_free, nullptr);
    _compile_context = pcre2_compile_context_create(_general_context);
    _match_context   = pcre2_match_context_create(_general_context);
    _jit_stack       = pcre2_jit_stack_create(4096, 1024 * 1024, nullptr); // 1 page min and 1MB max
    pcre2_jit_stack_assign(_match_context, nullptr, _jit_stack);
  }
  pcre2_general_context *_general_context = nullptr;
  pcre2_compile_context *_compile_context = nullptr;
  pcre2_match_context   *_match_context   = nullptr;
  pcre2_jit_stack       *_jit_stack       = nullptr;
};

} // namespace

//----------------------------------------------------------------------------
struct RegexMatches::_MatchData {
  static pcre2_match_data *
  get(_MatchDataPtr const &p)
  {
    return static_cast<pcre2_match_data *>(p._ptr);
  }
  static void
  set(_MatchDataPtr &p, pcre2_match_data *ptr)
  {
    p._ptr = ptr;
  }
};

//----------------------------------------------------------------------------
RegexMatches::RegexMatches(uint32_t size)
{
  pcre2_general_context *ctx = pcre2_general_context_create(&RegexMatches::malloc, &RegexMatches::free, static_cast<void *>(this));

  pcre2_match_data *match_data = pcre2_match_data_create(size, ctx);
  debug_assert_message(match_data, "Failed to allocate pcre2 match data from custom context");

  _MatchData::set(_match_data, match_data);
}

//----------------------------------------------------------------------------
void *
RegexMatches::malloc(size_t size, void *caller)
{
  auto *matches = static_cast<RegexMatches *>(caller);

  // allocate from the buffer if possible
  if (size <= sizeof(matches->_buffer) - matches->_buffer_bytes_used) {
    void *ptr                    = matches->_buffer + matches->_buffer_bytes_used;
    matches->_buffer_bytes_used += size;
    return ptr;
  }

  return ::malloc(size);
}

void
RegexMatches::free(void *p, void *caller)
{
  auto *matches = static_cast<RegexMatches *>(caller);

  // Call free for any p outside _buffer
  // If the pcre2 context requests more data than fits in our builtin buffer, we will call malloc
  // to fulfil that request.
  // !his checks for any pointers outside of our buffer in order to free that memory up.
  //
  // nullptr is outside of our buffer, but its ok to call ::free with nullptr.
  if (!(p >= matches->_buffer && p < matches->_buffer + sizeof(matches->_buffer))) {
    ::free(p);
  }
}

//----------------------------------------------------------------------------
RegexMatches::~RegexMatches()
{
  auto ptr = _MatchData::get(_match_data);
  if (ptr != nullptr) {
    pcre2_match_data_free(ptr);
  }
}

//----------------------------------------------------------------------------
size_t *
RegexMatches::get_ovector_pointer()
{
  return pcre2_get_ovector_pointer(_MatchData::get(_match_data));
}

//----------------------------------------------------------------------------
int32_t
RegexMatches::size() const
{
  return _size;
}

//----------------------------------------------------------------------------
std::string_view
RegexMatches::operator[](size_t index) const
{
  // check if the index is valid
  if (index >= pcre2_get_ovector_count(_MatchData::get(_match_data))) {
    return std::string_view();
  }

  PCRE2_SIZE *ovector = pcre2_get_ovector_pointer(_MatchData::get(_match_data));
  return std::string_view(_subject.data() + ovector[2 * index], ovector[2 * index + 1] - ovector[2 * index]);
}

//----------------------------------------------------------------------------
struct RegexMatchContext::_MatchContext {
  static pcre2_match_context *
  get(_MatchContextPtr const &p)
  {
    return static_cast<pcre2_match_context *>(p._ptr);
  }
  static void
  set(_MatchContextPtr &p, pcre2_match_context *ptr)
  {
    p._ptr = ptr;
  }
};

//----------------------------------------------------------------------------
RegexMatchContext::RegexMatchContext()
{
  auto ctx = pcre2_match_context_create(nullptr);
  debug_assert_message(ctx, "Failed to allocate custom pcre2 match context");
  _MatchContext::set(_match_context, ctx);
}

//----------------------------------------------------------------------------
RegexMatchContext::RegexMatchContext(RegexMatchContext const &other)
{
  auto ptr = _MatchContext::get(other._match_context);
  if (nullptr != ptr) {
    pcre2_match_context *const ctx = pcre2_match_context_copy(ptr);
    _MatchContext::set(_match_context, ctx);
  }
}

//----------------------------------------------------------------------------
RegexMatchContext &
RegexMatchContext::operator=(RegexMatchContext const &other)
{
  if (&other != this) {
    auto ptr = _MatchContext::get(other._match_context);
    if (nullptr != ptr) {
      pcre2_match_context *const ctx = pcre2_match_context_copy(ptr);
      _MatchContext::set(_match_context, ctx);
    } else {
      _MatchContext::set(_match_context, nullptr);
    }
  }
  return *this;
}

//----------------------------------------------------------------------------
RegexMatchContext::~RegexMatchContext()
{
  auto ptr = _MatchContext::get(_match_context);
  if (ptr != nullptr) {
    pcre2_match_context_free(ptr);
  }
}

//----------------------------------------------------------------------------
void
RegexMatchContext::set_match_limit(uint32_t limit)
{
  auto ptr = _MatchContext::get(_match_context);
  debug_assert_message(ptr, "Failed to get the match context");
  if (ptr != nullptr) {
    pcre2_set_match_limit(ptr, limit);
  }
}

//----------------------------------------------------------------------------
struct Regex::_Code {
  static pcre2_code *
  get(_CodePtr const &p)
  {
    return static_cast<pcre2_code *>(p._ptr);
  }
  static void
  set(_CodePtr &p, pcre2_code *ptr)
  {
    p._ptr = ptr;
  }
};

//----------------------------------------------------------------------------
Regex::Regex(Regex &&that) noexcept
{
  _code = that._code;
  _Code::set(that._code, nullptr);
}

//----------------------------------------------------------------------------
Regex::~Regex()
{
  auto ptr = _Code::get(_code);
  if (ptr != nullptr) {
    pcre2_code_free(ptr);
  }
}

//----------------------------------------------------------------------------
bool
Regex::compile(std::string_view pattern, uint32_t flags)
{
  std::string error;
  int         erroroffset;

  return this->compile(pattern, error, erroroffset, flags);
}

//----------------------------------------------------------------------------
bool
Regex::compile(std::string_view pattern, std::string &error, int &erroroffset, uint32_t flags)
{
  // free the existing compiled regex if there is one
  if (auto ptr = _Code::get(_code); ptr != nullptr) {
    pcre2_code_free(ptr);
  }

  // get the RegexContext instance - should only be null when shutting down
  RegexContext *regex_context = RegexContext::get_instance();
  if (regex_context == nullptr) {
    return false;
  }

  // On PCRE2 < 10.30 the ENDANCHORED bit is not a valid pcre2_compile option. Rewrite
  // the pattern to "(?:pattern)\z" and strip the bit so pcre2 enforces end-of-subject
  // natively (including proper alternation backtracking). Zero overhead on modern PCRE2
  // where the bit is passed through unchanged.
  std::string      rewritten_pattern;
  std::string_view effective_pattern = pattern;
  bool             pattern_wrapped   = false;
  if constexpr (!ATS_PCRE2_HAS_ENDANCHORED) {
    if ((flags & RE_ENDANCHORED) != 0) {
      rewritten_pattern.reserve(pattern.size() + 6);
      rewritten_pattern.append("(?:").append(pattern).append(")\\z");
      effective_pattern  = rewritten_pattern;
      flags             &= ~static_cast<uint32_t>(RE_ENDANCHORED);
      pattern_wrapped    = true;
    }
  }

  PCRE2_SIZE error_offset;
  int        error_code;
  auto code = pcre2_compile(reinterpret_cast<PCRE2_SPTR>(effective_pattern.data()), effective_pattern.size(), flags, &error_code,
                            &error_offset, regex_context->get_compile_context());
  if (!code) {
    // Compensate for the "(?:" prefix so callers see offsets into their pattern, not ours.
    // If the offset is inside the prefix itself, clamp to 0.
    if (pattern_wrapped && error_offset >= 3) {
      erroroffset = static_cast<int>(error_offset - 3);
    } else {
      erroroffset = static_cast<int>(error_offset);
    }

    // get pcre2 error message
    PCRE2_UCHAR buffer[256];
    pcre2_get_error_message(error_code, buffer, sizeof(buffer));
    error.assign((char *)buffer);
    return false;
  }

  // support for JIT
  pcre2_jit_compile(code, PCRE2_JIT_COMPLETE);

  _Code::set(_code, code);

  return true;
}

//----------------------------------------------------------------------------
bool
Regex::exec(std::string_view subject) const
{
  if (_Code::get(_code) == nullptr) {
    return false;
  }
  RegexMatches matches;

  int count = this->exec(subject, matches);
  return count > 0;
}

//----------------------------------------------------------------------------
int32_t
Regex::exec(std::string_view subject, RegexMatches &matches) const
{
  auto code = _Code::get(_code);

  // check if there is a compiled regex
  if (code == nullptr) {
    return 0;
  }
  int count = pcre2_match(code, reinterpret_cast<PCRE2_SPTR>(subject.data()), subject.size(), 0, 0,
                          RegexMatches::_MatchData::get(matches._match_data), RegexContext::get_instance()->get_match_context());

  matches._size = count;

  if (count < 0) {
    return count;
  }

  if (count > 0) {
    matches._subject = subject;
  }

  return count;
}

//----------------------------------------------------------------------------
int32_t
Regex::exec(std::string_view subject, RegexMatches &matches, uint32_t flags, RegexMatchContext const *const matchContext) const
{
  auto code = _Code::get(_code);

  // check if there is a compiled regex
  if (code == nullptr) {
    return PCRE2_ERROR_NULL;
  }

  // Use the provided match context or fall back to the thread-global one.
  pcre2_match_context *match_context;
  if (nullptr == matchContext) {
    match_context = RegexContext::get_instance()->get_match_context();
  } else {
    match_context = RegexMatchContext::_MatchContext::get(matchContext->_match_context);
  }

  bool const     full_match  = (flags & RE_FULL_MATCH) != 0;
  uint32_t const pcre2_flags = flags & ~RE_FULL_MATCH;

  int rc = pcre2_match(code, reinterpret_cast<PCRE2_SPTR>(subject.data()), subject.size(), 0, pcre2_flags,
                       RegexMatches::_MatchData::get(matches._match_data), match_context);

  matches._size = rc;

  // match was successful
  if (rc >= 0) {
    matches._subject = subject;

    // match but the output vector was too small, adjust the size of the matches
    if (rc == 0) {
      matches._size = pcre2_get_ovector_count(RegexMatches::_MatchData::get(matches._match_data));
    }

    // Enforce full-subject consumption when requested (post-match length check, see RE_FULL_MATCH).
    if (full_match && matches[0].size() != subject.size()) {
      matches._size = PCRE2_ERROR_NOMATCH;
      rc            = PCRE2_ERROR_NOMATCH;
    }
  }

  return rc;
}

//----------------------------------------------------------------------------
int32_t
Regex::get_capture_count()
{
  int captures = -1;
  if (pcre2_pattern_info(_Code::get(_code), PCRE2_INFO_CAPTURECOUNT, &captures) != 0) {
    return -1;
  }
  return captures;
}

//----------------------------------------------------------------------------
DFA::~DFA() {}

//----------------------------------------------------------------------------
bool
DFA::build(const std::string_view pattern, unsigned flags)
{
  Regex       rxp;
  std::string string{pattern};

  if (!(flags & RE_UNANCHORED)) {
    flags |= RE_ANCHORED;
  }

  if (!rxp.compile(pattern, flags)) {
    return false;
  }
  _patterns.emplace_back(std::move(rxp), std::move(string));
  return true;
}

//----------------------------------------------------------------------------
int32_t
DFA::compile(const std::string_view pattern, unsigned flags)
{
  release_assert(_patterns.empty());
  this->build(pattern, flags);
  return _patterns.size();
}

//----------------------------------------------------------------------------
int32_t
DFA::compile(const std::string_view *const patterns, int npatterns, unsigned flags)
{
  _patterns.reserve(npatterns); // try to pre-allocate.
  for (int i = 0; i < npatterns; ++i) {
    this->build(patterns[i], flags);
  }
  return _patterns.size();
}

//----------------------------------------------------------------------------
int32_t
DFA::compile(const char *const *patterns, int npatterns, unsigned flags)
{
  _patterns.reserve(npatterns); // try to pre-allocate.
  for (int i = 0; i < npatterns; ++i) {
    this->build(patterns[i], flags);
  }
  return _patterns.size();
}

//----------------------------------------------------------------------------
int32_t
DFA::match(std::string_view str) const
{
  for (auto spot = _patterns.begin(), limit = _patterns.end(); spot != limit; ++spot) {
    if (spot->_re.exec(str)) {
      return spot - _patterns.begin();
    }
  }

  return -1;
}
