// See License.txt for copyright and terms of use
#include "squash_expr_rewriter.h"
#include "squash_patcher.h"
#include "squash_cmd.h"
#include "expr_visitor.h"
#include <iostream>
#include <vector>
#include <string>
#include <list>
#include <map>
#include <sstream>

static bool rewriteType(Type *t, std::string const &fromType,
                        std::string const &toType ) {
  AtomicType_list als;
  listAtomicTypes(t, als);
  bool changed = false;
  for (AtomicType_list::iterator i = als.begin(); i != als.end(); i++) {
    AtomicType *t = *i;
    if (!t->isCompoundType()) continue;
    CompoundType *c = t->asCompoundType();
    if (fromType == c->instName)
    {
      // don't like this c_str() one big
      c->typedefVar->name = c->instName = toType.c_str();
      changed = true;
    } else if (toType == c->instName) {
      return true;
    }
  }
  return changed;
}

// will break this  up into 2 visitors
// one that does renames and one that does decom
class ExprRewriterVisitor : public ExpressionVisitor {
public:
  ExprRewriterVisitor(ExprRewriter &rewriter) : rewriter(rewriter) {
  }

  virtual bool visitE_fieldAcc(E_fieldAcc *f) {
    str_map::const_iterator it =
        rewriter.cmd.memberRenames.find(f->fieldName->getName());
    if (it != rewriter.cmd.memberRenames.end()
	&& matchesType(f->obj->type))
    {
      MinMaxLineVisitor mmlv = traverseAndGuess(rewriter.patcher,
                                                f->fieldName, ";{}");
      rewriter.patcher.replaceWord(mmlv, it->first, it->second);
    }
    return true;
  }

  virtual bool visitE_arrow(E_arrow *a) {
    str_map::const_iterator it =
        rewriter.cmd.memberRenames.find(a->fieldName->getName());
    if (it != rewriter.cmd.memberRenames.end()
       && matchesType(a->obj->type))
    {
      MinMaxLineVisitor mmlv = traverseAndGuess(rewriter.patcher, 
                                                a->fieldName, ";{}");
      rewriter.patcher.replaceWord(mmlv, it->first, it->second);
    }
    return true;
  }

  template<class T> bool visitCast(T *c) {
    if (rewriter.cmd.squash)
    {
      if (rewriteType(c->ctype->decl->var->type,
                     rewriter.implementationClassName,
                     rewriter.interfaceClassName))
      {
        rewriter.rewriteSurrounding(c);
      }
    }
    return true;
  }

  virtual bool visitE_cast(E_cast *c) {
    return visitCast(c);
  }

  virtual bool visitE_keywordCast(E_keywordCast *c) {
    return visitCast(c);
  }

  virtual bool visitS_decl(S_decl *s) {
    Declaration *d = s->decl;
    if (rewriter.cmd.squash) {
      bool rewrite = false;
      for (FakeList<Declarator > *bs = d->decllist; !bs->isEmpty();
          bs = bs->butFirst())
      {
        rewrite = rewrite || rewriteType(bs->first()->var->type,
                                         rewriter.implementationClassName,
                                         rewriter.interfaceClassName);
      }
      if(rewrite) rewriter.rewriteSurrounding(s);
    }
    return true;
  }
  
  virtual bool visitE_variable(E_variable *v) {
    Variable *var = v->var;
    if (!var) return true;
    if (rewriter.cmd.squash) {
      if (v->name->isPQ_qualifier()) {
        PQ_qualifier *q = v->name->asPQ_qualifier();
        if (q->qualifier && rewriter.implementationClassName == q->qualifier) {
          q->qualifier = rewriter.interfaceClassName.c_str();
          rewriter.rewriteSurrounding(v);
          return true;
        }
      } 
    } else if (rewriter.interfaceClassName.empty() 
               && var->type->isFunctionType()) {
      str_map::const_iterator it =
      rewriter.cmd.memberRenames.find(var->name);
      if (it != rewriter.cmd.memberRenames.end())
      {
        MinMaxLineVisitor mmlv = traverseAndGuess(rewriter.patcher, v, ";{}");
        rewriter.patcher.replaceWord(mmlv, it->first, it->second);
      }      
    }
    return true;
  }
  
