#include "squash_outparam_remover.h"
#include "squash_expr_rewriter.h"
#include "squash_cmd.h"
#include <fstream>
#include <sstream>

static const std::string eod = ";{}";
static const std::string eop = ",)";
static const std::string sop = ",(";
static const std::string soop = "&*";
static const std::string eos = ";";  
static const std::string s_return = "return";  

OutparamRemover::OutparamRemover(SquashCmd const &cmd, Patcher &patcher):
  cmd(cmd),
  patcher(patcher),
  curr_opi(NULL)
{
  if (cmd.outparamFile.empty()) return;
  std::ifstream f(cmd.outparamFile.c_str());
  if (!f.is_open()) {
    cerr << "Could not open file:" << cmd.outparamFile << endl;
    exit(1);
  }
  std::string line;
  while (std::getline(f, line)) {
    OutParamInfo opi;
    std::string::size_type old_n, n;

    if (line.empty() || line[0] == '#') continue;

    // skip dehydra location info
    n = line.find(" ");
    if (n != std::string::npos) line.erase(0, n + 1);

    n = line.find("::");
    if(n == std::string::npos)
      goto error;
    opi.className = line.substr(0, n);
    n = line.find(',', old_n = n + 2);
    if(n == std::string::npos)
      goto error;
    opi.memberName = line.substr(old_n, n - old_n);
    n = line.find(',', old_n = n + 1);
    if(n == std::string::npos)
      goto error;
    std::string param = line.substr(old_n, n - old_n);
    n = param.find('=');
    if(n == std::string::npos)
      goto error;
    opi.assign = param.substr(n+1, param.size() - n - 1);
    opi.param = atoi(param.c_str());
    cerr << opi.className << "::"<<opi.memberName<<","<<opi.param << "=" 
	 << opi.assign << endl;
    ls.push_back(opi);
  }
  f.close();
  return;
 error:
  cerr << "Invalid input file:" << cmd.outparamFile << endl;
  exit(1);
}

// Based on a func with the same name in squash_expr_rewriter.cc
static bool matchesType(std::string const &type, Type *t) {
  if (t->isPtrOrRef() || t->isArrayType()) {
    return matchesType(type, t->getAtType());
  }
  if (!t->isCompoundType()) {
    return false;
  }
  CompoundType *ct = t->asCompoundType();
  TS_classSpec *c = ct->syntax;
  return type == c->name->getName();
}

template<class T> MinMaxLineVisitor
OutparamRemover::traverseWithMaxCol(T *obj, 
				    std::string const &chars,
				    int offset) const {
  MinMaxLineVisitor mmlv;
  obj->traverse(mmlv);
  mmlv.setIndent(mmlv.getIndent() + offset);
  std::string line = patcher.getLine(mmlv.maxLine, mmlv.getFile());
  std::string::size_type last = line.find_first_of(chars, mmlv.getIndent() + 1);
  mmlv.maxCol = (last != std::string::npos) ? last : 0;
  return mmlv;
}

void OutparamRemover::rewriteFunctionDeclaration(MR_decl *m,
						 OutParamInfo &opi, 
						 ASTTypeId *outparam) {
  Declarator *d = m->d->decllist->first();
  MinMaxLineVisitor mmlvDecl(traverseWithMaxCol(m, eod));

  MinMaxLineVisitor mmlvFunc;
  d->decl->getDeclaratorId()->traverse(mmlvFunc);

  rewriteFuncDeclOrDef(mmlvDecl, mmlvFunc, 
		       d->var->name,
		       outparam);
}

void OutparamRemover::postvisitMember(Member *mem) {
  MR_decl *m = mem->ifMR_decl();
  if (!m) return;
  Declarator *d = m->d->decllist->first();
  if (!d) return;
  D_func *df = d->decl->ifD_func();

  if (!df) return; 
  
  Variable *v = d->var;
  // these are from the type system
  if (!d->type->isFunctionType()) return;
  FunctionType *f = d->type->asFunctionType();
  SObjListIter<Variable> typedParams(f->params);
  ASTTypeId *outparam = NULL;
  Type *thisType = NULL;
  if (!typedParams.isDone()) thisType = typedParams.data()->type;
  OutParamInfo *opi = isMatch(df, "", v, outparam, thisType);
  if(opi && outparam)
    rewriteFunctionDeclaration(m, *opi, outparam);
}

