// This file is part of the AspectC++ compiler 'ac++'.
// Copyright (C) 1999-2013  The 'ac++' developers (see aspectc.org)
//
// This program is free software;  you can redistribute it and/or
// modify it under the terms of the GNU General Public License as
// published by the Free Software Foundation; either version 2 of
// the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public
// License along with this program; if not, write to the Free
// Software Foundation, Inc., 59 Temple Place, Suite 330, Boston,
// MA  02111-1307  USA

#include "ClangWeaverBase.h"
#include "ACIntroducer.h"
#include "IntroductionUnit.h"
#include "LineDirectiveMgr.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Lex/Lexer.h"
#include "clang/Lex/Preprocessor.h"
#include "llvm/Support/MemoryBuffer.h"
#include "ACModel/Utils.h"
#include "version.h"

// Returns the location one past the end of the token.
static clang::SourceLocation getLocAfterToken(ClangToken tok) {
  unsigned len = tok.get ().getLength ();
  return tok.location ().getLocWithOffset (len);
}

const WeavePos &WeaverBase::get_pos_after_loc (clang::SourceLocation loc,
    WeavePos::Pos pos) {
  clang::SourceLocation result;
  if (loc.isValid() && loc.isMacroID ()) {
    clang::SourceLocation macro_end_loc;
    // See if we're at the end of a macro and get the location for the last
    // token if we are.
    bool is_at_end = clang::Lexer::isAtEndOfMacroExpansion(loc,
        getRewriter().getSourceMgr(), getRewriter().getLangOpts (), &macro_end_loc);
    if (is_at_end) {
      loc = macro_end_loc;
      // Now go on and get the location after this token.
    }
    clang::SourceLocation spell_loc = getRewriter().getSourceMgr().getSpellingLoc(loc);
    unsigned len = clang::Lexer::MeasureTokenLength(spell_loc, getRewriter().getSourceMgr(),
        getRewriter().getLangOpts ());
    result = loc.getLocWithOffset(len);
  }
  else {
    result = clang::Lexer::getLocForEndOfToken (loc, 0,
        getRewriter().getSourceMgr(), getRewriter().getLangOpts ());
  }

  assert (!result.isInvalid ());
  return weave_pos (result, pos);
}

const WeavePos &WeaverBase::weave_pos (ClangToken t, int p) {
  if (p == (int)WeavePos::WP_AFTER)
    return weave_pos (getLocAfterToken(t), p);
  else
    return weave_pos (t.location(), p);
}

const WeavePos &WeaverBase::weave_pos (clang::SourceLocation loc, int p) {
  assert(loc.isValid() && "Invalid SourceLocation!");

  // Optimization: If a location is directly before or after a macro we can turn
  // it into a non-macro location and avoid the extra work for macro expansion.
  // TODO: This is currently disabled because clang returns wrong locations for
  // some macro arguments.
  /*if (loc.isValid() && loc.isMacroID()) {
    clang::SourceManager &SM = _rewriter->getSourceMgr();
    clang::SourceLocation fixed_loc;
    if (SM.isAtStartOfImmediateMacroExpansion (loc, &fixed_loc))
      loc = fixed_loc;
    else if (SM.isAtEndOfImmediateMacroExpansion (loc, &fixed_loc))
      loc = fixed_loc;
  }*/

  // check whether the weaving position has already been created
  WPSet::iterator i = _positions.find (WeavePos (loc, (WeavePos::Pos)p));
  if (i != _positions.end ())
    return *i;

  // insert the position as a new one
  pair<WPSet::iterator, bool> result =
      _positions.insert (WeavePos (loc, (WeavePos::Pos)p));

  // the resulting iterator points to the new entry
  return *result.first;
}

const WeavePos &WeaverBase::header_pos () {
  return header_pos(ACFileID::main_file(_rewriter->getSourceMgr()));
}

const WeavePos &WeaverBase::header_pos (ACFileID acfid) {
  const clang::FileEntry *fentry = acfid.file_entry ();
  clang::FileID fid = _rewriter->getSourceMgr().translateFile (fentry);
  clang::SourceLocation first =
      _rewriter->getSourceMgr().getLocForStartOfFile(fid);
  return weave_pos (first, WeavePos::WP_HEADER);
}

const WeavePos &WeaverBase::footer_pos () {
  return footer_pos(ACFileID::main_file(_rewriter->getSourceMgr()));
}

