#include "SharedLibraryManipulation.hpp"
#include "b64/decode.h"

#include <llvm/Linker.h>
#include <llvm/Analysis/Verifier.h>
#include <llvm/Bitcode/ReaderWriter.h>
#include <llvm/Support/Debug.h>
#include <llvm/Support/MemoryBuffer.h>
#include <llvm/System/Path.h>
#include <llvm/System/Program.h>

#include <boost/assign/list_of.hpp>
#include <boost/format.hpp>
#include <boost/iostreams/concepts.hpp>
#include <boost/iostreams/stream.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/regex.hpp>
#include <boost/tokenizer.hpp>

#include <algorithm>
#include <fstream>
#include <iosfwd>
#include <iostream>
#include <iterator>
#include <sstream>
#include <vector>

using namespace boost;
using namespace boost::assign;
using namespace llvm;

static std::string readContentsOfFile(const sys::Path & filename)
{
  std::ifstream ifs(filename.c_str(), std::ios::binary);
  ifs.seekg(0, std::ios::end);
  size_t length = ifs.tellg();
  ifs.seekg(0, std::ios::beg);

  // Allocate an extra byte for a null terminator
  std::vector<char> buffer(length+1);
  ifs.read(&buffer[0], length);
  ifs.close();

  // Make sure the buffer is null terminated
  buffer[length] = 0;

  // Implicit string conversion
  return &buffer[0];
}

/*!
  Execute a program with arguments and the supplied redirects.

  redirects is a three element array:
  0: stdin
  1: stdout
  2: stderr

  redirects or any of its entries may be NULL.

  The return value is the exit code of the program.

  The first argument should be the program name.
 */
static int executeWithRedirects(const std::vector<std::string> & args, const sys::Path** redirects)
{
  assert(args.size() > 0);
  sys::Path prog = sys::Program::FindProgramByName(args[0]);

  std::vector<const char*> raw_args;
  for(std::vector<std::string>::const_iterator it = args.begin(),
        ed = args.end(); it != ed; ++it)
  {
    char* buffer = new char[it->size()+1];
    std::strncpy(buffer, it->c_str(), it->size()+1);
    raw_args.push_back(buffer);
  }

  // NULL terminate the arguments
  raw_args.push_back(NULL);

  int retval = sys::Program::ExecuteAndWait(prog, &raw_args[0], NULL, redirects);
  for(std::vector<const char*>::iterator it = raw_args.begin(),
        ed = raw_args.end(); it != ed; ++it)
  {
    delete [] *it;
  }
  
  return retval;
}

/*!
  Executes a program with arguments and returns a string containing the output
  from stdout.

  The first argument should be the program name
 */
static std::string executeCapturingSTDOUT(const std::vector<std::string> & args)
{
  sys::Path tmpFile;
  tmpFile.createTemporaryFileOnDisk();

  // const sys::Path* redirects[] = {NULL, &tmpFile, NULL};
  std::vector<const sys::Path*> redirects = list_of((sys::Path*)NULL)(&tmpFile)((sys::Path*)NULL);

  executeWithRedirects(args, &redirects[0]);

  std::string output(readContentsOfFile(tmpFile));
  tmpFile.eraseFromDisk();
  return output;
}

static size_t hexStringToNumber(const std::string & hex)
{
  std::istringstream is(hex);
  size_t ret;
  is >> std::hex >> ret;

  return ret;
}

/*!
  We match the output of objdump -h -w to find the line with the desired section and
  extract its offset and size.  Then we just open the file and read out exactly the
  bytes we want.
*/
static std::string extractELFSection(const std::string & file, std::string section)
{
  std::vector<std::string> objdumpCommand = list_of("objdump")("-h")("-w")(file.c_str());
  std::string objdumpOutput = executeCapturingSTDOUT(objdumpCommand);
  
  // # section size VMA LMA offset
  // The addresses are all hex fields, so use \w instead of \d
  // The section name begins with a dot, which will match itself, but it would be best to escape
  // it anyway
  if(section[0] == '.')
    section = "\\" + section;
  
  std::string rx = (format("^[[:space:]]+[[:digit:]]+[[:space:]]+%s[[:space:]]+([[:xdigit:]]+)[[:space:]]+[[:xdigit:]]+[[:space:]]+[[:xdigit:]]+[[:space:]]+([[:xdigit:]]+)[[:space:]]+") % section).str();
  regex sectionExtractor(rx);
  match_results<std::string::const_iterator> what;
  
  if(!boost::regex_search(objdumpOutput, what, sectionExtractor))
    throw "Missing ELF section " + section + " in " + file;

  std::string section_size_str = what[1];
  std::string section_offset_str = what[2];

  size_t ssize = hexStringToNumber(section_size_str);
  size_t soff = hexStringToNumber(section_offset_str);

  // Now read the section directly from the file using this size and offset
  // Allocate a buffer with space for a null terminator
  std::vector<char> buffer(ssize+1);
  std::ifstream ifs(file.c_str(), std::fstream::binary);
  ifs.seekg(soff, std::ios::beg);
  ifs.read(&buffer[0], ssize);
  // NULL terminator
  buffer[ssize] = NULL;
  ifs.close();
  
  return &buffer[0];
}