OutparamRemover::OutParamInfo* OutparamRemover::isMatch(D_func *df, 
							std::string const &className, 
							Variable *member,
							ASTTypeId *&outparam,
							Type *type) {
  for (std::list<OutParamInfo>::iterator i = ls.begin();
       i != ls.end();
       ++i) {
    if (member->name != i->memberName) continue;
     
    // type of "this" parameter is the same as the class name
    if (! (className == i->className ||
	   (type && matchesType(i->className, type)))) {
      continue; 
    }
    unsigned int param = i->param;
    for (FakeList<ASTTypeId>* params = df->params;
	 ! params->isEmpty();
	 params = params->butFirst()) {
      ASTTypeId *p = params->first();
      if (param--) continue;
      outparam = p;
      break;
    }
    
    if (!outparam) continue;
    i->function = member;
    return &(*i);
    }
  return NULL;
}

void OutparamRemover::rewriteFunction(Function *f) {
  OutParamInfo* opi = NULL;
  // a goto-by any other name
  // start a rewrite if f is a candidate
  // otherwise rewrite contents of f if it calls other
  // rewritten funcs
  do {
    PQName  *pqname = f->nameAndParams->getDeclaratorId();
    PQ_qualifier *q = pqname->ifPQ_qualifier();
    if (!q) break;
    PQ_name *pq = q->rest->ifPQ_name();
    if (!pq) break;
    D_func *df = f->nameAndParams->decl->ifD_func();
    ASTTypeId *outparam = NULL;
    opi = isMatch(df, q->qualifier, f->nameAndParams->var, outparam);
    if (!opi) break;
    
    // Capture return type of the func
    MinMaxLineVisitor mmlvDecl(traverseWithMaxCol(f->nameAndParams, eod));
    f->retspec->traverse(mmlvDecl);
  
    MinMaxLineVisitor mmlvFunc;
    df->getDeclaratorId()->traverse(mmlvFunc);

    std::string rtype = rewriteFuncDeclOrDef(mmlvDecl, mmlvFunc, 
					     opi->className.empty() ? pq->name : q->qualifier,
					     outparam);
    std::stringstream ss;
  
    MinMaxLineVisitor mmlv;
    mmlv.setIndent(sourceLocManager->getCol(f->body->loc));
    mmlv.maxLine = mmlv.minLine = sourceLocManager->getLine(f->body->loc);
    std::string strFile = sourceLocManager->getFile(f->body->loc);
    std::string line = patcher.getLine(mmlv.maxLine, strFile);
    std::string::size_type pos = line.find('{');
    if (pos != std::string::npos) {
      ss << line.substr(0, pos + 1) << endl;
      line.erase(0, pos + 1);
    }
    opi->paramVar = outparam->decl->var;
    ss << "  " << rtype << ' ' << opi->paramVar->name <<  " = nsnull;" << endl
       << line;
    patcher.printPatch(ss.str(), mmlv, strFile);  
  } while(false);

  Restorer<OutParamInfo*> restorer(this->curr_opi, opi);
  f->body->traverse(this->loweredVisitor);
}

std::string OutparamRemover::rewriteFuncDeclOrDef(MinMaxLineVisitor mmlvDecl, 
						  MinMaxLineVisitor mmlvFunc,
						  std::string const &startOfFunc,
						  ASTTypeId *outparam) {
  mmlvFunc.maxCol = mmlvDecl.maxCol;
  mmlvFunc.maxLine = mmlvDecl.maxLine;
  int offset = mmlvFunc.getIndent();
  mmlvFunc.setIndent(0);
  
  // calculate offset to compensate for real position vs CPP-expanded position
  std::string str = patcher.getRange(mmlvFunc);
  mmlvFunc.setIndent(str.find(startOfFunc));
  offset = mmlvFunc.getIndent() - offset;
  
  MinMaxLineVisitor mmlvPrefix = mmlvDecl;
  mmlvPrefix.maxLine = mmlvPrefix.minLine;
  mmlvPrefix.setIndent(0);
  mmlvPrefix.maxCol = mmlvDecl.getIndent();
  
  MinMaxLineVisitor mmlvPostfix = mmlvDecl;
  mmlvPostfix.minLine = mmlvPostfix.maxLine;
  mmlvPostfix.setIndent(mmlvDecl.maxCol);
  mmlvPostfix.maxCol = 0;
  
  MinMaxLineVisitor mmlvParamType, mmlvParamName, mmlvParam;
  // get the type
  mmlvParam = traverseWithMaxCol(outparam, eop, offset);
  // steal pointer *& stuff from name
  mmlvParamName = traverseWithMaxCol(outparam->decl, eop, offset);
  std::string paramNameString = patcher.getRange(mmlvParamName);
  
  mmlvParamType = mmlvParam;
  mmlvParamType.maxCol =  mmlvParamName.getIndent() 
    + paramNameString.find_last_of(soop, paramNameString.find(outparam->decl->var->name));
  mmlvParamType.maxLine = mmlvParamName.minLine;
  
  std::string strFunc = patcher.getRange(mmlvFunc);
  // this is the laziest way ever to remove the param
  std::string strParam = patcher.getRange(mmlvParam);
  std::string::size_type start = strFunc.find(strParam),
    stop = start + strParam.size();
  start = strFunc.find_last_of(sop, start)+1;
  strFunc.erase(start, stop - start);
  
  if (strFunc[start] == ',') strFunc.erase(start, 1);
  else if (strFunc[start-1] == ',') strFunc.erase(start - 1, 1);
  
  // now reassemble the decl
  std::stringstream ss;
  if (mmlvDecl.getIndent()) ss << patcher.getRange(mmlvPrefix);
  
  std::string strParamType = patcher.getRange(mmlvParamType);
  ss << strParamType
     << (mmlvDecl.minLine == mmlvFunc.minLine ? ' ' : '\n')
     << strFunc;
  if (mmlvDecl.maxCol) ss << patcher.getRange(mmlvPostfix);
  
  ss << '\n';
  patcher.printPatch(ss.str(), mmlvDecl, mmlvDecl.getFile());
  return strParamType;
}

