#include <cassert>
#include <fstream>
#include <iomanip>
#include <iostream>
#include "AST.h"
#include "ClassDecl.h"
#include "FieldDecl.h"
#include "Substitution.h"
#include "compiler.h"
#include "decls.h"
#include "template.h"
#include "utils.h"

// This file implements the "pseudo-template hack" used for 
// typechecking of Point<N>, RectDomain<N>, Domain<N>, TiArray<N> and JavaArray
// It has nothing to do with "real" templates, and will hopefully 
// eventually be replaced with real templates (whose functionality subsume it)

// Instantiated template environment

TemplateEnv templateEnv;

Decl *TemplateEnv::lookup(const TypeNode *type)
{
  foreach (instance, llist<instance *>, *instances)
    if ((*instance)->type->typeIdentNM(type)) return (*instance)->d;

  return NULL;
}

void TemplateEnv::add(const TypeNode *type, Decl *inst)
{
  instance *i = new instance;

  i->type = type->deepClone();
  i->d = inst;
  instances = cons(i, instances);
}

// Template instantiation

static TreeNode *substitute(TreeNode *from, Subst *args)
{
  Decl *fromd = from->decl();

  foreach (sub, Subst, *args)
    if ((*sub)->d == fromd) return (*sub)->with;

  return from;
}

Decl *Decl::forwardto(TreeNode *)
{
  invalidOperation("forwardto");
  return NULL;
}

Decl *LocalVarDecl::forwardto(TreeNode *src)
{
  forward(new LocalVarDecl(name(), src->dtype(), src, assignable()));

  return forward();
}

Decl *FormalParameterDecl::forwardto(TreeNode *src)
{
  forward(new FormalParameterDecl(name(), src->dtype(), src, assignable()));

  return forward();
}

Decl *StmtLblDecl::forwardto(TreeNode *src)
{
  forward(new StmtLblDecl(name(), src));

  return forward();
}

// memberforward is used for early substitution (before the source is
// copied), for name-resolution & type-checking of uses of the
// instantiated template.
Decl *Decl::memberforward(Subst *)
{
  invalidOperation("memberforward");
  return NULL;
}

FieldDecl *FieldDecl::memberforward(Subst *args)
{
  ClassDecl *newContainer = static_cast< ClassDecl * >(container()->forward());
  assert( newContainer->category() & (Decl::Class | Decl::Interface) );
  
  return new FieldDecl(name(), (TypeNode *)type()->instantiate(args),
		       newContainer, modifiers(), source());
}

MethodDecl *MethodDecl::memberforward(Subst *args)
{
  const string *newName = name();

  ClassDecl *newContainer = static_cast< ClassDecl * >(container()->forward());
  assert( newContainer->category() & (Decl::Class | Decl::Interface) );

  if (category() == Constructor)
    newName = newContainer->name();

  return new MethodDecl(newName, (TypeNode *)type()->instantiate(args),
			category(), newContainer, modifiers(), source());
}

TreeNode *TreeNode::instantiate(Subst *args)
{
  TreeNode *me = clone();

  const int numChildren = arity();
  for (int sweep = 0; sweep < numChildren; ++sweep)
    me->child(sweep, me->child(sweep)->instantiate(args));

  return me;
}

TreeNode *TypeNameNode::instantiate(Subst *args)
{
  TypeNode *newType = (TypeNode *)substitute(this, args);
  Modifiers mods = (Modifiers) ((newType->modifiers() & ~Single) | modifiers());
  return newType->withModifiers(mods);
}

// To support 'arity' instantiation
TreeNode *TypeFieldAccessNode::instantiate(Subst *args)
{
  return substitute(this, args);
}

TreeNode *NameNode::instantiate(Subst *args)
{
  NameNode *me = clone();

  /* Forward the decl's of name uses */
  me->decl(decl()->forward());

  return me;
}

TreeNode *VarDeclNode::instantiate(Subst *args)
{
  VarDeclNode *me = clone();

  /* Copy the decl's of declarations */
  me->simpName()->decl(decl()->forwardto(this));
  me->dtype((TypeNode *)me->dtype()->instantiate(args));
  me->initExpr(me->initExpr()->instantiate(args));

  return me;
}

TreeNode *ParameterNode::instantiate(Subst *args)
{
  ParameterNode *me = clone();

  /* another declaration */
  me->simpName()->decl(decl()->forwardto(this));
  me->dtype((TypeNode *)me->dtype()->instantiate(args));

  return me;
}

TreeNode *LabeledStmtNode::instantiate(Subst *args)
{
  LabeledStmtNode *me = clone();

  /* another declaration */
  me->label()->decl(decl()->forwardto(this));
  me->stmt(me->stmt()->instantiate(args));

  return me;
}

