#include "ArgumentHelpers.hpp"
#include "Results.hpp"
#include "Scheme.hpp"
#include "SchemeDepPaths.hpp"

#include <libguile.h>

#include <llvm/Module.h>
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/Debug.h>
#include <llvm/System/Path.h>

#include <boost/bind.hpp>
#include <boost/format.hpp>

using boost::format;
using namespace llvm;

struct guile_box
{
  results_t * res;
  argument_t arg;
};

static scm_t_bits box_tag;

static SCM mark_box(SCM smob)
{
  return SCM_BOOL_F;
}

static size_t free_box(SCM smob)
{
  guile_box * box = (guile_box*)SCM_SMOB_DATA(smob);
  scm_gc_free(box, sizeof(guile_box), "box");
  return 0;
}

static int print_box(SCM smob, SCM port, scm_print_state * pstate)
{
  scm_puts("#<SignatureDumpBox>", port);
  return 1;
}

static SCM makeBox(results_t * res, argument_t arg)
{
  SCM smob;

  guile_box * box = (guile_box*) scm_gc_malloc(sizeof(guile_box), "box");
  box->res = res;
  box->arg = arg;

  SCM_NEWSMOB(smob, box_tag, box);

  return smob;
}

static SCM getFunction(const std::string & funcname)
{
  SCM sym = scm_c_lookup(funcname.c_str());
  return scm_variable_ref(sym);
}

static SCM registerAllocator(SCM box)
{
  guile_box * b = (guile_box*)SCM_SMOB_DATA(box);
  const Function * fn = b->arg._fn;
  results_t * res = b->res;

  res->markAllocator(fn);

  DOUT << format("Registering [%s] as an allocator\n")
    % fn->getName();

  return SCM_BOOL_F;
}

static SCM registerNotDeallocator(SCM box)
{
  guile_box * b = (guile_box*)SCM_SMOB_DATA(box);
  const Function * fn = b->arg._fn;
  const Argument * arg = b->arg._arg;
  results_t * res = b->res;

  res->markNotDeallocator(fn, arg);
  
  return SCM_BOOL_F;
}

static SCM registerArray(SCM box, SCM depth)
{
  guile_box * b = (guile_box*)SCM_SMOB_DATA(box);
  const Function * fn = b->arg._fn;
  const Argument * arg = b->arg._arg;
  results_t * res = b->res;

  size_t idepth = scm_to_size_t(depth);
  
  res->markArray(fn, arg, idepth);
  
  return SCM_BOOL_F;
}

static SCM registerEscapeStrong(SCM box)
{
  guile_box * b = (guile_box*)SCM_SMOB_DATA(box);
  const Function * fn = b->arg._fn;
  const Argument * arg = b->arg._arg;
  results_t * res = b->res;

  res->markStrongEscape(fn, arg);
  
  return SCM_BOOL_F;
}

static SCM registerEscapeWeak(SCM box)
{
  guile_box * b = (guile_box*)SCM_SMOB_DATA(box);
  const Function * fn = b->arg._fn;
  const Argument * arg = b->arg._arg;
  results_t * res = b->res;

  res->markWeakEscape(fn, arg);
  
  return SCM_BOOL_F;
}

static SCM registerIn(SCM box)
{
  guile_box * b = (guile_box*)SCM_SMOB_DATA(box);
  const Function * fn = b->arg._fn;
  const Argument * arg = b->arg._arg;
  results_t * res = b->res;

  res->markIn(fn, arg);
  
  return SCM_BOOL_F;
}

static SCM registerOut(SCM box)
{
  guile_box * b = (guile_box*)SCM_SMOB_DATA(box);
  const Function * fn = b->arg._fn;
  const Argument * arg = b->arg._arg;
  results_t * res = b->res;

  res->markOut(fn, arg);
  
  return SCM_BOOL_F;
}

static void annotateAllocator(const Function * fn)
{
  // Don't register e.g. malloc as an allocator, since
  // we only have a declaration, not a full definition.
  if(fn->isDeclaration())
    return;
  
  SCM func = getFunction("partial-mark-allocator");
  SCM name = scm_from_locale_string(fn->getName().c_str());
  scm_call_1(func, name);
}

static void annotateArrayArgument(results_t * res,
                                  std::pair<argument_t, short> p)
{
  if(p.first._fn->isDeclaration())
    return;

  // arg will be NULL for the varg parameter of
  // a function.  Make sure we don't segfault
  // while just getting its name.
  const Argument * arg = p.first._arg;
  
  // We can just bail out of we hit a vararg parameter; we don't
  // need to set any properties for those, since we acknowledge that
  // inferring anything is impractical.
  if(!arg)
    return;

  SCM func = getFunction("partial-mark-array");
  
  std::string name = p.first._fn->getName();
  SCM sName = scm_from_locale_string(name.c_str());
  
  SCM argno = scm_from_uint32(realArgno(arg));
  SCM depth = scm_from_short(p.second);

  scm_call_3(func, sName, argno, depth);  
}

