/* Tuplr: A programming language
 * Copyright (C) 2008-2009 David Robillard <dave@drobilla.net>
 *
 * Tuplr is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero General Public License as published by the
 * Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 *
 * Tuplr is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General
 * Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with Tuplr.  If not, see <http://www.gnu.org/licenses/>.
 */

/** @file
 * @brief Compile to C
 */

#include <map>
#include <sstream>
#include <boost/format.hpp>
#include "tuplr.hpp"

using namespace std;
using boost::format;

typedef string Type;
typedef string Value;

struct Function {
	string returnType;
	string name;
	string text;
};

static inline Value*    llVal(CVal     v) { return static_cast<Value*>(v); }
static inline Function* llFunc(CFunc f)   { return static_cast<Function*>(f); }

static const Type*
llType(const AType* t)
{
	if (t->kind == AType::PRIM) {
		if (t->head()->str() == "Nothing") return new string("void");
		if (t->head()->str() == "Bool")    return new string("bool");
		if (t->head()->str() == "Int")     return new string("int");
		if (t->head()->str() == "Float")   return new string("float");
		throw Error(t->loc, string("Unknown primitive type `") + t->str() + "'");
	} else if (t->kind == AType::EXPR && t->head()->str() == "Fn") {
		AType::const_iterator i     = t->begin();
		const ATuple*         protT = (*++i)->to<const ATuple*>();
		const AType*          retT  = (*i)->as<const AType*>();
		if (!llType(retT))
			return NULL;

		Type* ret = new Type(*llType(retT) + " (*)(");
		FOREACHP(ATuple::const_iterator, i, protT) {
			const AType* at = (*i)->to<const AType*>();
			const Type*  lt = llType(at);
			if (!lt)
				return NULL;
			*ret += *lt;
		}
		*ret += ")";

		return ret;
	}
	return NULL; // non-primitive type
}


/***************************************************************************
 * LLVM Engine                                                             *
 ***************************************************************************/

struct CEngine : public Engine {
	CEngine()
		: out(
			"#include <stdint.h>\n"
			"#include <stdbool.h>\n"
			"void* tuplr_gc_allocate(unsigned size, uint8_t tag);\n\n")
	{
	}

	CFunc startFunction(CEnv& cenv,
			const std::string& name, const AType* retT, const ATuple& argsT,
			const vector<string> argNames)
	{
		vector<const Type*> cprot;
		FOREACH(ATuple::const_iterator, i, argsT) {
			AType* at = (*i)->as<AType*>();
			THROW_IF(!llType(at), Cursor(), string("non-concrete parameter :: ")
				+ at->str())
			cprot.push_back(llType(at));
		}

		THROW_IF(!llType(retT), Cursor(),
				(format("return has non-concrete type `%1%'") % retT->str()).str());

		Function* f = new Function();
		f->returnType = *llType(retT);
		f->name = name;
		f->text += f->returnType + "\n" + f->name + "(";
		ATuple::const_iterator ai         = argsT.begin();
		vector<string>::const_iterator ni = argNames.begin();
		for (; ai != argsT.end(); ++ai, ++ni) {
			if (ai != argsT.begin())
				f->text += ", ";
			f->text += *llType((*ai)->as<const AType*>()) + " " + *ni;
		}
		f->text += ")\n{\n";

		out += f->text;
		return f;
	}

	void finishFunction(CEnv& cenv, CFunc f, const AType* retT, CVal ret) {
		out += "return " + *(Value*)ret + ";\n}\n\n";
	}

	void eraseFunction(CEnv& cenv, CFunc f) {
		cenv.err << "C backend does not support JIT (eraseFunction)" << endl;
	}

	CVal compileCall(CEnv& cenv, CFunc func, const vector<CVal>& args) {
		Value* varname = new string(cenv.penv.gensymstr("x"));
		Function* f = llFunc(func);
		out += (format("const %s %s = %s(") % f->returnType % *varname % f->name).str();
		FOREACH(vector<CVal>::const_iterator, i, args)
			out += *llVal(*i);
		out += ");\n";
		return varname;
	}

	CFunc compileFunction(CEnv& cenv, AFn* fn, const AType& argsT);

	CVal compileTup(CEnv& cenv, const AType* type, const vector<CVal>& fields);
	CVal compileDot(CEnv& cenv, CVal tup, int32_t index);
	CVal compileLiteral(CEnv& cenv, AST* lit);
	CVal compilePrimitive(CEnv& cenv, APrimitive* prim);
	CVal compileIf(CEnv& cenv, AIf* aif);
	CVal compileGlobal(CEnv& cenv, AType* type, const string& name, CVal val);
	CVal getGlobal(CEnv& cenv, CVal val);

	void writeModule(CEnv& cenv, std::ostream& os) {
		os << out;
	}

	const string call(CEnv& cenv, CFunc f, AType* retT) {
		cenv.err << "C backend does not support JIT (call)" << endl;
		return "";
	}

	std::string out;
};

Engine*
tuplr_new_c_engine()
{
	return new CEngine();
}

/***************************************************************************
 * Code Generation                                                         *
 ***************************************************************************/