/*!
  Accepts an ELF section of base64 encoded bitcode files (one per compilation unit)
  and reconsitutes them into a single bitcode file in a temporary file on disk with
  the llvm-link utility.

  Sections are base64 encoded and enclosed between % characters and terminated with
  a newline: %b64data%\n
 */
static Module* linkPerCompilationUnitBitcodeFiles(const std::string & bitcodeSections)
{
  // Perform linking with the LLVM linker
  Linker L("wholeLibrary", "wholeLibrary");
  
  // Separate the sections on %...%\n boundaries
  typedef tokenizer<char_separator<char> > bitcode_tokenizer;
  char_separator<char> sep("%");
  bitcode_tokenizer tokens(bitcodeSections, sep);
  
  // For each section, base64 decode it into a memory buffer
  // and link it into a composite library
  for(bitcode_tokenizer::iterator tok_it = tokens.begin(),
        tok_ed = tokens.end(); tok_it != tok_ed; ++tok_it)
  {
    const std::string & token = *tok_it;

    // Not a real section
    if(token.length() < 3)
      continue;
    
    std::istringstream is(token);
    std::vector<char> buffer;
    iostreams::filtering_ostream os(std::back_inserter(buffer));
    base64::decoder D;
    
    // Decode the base64 string (in the stringstream) into a
    // temporary stringstream
    D.decode(is, os);

    os.flush();
    
    std::ofstream fout("/tmp/test.bc");
    fout.write(&buffer[0], buffer.size());
    fout.close();
    
    MemoryBuffer* mem_buf =
      MemoryBuffer::getMemBufferCopy(&buffer[0], &buffer[0] + buffer.size());
    
    Module* M = ParseBitcodeFile(mem_buf);
    assert(M && "Bad bitcode");
    L.LinkInModule(M); // Do we need to free M after this? The linker probably
                       // takes ownership

    delete M;
  }
  
  return L.releaseModule();
}

std::auto_ptr<Module> assembleBitcodeFile(const std::string & sharedLib)
{
  std::string rawBitcodeSections = extractELFSection(sharedLib, ".llvm_bc");
  std::auto_ptr<Module> M(linkPerCompilationUnitBitcodeFiles(rawBitcodeSections));
  // M->dump();
  // llvm::verifyModule(*M.get());
  return M;
}

std::string extractPartialInterfaceDescription(const std::string & sharedLib)
{
  return extractELFSection(sharedLib, ".llvm_pid");
}

std::vector<std::string> findImplicitDependencies(const std::string & sharedLib)
{
  // libm is included in libc on most systems... need to be careful with this
  // linux-gate is a virtual shared object
  // /lib/ld-linux.so is common, but the regex should ignore that since it
  //  begins with a /
  std::set<std::string> ignoredDependencies =
    list_of("m")("linux-gate").to_container(ignoredDependencies);
  
  std::vector<std::string> deps;
  std::vector<std::string> commandLine = list_of("ldd")(sharedLib.c_str());
  std::string output = executeCapturingSTDOUT(commandLine);

  regex rx("^\\s*lib(.+)\\.so+.*=>");
  
  // Break the output into lines and match against a regex
  typedef tokenizer<char_separator<char> > line_tokenizer;
  char_separator<char> sep("\n");
  line_tokenizer tokens(output, sep);

  for(line_tokenizer::iterator tok_it = tokens.begin(),
        tok_ed = tokens.end(); tok_it != tok_ed; ++tok_it)
  {
    match_results<std::string::const_iterator> what;
    if(!regex_search(*tok_it, what, rx))
      continue;

    std::string libname = what[1];
    if(ignoredDependencies.find(libname) != ignoredDependencies.end())
      continue;
    deps.push_back("lib" + libname + ".id");
  }

  return deps;
}