/* Member & field decl's are forwarded earlier (see memberforward), so
   only need to fix source pointer */
TreeNode *FieldDeclNode::instantiate(Subst *args)
{
  TreeNode *me = TreeNode::instantiate(args);

  me->simpName()->decl()->source(this);

  return me;
}

TreeNode *MethodNode::instantiate(Subst *args)
{
  TreeNode *me = TreeNode::instantiate(args);

  me->simpName()->decl()->source(this);

  return me;
}

TreeNode *ConstructorDeclNode::instantiate(Subst *args)
{
  TreeNode *me = TreeNode::instantiate(args);

  me->simpName()->decl()->source(this);

  return me;
}


TreeNode *ExprNode::instantiate(Subst *args)
{
  ExprNode *me = (ExprNode *)TreeNode::instantiate(args);

  /* It would be nice to avoid the copy when possible and/or to do
     some sharing of copied types */
  /*me->_type = (TypeNode *)type()->instantiate(args);*/
  me->_type = NULL; // force type recomputation - should work

  return me;
}

/* Template handling is divided into 2 phases:

1) When the template instance is seen, we create a Decl for the template
   instance with the correct name and a correct environment

2) Once all templates are type-checked, widened, etc, we take all
   template instances and copy the template code with appropriate
   substitutions. The copies of the code have correct types and
   copies of all the Decl's.

TBD for real templates:
     - where clauses need decl substitution
     - template instances used by tmplate are special if they use some of the
       template arguments
     - other issues I'm sure...
   */

ClassDecl *instantiateClass(ClassDecl *tmplate,
			    const TypeNode *type,
			    TemplateEnv *templateEnv, 
			    Subst *args,
			    TypeNode* classType)
{
  /* Create a decl for a template instance: 
     - a) make a name for the new class
     - b) make a decl for this class
     - c) substitute all decls in the ClassDecl
          (fields, methods, super class, and (TBD) super interfaces)
     - d) save the substitution
  */

  /* Make the instance name */
  string *instanceName = new string(*(tmplate->name()));
  bool first = true;

  *instanceName += "<";
  foriter (arg, type->allChildren(), TreeNode::ConstChildIter)
    {
      if (!first) *instanceName += ", ";
      first = false;
      *instanceName += (*arg)->typeName();
    }
  *instanceName += ">";

  if (DEBUG)
    cout << "instantiateClass for " << *instanceName << '\n';

  /* Make the decl representing the instance, add it to the template mapping
     (do the add now to avoid infinite recursion) */
  ClassDecl *inst = new ClassDecl(instanceName, tmplate->container(),
				  tmplate->modifiers(),
				  tmplate->source(),
				  classType);
  tmplate->forward(inst);
  templateEnv->add(type, inst);
  Environ *newEnv = new Environ();
  inst->environ(newEnv);

  /* Fill in inst decl */
  if (tmplate->superClass())
    inst->superClass(static_cast<ClassDecl *>(tmplate->superClass()->forward()));

  llist<Decl *> *ifs = NULL;
  foreach (d, llist<Decl *>, *tmplate->interfaces())
    ifs = cons((*d)->forward(), ifs);
  inst->interfaces(ifs);

  foriter (member, tmplate->environ()->allProperDecls(), EnvironIter)
    newEnv->add(member->container() == tmplate
		? member->memberforward(args)
		: &*member);
	

  inst->templateArgs(args);
  
  inst->visits = 3; // all done.

  return inst;
}

llist<TreeNode *> *TemplateEnv::allInstances()
{
  llist<TreeNode *> *instsrc = NULL;

  foreach (i, llist<instance *>, *instances)
    instsrc = cons((*i)->d->source(), instsrc);

  return instsrc;
}

/* Fill in decl's and types in types instantiated from templates
   based on the decls and types in the actual template.
   (Can't do this when the template is instantiated as we haven't
   type-checked the template then)
   */
void TemplateEnv::resolveTemplates(void)
{
  foreach (i, llist<instance *>, *instances)
    {
      Decl *inst = (*i)->d;

      if (DEBUG)
	cout << "Instantiating " << *inst->name() << "\n";

      /* Make all field & method decl's in the template source point
	 to instance decls (the line in the iterator is a noop for
	 decls that don't come from the template) */
      foriter (member, inst->environ()->allProperDecls(), EnvironIter)
	member->source()->decl()->forward(&*member);

      // Same for the class decl for this template
      inst->source()->decl()->forward(inst);

      /* Copy source, decls & types */
      inst->source(inst->source()->instantiate(inst->templateArgs()));
    }
}