const WeavePos &WeaverBase::footer_pos (ACFileID acfid) {
  const clang::FileEntry *fentry = acfid.file_entry ();
  clang::FileID fid = _rewriter->getSourceMgr().translateFile (fentry);
  clang::SourceLocation last =
      _rewriter->getSourceMgr().getLocForEndOfFile(fid);
  return weave_pos (last, WeavePos::WP_FOOTER);
}

// insert a generated string a given position
void WeaverBase::insert (const WeavePos &pos, const string &str,
                              bool before) {
  // return immediately if the string to paste is empty
  if (str == "")
    return;

  // If this is a bulk code insert wrap it in line directives.
  if (str.find('\n') != string::npos) {
    clang::PresumedLoc loc = _rewriter->getSourceMgr().getPresumedLoc(pos._loc);

    // <ac> directives are fixed up later on.
    std::ostringstream os;
    _line_mgr.directive (os, clang::PresumedLoc ());
    os << str;
    _line_mgr.directive (os, loc);

    if (before)
      pos._strings.push_front(os.str());
    else
      pos._strings.push_back(os.str());
    return;
  }

  // FIXME: handle before

  // store the unit in the right order for insertion (later when committing)
  if (before)
    pos._strings.push_front(str);
  else
    pos._strings.push_back(str);
}

// replace the text between two positions with some new text
void WeaverBase::replace (const WeavePos &from, const WeavePos &to,
    const string &str) {
  insert (from, str);
  kill (from, to);
}

void WeaverBase::collect_macro_token (clang::Preprocessor &PP, clang::Token &Result) {

  clang::SourceLocation loc = Result.getLocation();
  // ignore annotation (getSpelling fails for them!) and non-macro tokens
  if (Result.isAnnotation() || !loc.isMacroID())
    return;

  // TODO: ignore macros that are expand in a file, which does not belong to the project
  //       => improved efficiency; make sure the cache the last check's result

  bool invalid;
  std::string text = PP.getSpelling(Result, &invalid);
  if (invalid)
    return;

//  cout << "tok: " << Result.getName() << " "<< text << endl;
  
  clang::SourceManager &SM = _rewriter->getSourceMgr();
  clang::SourceLocation expansion_loc = SM.getExpansionLoc(Result.getLocation());
  clang::SourceLocation end_loc = loc.getLocWithOffset (text.size ());
  MacroMap::iterator mm_iter = _macros.find (expansion_loc);
  if (mm_iter == _macros.end ()) {
    MacroExpansion macro_expansion;
    macro_expansion._token_loc_map.insert (std::pair<clang::SourceLocation, unsigned int> (loc, 0u));
    macro_expansion._token_end_loc_map.insert (std::pair<clang::SourceLocation, unsigned int> (end_loc, text.size ()));
    macro_expansion._somewhere_inside = loc;
    macro_expansion._expanded_macro_text = text;
    _macros.insert (std::pair<clang::SourceLocation, MacroExpansion> (expansion_loc, macro_expansion));
//    cout << "new: " << text << endl;
  }
  else {
    std::pair<TokenLocMap::iterator, bool> ret;
    unsigned int pos = mm_iter->second._expanded_macro_text.size () + 1u;
    ret = mm_iter->second._token_loc_map.insert (std::pair<clang::SourceLocation, unsigned int> (loc, pos));
    if (ret.second) {
      mm_iter->second._token_end_loc_map.insert (std::pair<clang::SourceLocation, unsigned int> (end_loc, pos + text.size ()));
      mm_iter->second._expanded_macro_text += " ";
      mm_iter->second._expanded_macro_text += text;
//      cout << "ext: " << mm_iter->second._expanded_macro_text << endl;
    }
  }
}

bool WeaverBase::lookup_macro_token (clang::SourceLocation loc, clang::SourceLocation &mapped_loc) {
  clang::SourceManager &SM = _rewriter->getSourceMgr();

  // find the expansion information stored for 'loc' or return false
  clang::SourceLocation expansion_loc = SM.getExpansionLoc(loc);
  MacroMap::iterator mm_iter = _macros.find (expansion_loc);
  if (mm_iter == _macros.end ()) {
    cout << "fatal: macro not found in expansion location map" << endl;
    return false;
  }

  // generate a memory buffer for the expansion if not done yet
  clang::FileID &fid = mm_iter->second._fid;
  if (fid.isInvalid ()) {
    std::unique_ptr<llvm::MemoryBuffer> mb (llvm::MemoryBuffer::getMemBufferCopy(mm_iter->second._expanded_macro_text + " ", "<macro-expansion>"));
    fid = SM.createFileID(std::move (mb), clang::SrcMgr::C_User, 0, 0, expansion_loc);
  }
  
  clang::SourceLocation start = SM.getLocForStartOfFile(fid);
  TokenLocMap::iterator tl_iter = mm_iter->second._token_loc_map.find (loc);
  if (tl_iter == mm_iter->second._token_loc_map.end ()) {
    tl_iter = mm_iter->second._token_end_loc_map.find (loc);
    if (tl_iter == mm_iter->second._token_end_loc_map.end ()) {
      cout << "fatal: token not found" << endl;
      return false;
    }
  }
  unsigned int offset = tl_iter->second;
  mapped_loc = start.getLocWithOffset(offset);
  
  return true;
}