CVal
CEngine::compileTup(CEnv& cenv, const AType* type, const vector<CVal>& fields)
{
	return NULL;
}

CVal
CEngine::compileDot(CEnv& cenv, CVal tup, int32_t index)
{
	return NULL;
}

CVal
CEngine::compileLiteral(CEnv& cenv, AST* lit)
{
	return new Value(lit->str());
}

CFunc
CEngine::compileFunction(CEnv& cenv, AFn* fn, const AType& argsT)
{
	CEngine* engine      = reinterpret_cast<CEngine*>(cenv.engine());
	AType*   genericType = cenv.type(fn);
	AType*   thisType    = genericType;
	Subst    argsSubst;

	// Build and apply substitution to get concrete type for this call
	if (!genericType->concrete()) {
		argsSubst = cenv.tenv.buildSubst(genericType, argsT);
		thisType = argsSubst.apply(genericType)->as<AType*>();
	}

	THROW_IF(!thisType->concrete(), fn->loc,
			string("call has non-concrete type %1%\n") + thisType->str());

	Object::pool.addRoot(thisType);
	CFunc f = fn->impls.find(thisType);
	if (f)
		return f;

	ATuple* protT = thisType->prot();

	vector<string> argNames;
	for (ATuple::const_iterator i = fn->prot()->begin(); i != fn->prot()->end(); ++i)
		argNames.push_back((*i)->str());

	// Write function declaration
	const string name = (fn->name == "") ? cenv.penv.gensymstr("_fn") : fn->name;
	f = llFunc(cenv.engine()->startFunction(cenv, name,
			thisType->last()->to<AType*>(),
			*protT, argNames));

	cenv.push();
	Subst oldSubst = cenv.tsubst;
	cenv.tsubst = Subst::compose(cenv.tsubst, argsSubst);

	// Bind argument values in CEnv
	vector<Value*> args;
	AFn::const_iterator    p  = fn->prot()->begin();
	ATuple::const_iterator pT = protT->begin();
	for (; p != fn->prot()->end(); ++p, ++pT) {
		AType* t = (*pT)->as<AType*>();
		const Type* lt = llType(t);
		THROW_IF(!lt, fn->loc, "untyped parameter\n");
		cenv.def((*p)->as<ASymbol*>(), *p, t, new string((*p)->str()));
	}

	// Write function body
	try {
		fn->impls.push_back(make_pair(thisType, f));
		CVal retVal = NULL;
		for (AFn::iterator i = fn->begin() + 2; i != fn->end(); ++i)
			retVal = (*i)->compile(cenv);
		cenv.engine()->finishFunction(cenv, f, cenv.type(fn->last()), retVal);
	} catch (Error& e) {
		cenv.pop();
		throw e;
	}
	cenv.tsubst = oldSubst;
	cenv.pop();
	return f;
}

CVal
CEngine::compileIf(CEnv& cenv, AIf* aif)
{
	CEngine* engine  = reinterpret_cast<CEngine*>(cenv.engine());
	Value*   varname = new string(cenv.penv.gensymstr("if"));
	out += (format("%s %s;\n") % *llType(cenv.type(aif)) % *varname).str();
	size_t idx = 1;
	for (AIf::iterator i = aif->begin() + 1; ; ++i, idx += 2) {
		AIf::iterator next = i;
		if (++next == aif->end())
			break;

		if (idx > 1)
			out += "else {\n";

		Value* condV = llVal((*i)->compile(cenv));
		out += (format("if (%s) {\n") % *condV).str();

		Value* thenV = llVal((*next)->compile(cenv));
		out += (format("%s = %s;\n}\n") % *varname % *thenV).str();
	}

	// Emit final else block
	out += "else {\n";
	Value* elseV = llVal(aif->last()->compile(cenv));
	out += (format("%s = %s;\n}\n") % *varname % *elseV).str();

	for (size_t i = 1; i < idx / 2; ++i)
		out += "}";

	return varname;
}

CVal
CEngine::compilePrimitive(CEnv& cenv, APrimitive* prim)
{
	APrimitive::iterator i = prim->begin();
	++i;

	CEngine*     engine = reinterpret_cast<CEngine*>(cenv.engine());
	Value*       a      = llVal((*i++)->compile(cenv));
	Value*       b      = llVal((*i++)->compile(cenv));
	const string n      = prim->head()->to<ASymbol*>()->str();
	string       op     = n;

	// Convert operator to C operator if they don't match
	if (n == "=")        op = "==";
	else if (n == "and") op = "&";
	else if (n == "or")  op = "|";
	else if (n == "xor") op = "^";

	op = string(" ") + op + " ";

	string val("(");
	val += *a + op + *b;
	while (i != prim->end())
		val += op + *llVal((*i++)->compile(cenv));
	val += ")";

	Value* varname = new string(cenv.penv.gensymstr("x"));
	out += (format("const %s %s = %s;\n") % *llType(cenv.type(prim)) % *varname % val).str();
	return varname;
}

CVal
CEngine::compileGlobal(CEnv& cenv, AType* type, const string& name, CVal val)
{
	return NULL;
}

CVal
CEngine::getGlobal(CEnv& cenv, CVal val)
{
	return NULL;
}