/*
 * Copyright (c) 2014 Javier Escalada Gómez
 * All rights reserved
 *
 * Based on: https://github.com/eliben/llvm-clang-samples/blob/master/src_clang/rewritersample.cpp
 */

#include <vector>
#include <map>
#include <fstream>
#include <sstream>

#include <clang/AST/ASTConsumer.h>
#include <clang/AST/RecursiveASTVisitor.h>
#include <clang/Basic/Diagnostic.h>
#include <clang/Basic/FileManager.h>
#include <clang/Basic/SourceManager.h>
#include <clang/Basic/TargetOptions.h>
#include <clang/Basic/TargetInfo.h>
#include <clang/Frontend/CompilerInstance.h>
#include <clang/Lex/Preprocessor.h>
#include <clang/Parse/ParseAST.h>
#include <clang/Rewrite/Core/Rewriter.h>
#include <clang/Rewrite/Frontend/Rewriters.h>
#include <clang/Frontend/TextDiagnosticPrinter.h>
#include <llvm/Support/Host.h>
#include <llvm/Support/raw_ostream.h>

#include "getopt.h" 

using namespace clang;
using namespace std;

char *prefix = "";
char *suffix = "";
string current_file_name;
map<string, string> function_return_type;
map<string, pair<bool, string>> function_return_info;
bool skip_this_function = false;
vector<string> include_paths;
vector<string> defines;
vector<string> undefines;
string output_file = "return_info.json";

/******************************************************************************
 * Log code
 ***/
void inline _log(string level, raw_ostream::Colors color, bool bold, string txt, string *location=nullptr)
{
	assert(!level.empty());
	assert(!txt.empty());

	if (location) {
		llvm::errs().changeColor(raw_ostream::WHITE, true);
		llvm::errs() << *location << ": ";
	} else if (!current_file_name.empty()) {
		llvm::errs().changeColor(raw_ostream::WHITE, true);
		llvm::errs() << current_file_name << ": ";
	}
	llvm::errs().changeColor(color, bold);
	llvm::errs() << level <<": ";
	llvm::errs().changeColor(raw_ostream::WHITE, true);
	llvm::errs() << txt << "\n";
	llvm::errs().resetColor();
	llvm::errs().flush();
}

void inline _fatal(string txt, string *location=nullptr) {
	_log("fatal", raw_ostream::RED, true, txt, location);
	exit(EXIT_FAILURE);
}

void inline _error(string txt, string *location=nullptr) {
	_log("error", raw_ostream::RED, false, txt, location);
}

void inline _warning(string txt, string *location=nullptr) {
	_log("warning", raw_ostream::MAGENTA, false, txt, location);
}

#ifdef NDEBUG
void inline _debug(string txt, string *location=nullptr) {
	_log("debug", raw_ostream::CYAN, true, txt, location);
}
#else
#define _debug(txt, ...)     ((void)0)
#endif


/******************************************************************************
 * LLVM code
 ***/
class MyASTVisitor : public RecursiveASTVisitor<MyASTVisitor>
{
public:
	MyASTVisitor(Rewriter &r) : rewriter(r)
	{
		label_id = 0;
	}

	bool VisitStmt(Stmt *statement)
    {
		if (skip_this_function) {
			return true;
		}

		// Debug info {{{
		/*string returnRepr;
		llvm::raw_string_ostream returnReprStream(returnRepr);
		statement->dump(returnReprStream, rewriter.getSourceMgr());
		_debug("STMT <<<" + returnReprStream.str() + ">>>");*/
		// }}}

        if (isa<ReturnStmt>(statement)) {
            ReturnStmt *return_statement = cast<ReturnStmt>(statement);
            Stmt::child_iterator returned_value_statement = 
            	return_statement->child_begin();
            Stmt::child_iterator end = return_statement->child_end();
			SourceLocation location = return_statement->getLocStart();

			_debug("Return", &location.printToString(rewriter.getSourceMgr()));

			stringstream buffer;
			buffer << "\n__RETURN" << label_id++ << "__:\n";
			rewriter.InsertText(location, buffer.str(), false, true);
        }

        return true;
    }

