// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
// Exceptions. See /LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "testing/file_test/autoupdate.h"

#include <fstream>

#include "absl/strings/string_view.h"
#include "common/check.h"
#include "common/ostream.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/FormatVariadic.h"

namespace Carbon::Testing {

// Converts a matched line number to an int, trimming whitespace.
static auto ParseLineNumber(absl::string_view matched_line_number) -> int {
  llvm::StringRef trimmed = matched_line_number;
  trimmed = trimmed.trim();
  // NOLINTNEXTLINE(google-runtime-int): API requirement.
  long long val;
  CARBON_CHECK(!llvm::getAsSignedInteger(trimmed, 10, val), "{0}",
               matched_line_number);
  return val;
}

FileTestAutoupdater::FileAndLineNumber::FileAndLineNumber(
    const LineNumberReplacement* replacement, int file_number,
    absl::string_view line_number)
    : replacement(replacement),
      file_number(file_number),
      line_number(ParseLineNumber(line_number)) {}

auto FileTestAutoupdater::CheckLine::RemapLineNumbers(
    const llvm::DenseMap<std::pair<int, int>, int>& output_line_remap,
    const llvm::SmallVector<int>& new_last_line_numbers) -> void {
  // Only need to do remappings when there's a line number replacement.
  if (!replacement_) {
    return;
  }

  bool found_one = false;
  // Use a cursor for the line so that we can't keep matching the same
  // content, which may occur when we keep a literal line number.
  int line_offset = 0;
  while (true) {
    // Rebuild the cursor each time because we're editing the line, which
    // could cause a reallocation.
    absl::string_view line_cursor = line_;
    line_cursor.remove_prefix(line_offset);
    // Look for a line number to replace. There may be multiple, so we
    // repeatedly check.
    absl::string_view matched_line_number;
    if (replacement_->has_file) {
      RE2::PartialMatch(line_cursor, *replacement_->re, nullptr,
                        &matched_line_number);
    } else {
      RE2::PartialMatch(line_cursor, *replacement_->re, &matched_line_number);
    }
    if (matched_line_number.empty()) {
      CARBON_CHECK(found_one, "{0}", line_);
      return;
    }
    found_one = true;

    // Update the cursor offset from the match.
    line_offset = matched_line_number.begin() - line_.c_str();

    // Calculate the new line number (possibly with new CHECK lines added, or
    // some removed).
    int old_line_number = ParseLineNumber(matched_line_number);
    int new_line_number = -1;
    if (auto remapped =
            output_line_remap.find({file_number(), old_line_number});
        remapped != output_line_remap.end()) {
      // Map old non-check lines to their new line numbers.
      new_line_number = remapped->second;
    } else {
      // We assume unmapped references point to the end-of-file.
      new_line_number = new_last_line_numbers[file_number()];
    }

    std::string replacement;
    if (output_file_number_ == file_number()) {
      int offset = new_line_number - output_line_number_;
      // Update the line offset in the CHECK line.
      const char* offset_prefix = offset < 0 ? "" : "+";
      replacement = llvm::formatv(
          replacement_->line_formatv.c_str(),
          llvm::formatv("[[@LINE{0}{1}]]", offset_prefix, offset));
    } else {
      // If the CHECK was written to a different file from the file that it
      // refers to, leave behind an absolute line reference rather than a
      // cross-file offset.
      replacement =
          llvm::formatv(replacement_->line_formatv.c_str(), new_line_number);
    }
    line_.replace(matched_line_number.data() - line_.data(),
                  matched_line_number.size(), replacement);
  }
}

auto FileTestAutoupdater::GetFileAndLineNumber(
    const llvm::DenseMap<llvm::StringRef, int>& file_to_number_map,
    int default_file_number, const std::string& check_line)
    -> FileAndLineNumber {
  for (const auto& replacement : line_number_replacements_) {
    if (replacement.has_file) {
      absl::string_view filename;
      absl::string_view line_number;
      if (RE2::PartialMatch(check_line, *replacement.re, &filename,
                            &line_number)) {
        if (auto it = file_to_number_map.find(filename);
            it != file_to_number_map.end()) {
          return FileAndLineNumber(&replacement, it->second, line_number);
        } else {
          return FileAndLineNumber(default_file_number);
        }
      }
    } else {
      // There's no file association, so we only look at the line, and assume
      // it refers to the default file.
      absl::string_view line_number;
      if (RE2::PartialMatch(check_line, *replacement.re, &line_number)) {
        return FileAndLineNumber(&replacement, default_file_number,
                                 line_number);
      }
    }
  }
  return FileAndLineNumber(default_file_number);
}

auto FileTestAutoupdater::BuildCheckLines(llvm::StringRef output,
                                          const char* label) -> CheckLines {
  if (output.empty()) {
    return CheckLines({});
  }

  // Prepare to look for filenames in lines.
  llvm::DenseMap<llvm::StringRef, int> file_to_number_map;
  for (auto [number, name] : llvm::enumerate(filenames_)) {
    file_to_number_map.insert({name, number});
  }

  // %t substitution means we may see TEST_TMPDIR in output.
  char* tmpdir_env = getenv("TEST_TMPDIR");
  CARBON_CHECK(tmpdir_env != nullptr);
  llvm::StringRef tmpdir = tmpdir_env;

  llvm::SmallVector<llvm::StringRef> lines(llvm::split(output, '\n'));
  // It's typical that output ends with a newline, but we don't want to add a
  // blank CHECK for it.
  if (lines.back().empty()) {
    lines.pop_back();
  }

  // `{{` and `[[` are escaped as a regex matcher.
  static RE2 double_brace_re(R"(\{\{)");
  static RE2 double_square_bracket_re(R"(\[\[)");
  // End-of-line whitespace is replaced with a regex matcher to make it visible.
  static RE2 end_of_line_whitespace_re(R"((\s+)$)");

  // The default file number for when no specific file is found.
  int default_file_number = 0;

  llvm::SmallVector<CheckLine> check_lines;
  for (const auto& line : lines) {
    // This code is relatively hot in our testing, and because when testing it
    // isn't run with an optimizer we benefit from making it use simple
    // constructs. For this reason, we avoid `llvm::formatv` and similar tools.
    std::string check_line;
    check_line.reserve(line.size() + strlen(label) + strlen("// CHECK:: "));
    check_line.append("// CHECK:");
    check_line.append(label);
    check_line.append(":");
    if (!line.empty()) {
      check_line.append(" ");
      check_line.append(line);
    }

    RE2::Replace(&check_line, double_brace_re, R"({{\\{\\{}})");
    RE2::Replace(&check_line, double_square_bracket_re, R"({{\\[\\[}})");
    RE2::Replace(&check_line, end_of_line_whitespace_re, R"({{\1}})");

    // Ignore TEST_TMPDIR in output.
    if (auto pos = check_line.find(tmpdir); pos != std::string::npos) {
      check_line.replace(pos, tmpdir.size(), "{{.+}}");
    }

    do_extra_check_replacements_(check_line);

    if (default_file_re_) {
      absl::string_view filename;
      if (RE2::PartialMatch(line, *default_file_re_, &filename)) {
        auto it = file_to_number_map.find(filename);
        CARBON_CHECK(it != file_to_number_map.end(),
                     "default_file_re had unexpected match in '{0}' (`{1}`)",
                     line, default_file_re_->pattern());
        default_file_number = it->second;
      }
    }
    auto file_and_line = GetFileAndLineNumber(file_to_number_map,
                                              default_file_number, check_line);
    check_lines.push_back(CheckLine(file_and_line, check_line));
  }
  return CheckLines(check_lines);
}

auto FileTestAutoupdater::AddRemappedNonCheckLine() -> void {
  new_lines_.push_back(non_check_line_);
  CARBON_CHECK(output_line_remap_
                   .insert({{non_check_line_->file_number(),
                             non_check_line_->line_number()},
                            ++output_line_number_})
                   .second);
}

auto FileTestAutoupdater::AddTips() -> void {
  CARBON_CHECK(tips_.empty(), "Should only add tips once");

  tips_.reserve(4);
  // This puts commands on a single line so that they can be easily copied.
  tips_.emplace_back("// TIP: To test this file alone, run:");
  tips_.emplace_back("// TIP:   " + test_command_);
  tips_.emplace_back("// TIP: To dump output, run:");
  tips_.emplace_back("// TIP:   " + dump_command_);

  for (const auto& tip : tips_) {
    new_lines_.push_back(&tip);
    ++output_line_number_;
  }
}

auto FileTestAutoupdater::ShouldAddCheckLine(const CheckLines& check_lines,
                                             bool to_file_end) const -> bool {
  return !autoupdate_split_ && check_lines.cursor != check_lines.lines.end() &&
         (check_lines.cursor->file_number() < output_file_number_ ||
          (check_lines.cursor->file_number() == output_file_number_ &&
           (to_file_end || check_lines.cursor->line_number() <=
                               non_check_line_->line_number())));
}

auto FileTestAutoupdater::AddCheckLines(CheckLines& check_lines,
                                        bool to_file_end) -> void {
  for (; ShouldAddCheckLine(check_lines, to_file_end); ++check_lines.cursor) {
    new_lines_.push_back(check_lines.cursor);
    check_lines.cursor->SetOutputLine(
        to_file_end ? "" : non_check_line_->indent(), output_file_number_,
        ++output_line_number_);
  }
}

auto FileTestAutoupdater::FinishFile(bool is_last_file) -> void {
  bool include_stdout = any_attached_stdout_lines_ || is_last_file;

  // At the end of each file, print any remaining lines which are associated
  // with the file.
  if (ShouldAddCheckLine(stderr_, /*to_file_end=*/true) ||
      (include_stdout && ShouldAddCheckLine(stdout_, /*to_file_end=*/true))) {
    // Ensure there's a blank line before any trailing CHECKs.
    if (!new_lines_.empty() && !new_lines_.back()->is_blank()) {
      new_lines_.push_back(&blank_line_);
      ++output_line_number_;
    }

    AddCheckLines(stderr_, /*to_file_end=*/true);
    if (include_stdout) {
      AddCheckLines(stdout_, /*to_file_end=*/true);
    }
  }

  new_last_line_numbers_.push_back(output_line_number_);
}

auto FileTestAutoupdater::StartSplitFile() -> void {
  // Advance the file.
  ++output_file_number_;
  output_line_number_ = 0;
  CARBON_CHECK(output_file_number_ == non_check_line_->file_number(),
               "Non-sequential file: {0}", non_check_line_->file_number());

  // Each following file has precisely one split line.
  CARBON_CHECK(non_check_line_->line_number() < 1,
               "Expected a split line, got {0}", *non_check_line_);
  // The split line is ignored when calculating line counts.
  new_lines_.push_back(non_check_line_);

  // Add any file-specific but line-unattached STDOUT messages here. STDERR is
  // handled through the main loop because it's before the next line.
  if (any_attached_stdout_lines_) {
    AddCheckLines(stdout_, /*to_file_end=*/false);
  }

  ++non_check_line_;
}

auto FileTestAutoupdater::Run(bool dry_run) -> bool {
  // Print everything until the autoupdate line.
  while (non_check_line_->line_number() != autoupdate_line_number_) {
    CARBON_CHECK(non_check_line_ != non_check_lines_.end() &&
                     non_check_line_->file_number() == 0,
                 "Missed autoupdate?");
    AddRemappedNonCheckLine();
    ++non_check_line_;
  }

  // Add the AUTOUPDATE line along with any early STDERR lines, so that the
  // initial batch of CHECK lines have STDERR before STDOUT. This also ensures
  // we don't insert a blank line before the STDERR checks if there are no more
  // lines after AUTOUPDATE.
  AddRemappedNonCheckLine();
  AddTips();
  if (!autoupdate_split_) {
    AddCheckLines(stderr_, /*to_file_end=*/false);
    if (any_attached_stdout_lines_) {
      AddCheckLines(stdout_, /*to_file_end=*/false);
    }
  }
  ++non_check_line_;

  // Loop through remaining content.
  while (non_check_line_ != non_check_lines_.end()) {
    if (output_file_number_ < non_check_line_->file_number()) {
      FinishFile(/*is_last_file=*/false);
      StartSplitFile();
      if (autoupdate_split_ &&
          output_file_number_ == static_cast<int>(filenames_.size())) {
        break;
      }
      continue;
    }

    // STDERR check lines are placed before the line they refer to, or as
    // early as possible if they don't refer to a line. Include all STDERR
    // lines until we find one that wants to go later in the file.
    AddCheckLines(stderr_, /*to_file_end=*/false);
    AddRemappedNonCheckLine();

    // STDOUT check lines are placed after the line they refer to, or at the
    // end of the file if none of them refers to a line.
    if (any_attached_stdout_lines_) {
      AddCheckLines(stdout_, /*to_file_end=*/false);
    }

    ++non_check_line_;
  }

  // When autoupdate_split_ was true, this will result in all check lines (and
  // only check lines) being added to the split by FinishFile. We don't use
  // autoupdate_split_ past this point.
  autoupdate_split_ = false;

  FinishFile(/*is_last_file=*/true);

  for (auto& check_line : stdout_.lines) {
    check_line.RemapLineNumbers(output_line_remap_, new_last_line_numbers_);
  }
  for (auto& check_line : stderr_.lines) {
    check_line.RemapLineNumbers(output_line_remap_, new_last_line_numbers_);
  }

  // Generate the autoupdated file.
  std::string new_content;
  llvm::raw_string_ostream new_content_stream(new_content);
  for (const auto& line : new_lines_) {
    new_content_stream << *line << '\n';
  }

  // Update the file on disk if needed.
  if (new_content == input_content_) {
    return false;
  }
  if (!dry_run) {
    std::ofstream out(file_test_path_);
    out << new_content;
  }
  return true;
}

}  // namespace Carbon::Testing