// kill the text between two positions
void WeaverBase::kill (const WeavePos &from, const WeavePos &to) {
  // TODO: this should not matter, the code shoud be able to deal with it
  assert (from._pos == WeavePos::WP_BEFORE);
  assert (to._pos == WeavePos::WP_AFTER);

  _deletions.push_back(LocPair(from._loc, to._loc));
}


void WeaverBase::commit_kills() {
  for (LocPairList::iterator i = _deletions.begin(); i != _deletions.end(); ++i)
    commit_kill (i->first, i->second);
}


// kill the text between two positions
void WeaverBase::commit_kill (clang::SourceLocation from, clang::SourceLocation to) {

  /*if (macro_problem (from) || macro_problem (to))
    return;

  bool has_nl = (strstr (from._token.text (), "\n") != NULL);
  if ((to._token && from._token != to._token) || has_nl)
    _line_mgr.insert ((Unit*)to._token.unit (), to._token, true);*/

  clang::Rewriter::RewriteOptions opts;
  opts.IncludeInsertsAtBeginOfRange = false;
  opts.IncludeInsertsAtEndOfRange   = false;
  clang::SourceManager &SM = _rewriter->getSourceMgr();

  // Handle kills inside a macro.

  if (from.isMacroID() || to.isMacroID()) {

    clang::SourceLocation loc1 = from;
    clang::SourceLocation expansion_loc1 = SM.getExpansionLoc(loc1);

    if (from.isMacroID() && !lookup_macro_token(from, loc1)) {
      cout << "fatal: 'from' location not found during kill" << endl;
      return;
    }

    clang::SourceLocation loc2 = to;
    clang::SourceLocation expansion_loc2 = SM.getExpansionLoc(loc2);
    if (to.isMacroID() && !lookup_macro_token(to, loc2)) {
      cout << "fatal: 'to' location not found during kill" << endl;
      return;
    }

    if (from.isMacroID() && !to.isMacroID()) {
      MacroMap::iterator macro_info = _macros.find (expansion_loc1);
      if (macro_info == _macros.end ()) {
        cout << "fatal: macro info not found in kill" << endl;
        return;
      }
      clang::SourceLocation macro_end = macro_info->second.endLoc (SM);
      _rewriter->RemoveText(clang::CharSourceRange::getCharRange(loc1, macro_end));

      clang::SourceLocation somewhere_inside = macro_info->second._somewhere_inside;
      clang::CharSourceRange source_range = SM.getExpansionRange(somewhere_inside);
      std::pair<clang::SourceLocation, clang::SourceLocation> range =
          std::pair<clang::SourceLocation, clang::SourceLocation> (source_range.getBegin(), source_range.getEnd());
      unsigned len = _rewriter->getRangeSize(clang::SourceRange(range.second, range.second));
      _rewriter->RemoveText(clang::CharSourceRange::getCharRange(
          range.second.getLocWithOffset(len), loc2), opts);
    }
    else if (!from.isMacroID() && to.isMacroID()) {
      MacroMap::iterator macro_info = _macros.find (expansion_loc2);
      if (macro_info == _macros.end ()) {
        cout << "fatal: macro info not found in kill" << endl;
        return;
      }
      clang::SourceLocation macro_start = macro_info->second.startLoc (SM);
      _rewriter->RemoveText(clang::CharSourceRange::getCharRange(macro_start, loc2));

      clang::SourceLocation somewhere_inside = macro_info->second._somewhere_inside;
      clang::CharSourceRange source_range = SM.getExpansionRange(somewhere_inside);
      std::pair<clang::SourceLocation, clang::SourceLocation> range =
          std::pair<clang::SourceLocation, clang::SourceLocation> (source_range.getBegin(), source_range.getEnd());
      _rewriter->RemoveText(clang::CharSourceRange::getCharRange(
          loc1, range.first), opts);
    }
    else {
      // both locations are in a macro instance

      // remove from 'from' to the macro end
      MacroMap::iterator macro_info1 = _macros.find (expansion_loc1);
      if (macro_info1 == _macros.end ()) {
        cout << "fatal: macro info not found in kill" << endl;
        return;
      }
      clang::SourceLocation macro_end = macro_info1->second.endLoc (SM);

      // remove from macro start to 'end'
      MacroMap::iterator macro_info2 = _macros.find (expansion_loc2);
      if (macro_info2 == _macros.end ()) {
        cout << "fatal: macro info not found in kill" << endl;
        return;
      }
      clang::SourceLocation macro_start = macro_info2->second.startLoc (SM);

      if (macro_info1 != macro_info2) {
        _rewriter->RemoveText(clang::CharSourceRange::getCharRange(loc1, macro_end));
        _rewriter->RemoveText(clang::CharSourceRange::getCharRange(macro_start, loc2));

        // remove from first macro expansion end to second macro expansion begin
        clang::SourceLocation somewhere_inside1 = macro_info1->second._somewhere_inside;
      clang::CharSourceRange source_range1 = SM.getExpansionRange(somewhere_inside1);
      std::pair<clang::SourceLocation, clang::SourceLocation> range1 =
          std::pair<clang::SourceLocation, clang::SourceLocation> (source_range1.getBegin(), source_range1.getEnd());
        unsigned len = _rewriter->getRangeSize(clang::SourceRange(range1.second, range1.second));

        clang::SourceLocation somewhere_inside2 = macro_info2->second._somewhere_inside;
      clang::CharSourceRange source_range2 = SM.getExpansionRange(somewhere_inside2);
      std::pair<clang::SourceLocation, clang::SourceLocation> range2 =
          std::pair<clang::SourceLocation, clang::SourceLocation> (source_range2.getBegin(), source_range2.getEnd());

        _rewriter->RemoveText(clang::CharSourceRange::getCharRange(
            range1.second.getLocWithOffset(len), range2.first), opts);
      }
      else {
        _rewriter->RemoveText(clang::CharSourceRange::getCharRange(loc1, loc2));
      }
    }
    return;
  }

  if (!from.isValid () || from.isMacroID ()||
      !to.isValid () || to.isMacroID ()) {
    cout << "Invalid or macro source loc in kill" << endl;
    return;
  }
  _rewriter->RemoveText (clang::CharSourceRange::getCharRange (from, to), opts);
}