static void annotateArgument(results_t * res,
                             std::pair<argument_t, state_t> p)
{
  cerr << format("Annotating: %s\n") % p.first._fn->getName();
  if(p.first._fn->isDeclaration())
    return;

  // arg will be NULL for the varg parameter of
  // a function.  Make sure we don't segfault
  // while just getting its name.
  const Argument * arg = p.first._arg;

  // We can just bail out of we hit a vararg parameter; we don't
  // need to set any properties for those, since we acknowledge that
  // inferring anything is impractical.
  if(!arg)
    return;

  cerr << format("  Argument %s\n") % arg->getName();
  if(arg->getName().find(".") != std::string::npos)
    return;
  
  std::string name = p.first._fn->getName();
  SCM sName = scm_from_locale_string(name.c_str());
  SCM argno = scm_from_uint32(realArgno(arg));
  
  if(res->isIn(p.first._fn, p.first._arg))
  {
    SCM func = getFunction("partial-mark-in");
    scm_call_2(func, sName, argno);
  }
  
  if(res->isOut(p.first._fn, p.first._arg))
  {
    SCM func = getFunction("partial-mark-out");
    scm_call_2(func, sName, argno);
  }
  
  if(res->argumentEscapesStrong(p.first._fn, p.first._arg))
  {
    SCM func = getFunction("partial-mark-escape-strong");
    scm_call_2(func, sName, argno);
  }

  if(res->argumentEscapesWeak(p.first._fn, p.first._arg))
  {
    SCM func = getFunction("partial-mark-escape-weak");
    scm_call_2(func, sName, argno);
  }
  
  if(res->deallocatesArgument(p.first._fn, p.first._arg))
  {
    SCM func = getFunction("partial-mark-deallocate");
    scm_call_2(func, sName, argno);
  }
}

static void removeFromOutput(const Function * fn)
{
  SCM func = getFunction("remove-function-from-output");
  SCM sName = scm_from_locale_string(fn->getName().c_str());
  scm_call_1(func, sName);
}

static void renameFunction(const std::pair<const Function*,
                           std::string> & alias)
{
  SCM func = getFunction("apply-alias");
  SCM originalName =
    scm_from_locale_string(alias.first->getName().c_str());
  SCM newName = scm_from_locale_string(alias.second.c_str());

  scm_call_2(func, originalName, newName);
}

/*
  This is really a part of the results_t class, but is implemented
  here so that the guile dependnecy can stay localized.
 */
void results_t::annotateSignature() const
{
  // For each function that is hidden, call into Guile to remove it
  // from the output
  std::for_each(_hidden_functions.begin(),
                _hidden_functions.end(),
                removeFromOutput);

  // Apply aliases that we discovered
  std::for_each(_aliases.begin(),
                _aliases.end(),
                renameFunction);

  std::for_each(_arrays.begin(),
                _arrays.end(),
                bind(annotateArrayArgument, const_cast<results_t*>(this), _1));
  
  // Apply each relevant annotation that we have discovered to the
  // basic annotations already generated from Clang.
  std::for_each(_results.begin(),
                _results.end(),
                bind(annotateArgument, const_cast<results_t*>(this), _1));

  // Add the allocator annotation _LAST_ so that the function already
  // has all of its other annotations applied.  This ensures that we
  // don't have to update the function in (potentially) two places.
  std::for_each(_allocators.begin(), _allocators.end(),
                annotateAllocator);
}

void guileInit(const std::string & libDir, const std::string & inputSExps)
{
  scm_init_guile();

  // Register the SMOB we use to pass some C data to Scheme
  // with the runtime.  This is just an opaque object; it
  // is never manipulated in Scheme.
  box_tag = scm_make_smob_type("box", sizeof(guile_box));
  scm_set_smob_mark(box_tag, mark_box);
  scm_set_smob_free(box_tag, free_box);
  scm_set_smob_print(box_tag, print_box);

  // Register some callbacks that the Scheme code will use to
  // fill in dependency information as we request it.
  scm_c_define_gsubr("register-allocator", 1, 0, 0, (SCM(*)())registerAllocator);
  scm_c_define_gsubr("register-not-deallocator", 1, 0, 0, (SCM(*)())registerNotDeallocator);
  scm_c_define_gsubr("register-array", 2, 0, 0, (SCM(*)())registerArray);
  scm_c_define_gsubr("register-escape-strong", 1, 0, 0, (SCM(*)())registerEscapeStrong);
  scm_c_define_gsubr("register-escape-weak", 1, 0, 0, (SCM(*)())registerEscapeWeak);
  scm_c_define_gsubr("register-in", 1, 0, 0, (SCM(*)())registerIn);
  scm_c_define_gsubr("register-out", 1, 0, 0, (SCM(*)())registerOut);

  // Add our local module directory to the Guile load path (note the escaped %s)
  std::string libpath = (format("(set! %%load-path (cons \"%s\" %%load-path))") % libDir).str();
  SCM libpathEvalString = scm_from_locale_string(libpath.c_str());
  scm_eval_string(libpathEvalString);
  
  // Load the remaining functions as Scheme code.
  std::string loaderFile = (format("%s/loader.scm") % libDir).str();
  scm_c_primitive_load(loaderFile.c_str());

  // Load in all of the partial definitions that we will be annotating
  if(!inputSExps.empty())
  {
    SCM proc = getFunction("load-ffi-fragments");
    SCM inputString = scm_from_locale_string(inputSExps.c_str());
    scm_call_1(proc, inputString);
  }
}