bool OutparamRemover::visitS_return(S_return *s) {
  if (!curr_opi) return true;
  MinMaxLineVisitor mmlv = traverseWithMaxCol(s, eos);
  std::stringstream ss;
  ss << s_return << ' ';
  E_intLit *e = s->expr->expr->ifE_intLit();
  if (!e || (e && !strcmp(e->text, "0"))) ss << curr_opi->paramVar->name;
  else ss << "nsnull";
  patcher.replaceWord(mmlv, patcher.getRange(mmlv), ss.str());
  return false;
}

bool OutparamRemover::visitE_deref(E_deref *e) {
  E_variable *v = e->ptr->ifE_variable();
  if (!curr_opi || !v || v->var != curr_opi->paramVar) return true;
  MinMaxLineVisitor mmlv;
  v->traverse(mmlv);
  if (mmlv.getIndent()) mmlv.setIndent(mmlv.getIndent() - 1);
  //std::string line = patcher.getLine(mmlv.minLine, mmlv.getFile());
  //line.erase(line.rfind('*', mmlv.getIndent()), 1);
  //line.append("\n");
  //patcher.printPatch(line, mmlv, mmlv.getFile());
  std::string strDeref = v->var->name;
  strDeref.insert(0, "*");
  patcher.replaceWord(mmlv, strDeref, v->var->name); 
  return false;
}

void OutparamRemover::postvisitE_funCall(E_funCall *e) {
  E_fieldAcc *f = e->func->ifE_fieldAcc();
  if (!f || !f->field) return;
  OutParamInfo *opi = NULL;
  for(std::list<OutParamInfo>::iterator it = ls.begin(); 
      it != ls.end();
      ++it) {
    if (it->function != f->field) continue;
    opi = &(*it);
  }
  if (!opi) return;
  unsigned int counter = 0;
  Expression *paramExpr = NULL;
  for (FakeList<ArgExpression>* args = e->args;
	 ! args->isEmpty();
	 args = args->butFirst()) {
    if (counter != opi->param) {
      counter++;
      continue;
    }
    paramExpr = args->first()->expr;
    break;
  }
  E_addrOf *addrOfExpr = paramExpr->ifE_addrOf();
  E_variable *paramVarExpr = (addrOfExpr ? addrOfExpr->expr : paramExpr)->asE_variable();
  
  MinMaxLineVisitor mmlvOutparam = traverseWithMaxCol(paramExpr, eop);
  // remove the param
  std::string line = patcher.getLine(mmlvOutparam.minLine, mmlvOutparam.getFile());
  mmlvOutparam.setIndent(line.find_last_of(sop, mmlvOutparam.getIndent()) + 1);
  std::string param = patcher.getRange(mmlvOutparam);
  unsigned int start = mmlvOutparam.getIndent();
  if (line[mmlvOutparam.maxCol + 1] == ',') mmlvOutparam.maxCol++;
  else if (start && line[start - 1] == ',') mmlvOutparam.setIndent(start - 1);

  MinMaxLineVisitor mmlvFunc;
  f->traverse(mmlvFunc);
  line = patcher.getLine(mmlvFunc.maxLine, mmlvFunc.getFile());
  mmlvFunc.maxCol = line.rfind(opi->memberName) + opi->memberName.size();
  
  std::stringstream ss;
  std::string func = patcher.getRange(mmlvFunc);
  if (!addrOfExpr) ss << '*';
  ss << paramVarExpr->var->name << " = " << func;
  patcher.replaceWord(mmlvFunc, func, ss.str());
  
  
  patcher.replaceWord(mmlvOutparam, patcher.getRange(mmlvOutparam), "");


}