void WeaverBase::commit_internal () {

  // commit all code deletions and clear the operations
  commit_kills();
  _deletions.clear();

  // commit code insertions
  clang::SourceManager &SM = _rewriter->getSourceMgr();
  // handle code that has to be inserted at weaving positions
  for (WPSet::iterator i = _positions.begin (); i != _positions.end (); ++i) {
    const WeavePos &wp = *i;
    if (wp._loc.isInvalid ()) {
      cout << "Invalid source loc" << endl;
      continue;
    }
    bool nl = false;
    bool last_nl = false;

    clang::SourceLocation loc;
//    cout << "weaving at " << endl;
//    wp._loc.dump(SM);
    if (wp._loc.isMacroID()) {
      // find the location in the genrated memory buffer
      if (!lookup_macro_token (wp._loc, loc)) {
        cout << "fatal: location not found during commit" << endl;
        continue;
      }
      // mark this macro instance for expansion
      clang::SourceLocation expansion_loc = SM.getExpansionLoc(wp._loc);
      _macro_expansions.insert (expansion_loc);
    }
    else
      loc = wp._loc;

    // perform the code insertion
    for (list<string>::iterator si = wp._strings.begin();
	 si != wp._strings.end(); ++si) {
      const string &str = *si;
      _rewriter->InsertTextAfter(loc, str);
      nl |= (str.find("\n") != string::npos);
      last_nl = (str[str.size() - 1] == '\n');
    }
    if (nl) { // at least one new line character was inserted here
      // generate and insert a #line directive to make sure that debuggers will
      // be able to show the original source locations
      clang::FileID fid = SM.getFileID(wp._loc);
#if CLANG_VERSION_NUMBER >= VERSION_NUMBER_12_0_0
      llvm::MemoryBufferRef u = SM.getBufferOrFake(fid);
#else
      const llvm::MemoryBuffer *u = SM.getBuffer(fid);
#endif
      if (IntroductionUnit::cast(u)) {
	// TODO: Why don't we insert the #line directive into introduced code? Check!
      }
      else {
        clang::PresumedLoc ploc = SM.getPresumedLoc(wp._loc);
        ostringstream directive;
        // TODO: this seems to be unnecessary as the line mgr issues \n
        if (!last_nl) directive << endl;
        _line_mgr.directive (directive, ploc);
	// TODO: Why wp._loc and not 'loc'? ...
        _rewriter->InsertTextAfter (wp._loc, directive.str ());
      }
    }

  }

  // clear the list of weaving positions used for this transaction
  _positions.clear();
}