  bool matchesType(Type *t) {
    if (t->isPtrOrRef() || t->isArrayType()) {
      return matchesType(t->getAtType());
    }
    if (!t->isCompoundType()) {
      return false;
    }
    CompoundType *ct = t->asCompoundType();
    TS_classSpec *c = ct->syntax;
    if(rewriter.cmd.targetClasses.find(c->name->getName())
       != rewriter.cmd.targetClasses.end()) return true;
    
    // need to check base classes for inheritance involving templates
    for (FakeList<BaseClassSpec > *bs = c->bases;
	 !bs->isEmpty();
	 bs = bs->butFirst()) {
      BaseClassSpec *base = bs->first();
      if(matchesType(base->type->selfType))
	return true;
    }
    return false;
  }
  
  ExprRewriter &rewriter;
};

ExprRewriter::ExprRewriter(std::string const &interfaceClassName,
                           std::string const &implementationClassName,
                           Patcher &patcher,
                           SquashCmd const &cmd) :
   interfaceClassName(interfaceClassName),
   implementationClassName(implementationClassName),
   patcher(patcher),
   cmd(cmd),
   exprVisitor(new ExprRewriterVisitor(*this))
{
}

// Rewrites will be done on statement granularity within s
void ExprRewriter::rewrite(S_compound *s) {
  for (ASTListIterNC<Statement> i(s->stmts); !i.isDone(); i.adv()) {
    Statement *s = i.data();
    MinMaxLineVisitor mmlv;
    s->traverse(mmlv);
    stmLines.push_back(stMmlvList::value_type(s, mmlv));
  }
  ExprRewriterVisitor vis(*this);
  s->traverse(vis.loweredVisitor);

  if (rewrites.empty()) return;

  MinMaxLineVisitor mmlvPrev = rewrites.begin()->second;
  std::stringstream ss;
  for (stMmlvList::iterator i = rewrites.begin(); i != rewrites.end();
      i++)
  {
    Statement *s = i->first;
    MinMaxLineVisitor const &mmlv = i->second;
    if (mmlvPrev.maxLine < mmlv.minLine) {
      patcher.printPatch(ss.str(), mmlvPrev,
                         sourceLocManager->getFile(s->loc));
      ss.str("");
      mmlvPrev = mmlv;
    } else {
      mmlvPrev.maxLine = mmlv.maxLine;
    }
    print(ss, s);
  }
  if (!ss.str().empty()) {
    patcher.printPatch(ss.str(), mmlvPrev,
                       sourceLocManager->getFile(s->loc));
  }
}

template<class T> void ExprRewriter::rewriteSurrounding(T* expr) {
  MinMaxLineVisitor targetMmlv;
  expr->traverse(targetMmlv);
  for (stMmlvList::iterator i = stmLines.begin();
      i != stmLines.end();)
  {
    MinMaxLineVisitor const &mmlv = i->second;
    if (! (mmlv.minLine <= targetMmlv.minLine
           && targetMmlv.minLine <= mmlv.maxLine)) {
      i++;
      continue;
    }
    rewrites.push_back(*i);
    stmLines.erase(i++);
    // stop for now to avoid sourceLoc bugs
    break;
  }
}


void ExprRewriter::rewriteExpression(Expression *obj) {
  obj->traverse(exprVisitor->loweredVisitor);
}

void ExprRewriter::rewriteMemberInit(MemberInit *obj) {
  obj->traverse(exprVisitor->loweredVisitor);
}

ExprRewriter::~ExprRewriter() {
  delete exprVisitor;
}

void listAtomicTypes(AtomicType *t, AtomicType_list &ls )
{
  ls.push_back(t);
  if (t->isCompoundType()) {
    CompoundType* c = t->asCompoundType();
    TemplateInfo *ti = c->templateInfo();
    //if not a template can record class name
    if (ti) {
      ObjListIter<STemplateArgument> it(ti->arguments);
      for (;!it.isDone();it.adv()) {
        const STemplateArgument *ta = it.data();
        if (ta->isType()) {
          listAtomicTypes(ta->getType(), ls);
        }
      }
    }
  }
}

void listAtomicTypes(Type *t, AtomicType_list &ls) {
  //unpeel pointers
  if (t->isPtrOrRef()) {
    Type *bt = t->getAtType();
    listAtomicTypes(bt, ls);
  } else if (t->isCVAtomicType()) {
    listAtomicTypes(t->asCVAtomicType()->atomic, ls);
  } else if (t->isFunctionType()) {
    FunctionType* f = t->asFunctionType();
    listAtomicTypes(f->retType, ls);

    SObjListIter<Variable> it(f->params);

    for (;!it.isDone(); it.adv()) {
      listAtomicTypes(it.data()->type, ls);
    }
  }
}