void guileFetchDependency(results_t * res, argument_t & arg)
{
  // Call out to Guile with the res pointer and arg stuffed into a smob
  // and with the function name.  Guile will do the lookup on
  // that function and call back to some appropriate C functions
  // that will use the res pointer as needed.
  std::string function_name = arg._fn->getName();
  const Argument * a = arg._arg;

  // Nothing to say about varargs yet.
  if(!a)
    return;
  
  DOUT << format("Requesting argument information for [%s] from Scheme\n") %
    function_name;
  
  SCM fn = scm_from_locale_string(function_name.c_str());
  SCM box = makeBox(res, arg);

  SCM argno = scm_from_uint(a->getArgNo());
  SCM proc = getFunction("lookup-dependency");

  scm_call_3(proc, box, fn, argno);
}

void guileFetchAllocator(results_t * res, argument_t & arg)
{
  SCM proc = getFunction("lookup-allocator-dependency");
  std::string function_name = arg._fn->getName();

  DOUT << format("Requesting allocator information for [%s] from Scheme\n") %
    function_name;
  
  SCM fn = scm_from_locale_string(function_name.c_str());
  SCM box = makeBox(res, arg);

  scm_call_2(proc, box, fn);
}

void guileLoadDep(const std::string & loadPath, const std::string & dep)
{
  format fullPath = format("%s/%s") % loadPath % dep;
  sys::Path file(fullPath.str());

  if(!file.exists())
  {
    cerr << format("[warning] Dependency <%s> does not exist and will not be loaded\n") % dep;
    return;
  }
  
  DOUT << format("Loading dependencies from %s\n") % dep;
  
  SCM proc = getFunction("load-ffi-def");
  SCM fn = scm_from_locale_string(fullPath.str().c_str());
  scm_call_1(proc, fn);
}

void guileLoadAnnotation(const std::string & annot)
{
  DOUT << format("Reading annotations from %s\n") % annot;
  SCM proc = getFunction("load-escape-annotations");
  SCM filename = scm_from_locale_string(annot.c_str());
  scm_call_1(proc, filename);
}

void guileInternAnnotationData(Module * M, results_t & res)
{
  SCM proc;

  proc = getFunction("get-weak-fields");
  SCM weak_fields = scm_call_0(proc);

  while(scm_is_pair(weak_fields))
  {
    SCM item = scm_car(weak_fields);
    SCM name = scm_car(item);
    SCM index = scm_cdr(item);

    char* tmp_string = scm_to_locale_string(name);
    std::string sname(tmp_string);
    free(tmp_string);

    short fieldno = scm_to_short(index);
    res.registerWeakField(sname, fieldno);
    cerr << format("Registering Weak Field: %s(%d)\n")
      % sname % fieldno;
    
    weak_fields = scm_cdr(weak_fields);
  }

  proc = getFunction("get-allocator-pairs");
  SCM allocator_pairs = scm_call_0(proc);

  while(scm_is_pair(allocator_pairs))
  {
    SCM item = scm_car(allocator_pairs);
    SCM allocator = scm_car(item);
    SCM finalizer = scm_cdr(item);

    char* tmp_string = scm_to_locale_string(allocator);
    std::string allocator_name(tmp_string);
    free(tmp_string);

    tmp_string = scm_to_locale_string(finalizer);
    std::string finalizer_name(tmp_string);
    free(tmp_string);

    res.registerFinalizer(allocator_name, finalizer_name);
    cerr << format("Registering Allocator/Finalizer Pair: %s/%s\n")
      % allocator_name % finalizer_name;
    
    allocator_pairs = scm_cdr(allocator_pairs);    
  }

  proc = getFunction("get-complex-finalizers");
  SCM complex_finalizers = scm_call_0(proc);

  while(scm_is_pair(complex_finalizers))
  {
    SCM item = scm_car(complex_finalizers);
    SCM finalizer = scm_car(item);
    SCM index = scm_cdr(item);

    char * tmp_string = scm_to_locale_string(finalizer);
    std::string finalizer_name(tmp_string);
    free(tmp_string);

    short argno = scm_to_short(index);

    res.pinDeallocator(finalizer_name, argno);

    complex_finalizers = scm_cdr(complex_finalizers);
  }
}

void guileWriteOutput(const std::string & OutputFilename)
{
  SCM outPort = scm_current_output_port();

  if(!OutputFilename.empty())
  {
    SCM mode = scm_from_locale_string("w");
    SCM ofname = scm_from_locale_string(OutputFilename.c_str());
    
    outPort = scm_open_file(ofname, mode);
  }
    
  SCM writeFunc = getFunction("write-annotated-defs");
  scm_call_1(writeFunc, outPort);
}