void WeaverBase::resolve_macros() {
  clang::SourceManager &SM = _rewriter->getSourceMgr();

  for (set<clang::SourceLocation>::iterator i = _macro_expansions.begin ();
      i != _macro_expansions.end (); ++i) {
    MacroMap::iterator mm_iter = _macros.find (*i);
    if (mm_iter == _macros.end ()) {
      cout << "fatal: macro info not found in resolve_macros" << endl;
      continue;
    }
    clang::FileID fid = mm_iter->second._fid;
    clang::SourceLocation somewhere_inside = mm_iter->second._somewhere_inside;
//    cout << mm_iter->second._expanded_macro_text << endl;

    // expand macro
      clang::CharSourceRange source_range = SM.getExpansionRange(somewhere_inside);
      std::pair<clang::SourceLocation, clang::SourceLocation> range =
          std::pair<clang::SourceLocation, clang::SourceLocation> (source_range.getBegin(), source_range.getEnd());
    const auto *RB = _rewriter->getRewriteBufferFor(fid);
    if (!RB) {
      cout << "fatal: Rewrite buffer not found" << endl;
      continue;
    }
    string str(RB->begin(), RB->end());
//    cout << "Replace macro " << range.first.printToString(SM) << ": " << str << endl;
    _rewriter->ReplaceText(clang::SourceRange(range.first, range.second), str);
  }

  _macros.clear ();
  _macro_expansions.clear ();
}

bool WeaverBase::commit (clang::SourceManager &Target) {
  clang::SourceManager &SM = _rewriter->getSourceMgr();

  // Resolve weaving positions.
  commit_internal();

  // Put all macros back into the source.
  resolve_macros();

  for (clang::Rewriter::buffer_iterator i = _rewriter->buffer_begin(),
                                        e = _rewriter->buffer_end();
       i != e; ++i) {
#if CLANG_VERSION_NUMBER >= VERSION_NUMBER_18_1_3
    if (auto FE = SM.getFileEntryRefForID(i->first)) {
      std::string data(i->second.begin(), i->second.end());
      Target.overrideFileContents(FE.value(), llvm::MemoryBuffer::getMemBufferCopy(data, FE.value().getName()));
#else
    if (const clang::FileEntry *FE = SM.getFileEntryForID(i->first)) {
      std::string data(i->second.begin(), i->second.end());
      Target.overrideFileContents(FE, llvm::MemoryBuffer::getMemBufferCopy(data, FE->getName()));
#endif
    }
  }

  // renew rewrite
  const clang::LangOptions &LO = _rewriter->getLangOpts ();
  delete _rewriter;
  _rewriter = new clang::Rewriter (Target, LO);

  return true;
}

// postorder traversal of IntroductionUnits, inserting into their parents.
static void po_intros(WeaverBase &wb, ::ACIntroducer &intro,
                      IntroductionUnit *unit, unsigned depth = 0) {
  std::vector<IntroductionUnit *> &units = intro.get_intros_for(unit);

  for (unsigned i = 0, e = units.size(); i != e; ++i)
    po_intros(wb, intro, units[i], unit ? depth + 1 : depth);

  if (!unit)
    return;

  if (const auto *RB =
          wb.getRewriter().getRewriteBufferFor(unit->file_id())) {
    std::string str(RB->begin(), RB->end());
    if (depth == 0) // Insert the last unit via weave_pos, there may be macros.
      wb.insert(wb.weave_pos(unit->location(), WeavePos::WP_AFTER), str);
    else
      wb.getRewriter().InsertText(unit->location(), str);
  } else {
    std::unique_ptr<llvm::MemoryBuffer> buf (unit->buffer());
    llvm::StringRef str(buf->getBuffer());
    if (depth == 0) // Insert the last unit via weave_pos, there may be macros.
      wb.insert(wb.weave_pos(unit->location(), WeavePos::WP_AFTER), str.str());
    else
      wb.getRewriter().InsertText(unit->location(), str);
  }
}

bool WeaverBase::resolve_introduction_units(::ACIntroducer &intro) {
  // Resolve weaving positions.
  commit_internal();
  // Insert introduction units.
  po_intros(*this, intro, NULL);
  commit_internal();
  resolve_macros();
  return true;
}