	bool VisitFunctionDecl(FunctionDecl *func)
	{
		if (!func->hasBody()) {
			return true;
		}

		string name = func->getName();
		SourceLocation current_location = func->getLocation();
		// string location = rewriter.getSourceMgr().getFilename(func->getLocation()).str();
		string location = current_location.printToString(rewriter.getSourceMgr());
		FileID current_fid = rewriter.getSourceMgr().getFileID(current_location);
		FileID main_fid = rewriter.getSourceMgr().getMainFileID();
		skip_this_function = current_fid != main_fid;


		if (func->isExternC()) {
			name = prefix + name + suffix;
		}

		_debug("Function: " + name, &location);

		if (skip_this_function) {
			return true;
		}
	
		QualType return_type = func->getReturnType();

		// Search for collisions
		map<string, string>::iterator last = function_return_type.find(name);
		// If already exist but the return type is different, warn about it
		if (last != function_return_type.end()) {
			_warning("Function " + name + " alredy defined", &location);
			if (last->second.compare(return_type.getAsString()) != 0) {
				_fatal("Redefinition of " + func->getName().str() +
					" with different return type", &location);
			}
		}

		// If void return, add a custom return
        if (return_type.getTypePtr()->isVoidType()) {
           	Stmt *body = func->getBody();
            SourceLocation location = body->getLocEnd();
            
            stringstream buffer;
			buffer << "\n__RETURN" << label_id++ << "__:\n";
			buffer << "return;\n";
			rewriter.InsertText(location, buffer.str(), false, true);
        }

		function_return_type[name] = return_type.getAsString();
		function_return_info[name] = make_pair(return_type.getTypePtr()->isVoidType(),
			location);

		return true;
	}

private:
	Rewriter &rewriter;
	unsigned long long int label_id; // To ensure label uniqueness
};

// Implementation of the ASTConsumer interface for reading an AST produced
// by the Clang parser.
class MyASTConsumer : public ASTConsumer
{
public:
	MyASTConsumer(Rewriter &r) : visitor(r) {}

	// Override the method that gets called for each parsed top-level
	// declaration.
	virtual bool HandleTopLevelDecl(DeclGroupRef ref)
	{
		for (DeclGroupRef::iterator ref_it = ref.begin(), e = ref.end(); ref_it != e; ++ref_it) {
			// Traverse the declaration using our AST visitor.
			visitor.TraverseDecl(*ref_it);
		}
		return true;
	}

private:
	MyASTVisitor visitor;
};

/******************************************************************************
 * LLVM utils code
 ***/
void inline _add_include_path(CompilerInstance &compiler, StringRef path)
{

	compiler.getHeaderSearchOpts().AddPath(
		//StringRef("C:\\Archivos de programa\\Microsoft Visual Studio 10.0\\VC\\include"),
		path,
		frontend::Angled,   // Group
		false,              // IsFramework
		true);              // IgnoreSysRoot
}

void inline _define_macro(CompilerInstance &compiler, StringRef macro)
{
	compiler.getPreprocessorOpts().addMacroDef(macro);
}


void inline _undefine_macro(CompilerInstance &compiler, StringRef macro)
{
	compiler.getPreprocessorOpts().addMacroUndef(macro);
}

/******************************************************************************
 * Main code
 ***/
void inline _print_help(string progname)
{
	llvm::errs() << "\n";
	llvm::errs().changeColor(raw_ostream::YELLOW);
	llvm::errs() << "Usage: " << progname << " [-p prefix] [-s suffix] [-I path] [-D macro] [-U macro] [-o file ] <filenames list>\n";
	llvm::errs().resetColor();
	llvm::errs() << "    -p prefix      Prefix for exported symbols\n" <<
		            "    -s suffix      Suffix for exported symbols\n" <<
					"    -I path		Search headers in this paths\n" <<
					"    -D macro		Define this macro\n" <<
					"    -U macro		Undefine this macro\n" <<
					"    -o file        Write return information in this file\n" <<
		            "\n";
}

void inline _geopt_fatal(string progname, string msg)
{
	assert(!progname.empty());
	assert(!msg.empty());

	_error(msg);
	_print_help(progname);
	exit(EXIT_FAILURE);
}

int _parse_args(int argc, char *argv[])
{
	int o;
	while ((o = getopt(argc, argv, "D:hI:p:s:U:o:")) != -1) {
		switch (o) {
		case 'h':
			_print_help(argv[0]);
			exit(EXIT_SUCCESS);
		case 'p':
			prefix = optarg;
			break;
		case 's':
			suffix = optarg;
			break;
		case 'I':
			include_paths.push_back(optarg);
			break;
		case 'D':
			defines.push_back(optarg);
			break;
		case 'U':
			undefines.push_back(optarg);
			break;
		case 'o':
			output_file = string(optarg);
			break;
		case '?':
			if (optopt == 'p' || optopt == 's' || optopt == 'I' ||
				optopt == 'D' || optopt == 'U' || optopt == 'o')
			{
				string msg;
				msg += "Option '-";
				msg += optopt;
				msg += "' requires an argument";
				_geopt_fatal(argv[0], msg);
			} else {
				string msg;
				msg += "Unknown option '-";
				msg += optopt;
				msg += "'";
				_geopt_fatal(argv[0], msg);
			}
			break;
		default:
			_geopt_fatal(argv[0], "Unknown getopt case");
		}
	}

	if (optind == argc) {
		_geopt_fatal(argv[0], "Must provide at least one file");
	}

	return optind;
}

void _instrument_code(string file_name)
{
	assert(!file_name.empty());

	CompilerInstance compiler;

	DiagnosticOptions *dopts = new DiagnosticOptions();
	dopts->ShowColors = 1;
	TextDiagnosticPrinter *printer = new TextDiagnosticPrinter(
		llvm::errs(), dopts, false);
	compiler.createDiagnostics(printer, true);

	// Initialize target info with the default triple for our platform.
	TargetOptions *options = new TargetOptions();
	options->Triple = llvm::sys::getDefaultTargetTriple();
	TargetInfo *info =
		TargetInfo::CreateTargetInfo(compiler.getDiagnostics(), options);
	compiler.setTarget(info);

	compiler.createFileManager();
	FileManager &file_mgr = compiler.getFileManager();
	compiler.createSourceManager(file_mgr);
	SourceManager &src_mgr = compiler.getSourceManager();

	// Add additional include paths
	if (!include_paths.empty()) {
		for (std::vector<string>::iterator it = include_paths.begin();
			it != include_paths.end(); ++it)
		{
			_debug(string("New include path: ") + *it);
			_add_include_path(compiler, *it);
		}
	}

	// Add additional defines
	if (!defines.empty()) {
		for (std::vector<string>::iterator it = defines.begin();
			it != defines.end(); ++it)
		{
			_debug(string("New define: ") + *it);
			_define_macro(compiler, *it);
		}
	}

	// Add additional undefines
	if (!undefines.empty()) {
		for (std::vector<string>::iterator it = undefines.begin();
			it != undefines.end(); ++it)
		{
			_debug(string("New undefine: ") + *it);
			_undefine_macro(compiler, *it);
		}
	}

	compiler.createPreprocessor(TU_Module);
	compiler.createASTContext();

	// Set the main file handled by the source manager to the input file.
	const FileEntry *file_in = file_mgr.getFile(file_name);
	if (file_in == nullptr) {
		_fatal(string("Unkown file: ") + file_name);
	}
	src_mgr.setMainFileID(
		src_mgr.createFileID(file_in, SourceLocation(), SrcMgr::C_User));
	compiler.getDiagnosticClient().BeginSourceFile(
		compiler.getLangOpts(),
		&compiler.getPreprocessor());

	// Stype options
	LangOptions lo = compiler.getLangOpts();
	lo.MicrosoftExt = true;
	lo.CPlusPlus = false;

	Rewriter rewriter;
	rewriter.setSourceMgr(src_mgr, lo);
	MyASTConsumer TheConsumer(rewriter);

	// Parse
	printer->BeginSourceFile(lo, &compiler.getPreprocessor());
	ParseAST(compiler.getPreprocessor(), &TheConsumer, compiler.getASTContext());
	printer->EndSourceFile();

	const RewriteBuffer *rewrite_buffer =
		rewriter.getRewriteBufferFor(src_mgr.getMainFileID());
	if (rewrite_buffer) {
		ofstream dst_file(file_name + ".instrumented");

		for (RewriteBuffer::iterator b = rewrite_buffer->begin(),
			e = rewrite_buffer->end(); b != e; ++b)
		{
			if (*b == '\r') {
				continue;
			}
			dst_file << *b;
		}

		dst_file.close();

	} else {
		_warning("Any change has been made");
	}
}

void _write_return_info()
{
	ofstream return_info_file(output_file);
	return_info_file << "{";
	for (map<string, pair<bool, string>>::iterator it = function_return_info.begin();
		it != function_return_info.end(); ++it)
	{
		return_info_file << "\n  \"" << it->first << "\": { \"void_return\": ";
		return_info_file << (it->second.first ? "true" : "false");
		replace(it->second.second.begin(), it->second.second.end(), '\\', '/');
		return_info_file << ", \"_file\": \"" << it->second.second << "\" },";
	}
	long pos = return_info_file.tellp();
	return_info_file.seekp(pos - 1);
	return_info_file << "\n}\n";
	return_info_file.close();
}

int main(int argc, char *argv[])
{
	int file_pos = _parse_args(argc, argv);
	for (unsigned int i = file_pos; i < argc; i++) {
		current_file_name = string(argv[i]);
		_instrument_code(current_file_name);
	}
	_write_return_info();
	return EXIT_SUCCESS;
}