#include "codegen.hpp"
#include "node.hpp"
#include "grammar.hpp"

#include <vector>
#include <typeinfo>

#include <llvm/IR/Value.h>
#include <llvm/IR/Type.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/PassManager.h>
#include <llvm/IR/Instructions.h>
#include <llvm/Assembly/PrintModulePass.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/ExecutionEngine/ExecutionEngine.h>
#include <llvm/ExecutionEngine/JIT.h>
#include <llvm/ExecutionEngine/Interpreter.h>

using namespace std;
using namespace llvm;

CodeGenContext::CodeGenContext( )
{
	m_module = new Module( "main", getGlobalContext( ) );
}

void CodeGenContext::generateCode( NBlock& _root )
{
	/* Create the top level interpreter function to call as entry */
	vector<Type*> argTypes;
	FunctionType *ftype = FunctionType::get( Type::getVoidTy( getGlobalContext( ) ),
		makeArrayRef( argTypes ), false );
	m_mainFunction = Function::Create( ftype, GlobalValue::InternalLinkage, "main", m_module );
	BasicBlock *bblock = BasicBlock::Create( getGlobalContext( ), "entry", m_mainFunction, 0 );

	pushBlock( bblock );
	_root.codeGen( *this );
	ReturnInst::Create( getGlobalContext( ), bblock );
	popBlock( );
}

void CodeGenContext::printGeneratedCode( )
{
	PassManager pm;
	pm.add( createPrintModulePass( &outs( ) ) );
	pm.run( *m_module );
}

std::string CodeGenContext::getGeneratedCode( )
{
	PassManager pm;
	std::string s;
	raw_string_ostream os( s );
	pm.add( createPrintModulePass( &os ) );
	pm.run( *m_module );
	return s;
}

GenericValue CodeGenContext::runCode( )
{
	string error;
	GenericValue returnValue;

	ExecutionEngine *engine = EngineBuilder( m_module ).setErrorStr( &error ).create( );
	if( !engine ) {
		cerr << "ERROR: " << error << endl;
		return returnValue;
	}
	
	vector<GenericValue> args;
	returnValue = engine->runFunction( m_mainFunction, args );
	
	return returnValue;
}

void CodeGenContext::pushBlock( BasicBlock *block )
{
	m_blocks.push( new CodeGenBlock( ) );
	m_blocks.top( )->m_block = block;
}

void CodeGenContext::popBlock( )
{
	CodeGenBlock *top = m_blocks.top( );
	m_blocks.pop( );
	delete top;
}

BasicBlock *CodeGenContext::currentBlock( )
{
	return m_blocks.top( )->m_block;
}

llvm::Type *CodeGenContext::variableType( const NIdentifier& _type )
{
	if( _type.m_name.compare( "int" ) == 0 ) {
		return Type::getInt32Ty( getGlobalContext( ) );
	} else if( _type.m_name.compare( "double" ) == 0 ) {
		return Type::getDoubleTy( getGlobalContext( ) );
	} else if( _type.m_name.compare( "void" ) == 0 ) {
		return Type::getVoidTy( getGlobalContext( ) );
	} else {
		cerr << "illegal type '" << _type << "'" << endl;
		return 0;
	}
}

// code generators per ADT node

Value* Node::codeGen( CodeGenContext& context )
{
	cout << "Generating code for abstract node of type " << typeid( *this ).name( ) << endl;
	return 0;
}

Value *NBlock::codeGen( CodeGenContext& context )
{
	Value *code = 0;
	StatementList::const_iterator end = m_statements.end( );
	cout << "Starting block" << endl;
	for( StatementList::const_iterator it = m_statements.begin( );
		it != end; it++ ) {
		cout << "Generating code for " << typeid( *(*it) ).name( ) << endl;
		code = (*(*it)).codeGen( context );
	}
	cout << "End of block" << endl;
	return code;
}

Value *NAssignment::codeGen( CodeGenContext& context )
{
	cout << "Creating assignment for variable " << m_lhs.m_name << endl;
	CodeGenContext::Locals::const_iterator var;
	if( context.locals( ).find( m_lhs.m_name ) == context.locals( ).end( ) ) {
		cerr << "Undeclared variable '" << m_lhs.m_name << "'" << endl;
		return 0;
	}
	return new StoreInst( m_rhs.codeGen( context ),
		context.locals( )[m_lhs.m_name], false,
		context.currentBlock( ) );
}

Value *NVariableDeclaration::codeGen( CodeGenContext& context )
{
	cout << "Creating for a a variable declaration with name '"
		<< m_name.m_name << "' of type '" << m_type.m_name << "'" << endl;
	AllocaInst *alloca = new AllocaInst( CodeGenContext::variableType( m_type ),
		m_name.m_name.c_str( ), context.currentBlock( ) );
	context.locals( )[m_name.m_name] = alloca;
	if( m_assigned_expr ) {
		NAssignment assignment( m_name, *m_assigned_expr );
		assignment.codeGen( context );
	}
	return alloca;
}

Value *NInteger::codeGen( CodeGenContext& context )
{
	cout << "Creating code for integer constant '" << m_value << "'" << endl;
	return ConstantInt::get( Type::getInt32Ty( getGlobalContext( ) ), m_value, true );
}

Value *NFunctionCall::codeGen( CodeGenContext& context )
{
	Function *function = context.module( )->getFunction( m_name.m_name.c_str( ) );
	if( function == NULL ) {
		cerr << "No such function '" << m_name.m_name << "'" << endl;
	}
	
	vector<Value *> args;
	ExpressionList::const_iterator it, end = m_arguments.end( );
	for( it = m_arguments.begin( ); it != end; it++ ) {
		args.push_back( (**it).codeGen( context ) );
	}
	CallInst *call = CallInst::Create( function, makeArrayRef( args ), "", context.currentBlock( ) );
	cout << "Created function invocation for '" << m_name.m_name << "'" << endl;
	
	return call;
}

Value *NExpressionStatement::codeGen( CodeGenContext& context )
{
	cout << "Generating code for " << typeid( m_expression ).name( ) << endl;
	return m_expression.codeGen( context );
}

Value *NFunctionDeclaration::codeGen( CodeGenContext& context )
{
	vector<Type *> argTypes;
	VariableList::const_iterator it, end = m_arguments.end( );
	for( it = m_arguments.begin( );  it != end; it++ ) {
		argTypes.push_back( CodeGenContext::variableType( (**it).m_type ) );
	}
	FunctionType *ftype = FunctionType::get( CodeGenContext::variableType( m_type ),
		makeArrayRef( argTypes ), false );
	Function *function = Function::Create( ftype, GlobalValue::InternalLinkage,
		m_name.m_name.c_str( ), context.module( ) );
	BasicBlock *bblock = BasicBlock::Create( getGlobalContext( ), "entry", function, 0 );
	
	context.pushBlock( bblock );
	
	for( it = m_arguments.begin( ); it != end; it++ ) {
		(**it).codeGen( context );
	}
	
	m_block.codeGen( context );
	ReturnInst::Create( getGlobalContext( ), bblock );
	
	context.popBlock( );
	
	cout << "Created function '" << m_name.m_name << "'" << endl;
	
	return function;
}

Value *NIdentifier::codeGen( CodeGenContext& context )
{
	cout << "Creating identifier reference to variable '" << m_name << "'" << endl;
	if( context.locals( ).find( m_name ) == context.locals( ).end( ) ) {
		cerr << "undeclared variable '" << m_name << "'" << endl;
		return 0;
	}
	return new LoadInst( context.locals( )[m_name], "", false, context.currentBlock( ) );
}

Value *NBinaryOperator::codeGen( CodeGenContext& context )
{
	cout << "Creating binary operation '" << m_op << "'" << endl;
	
	Instruction::BinaryOps instr;
	switch( m_op ) {
		case TPLUS:	instr = Instruction::Add; break;
		case TMUL:	instr = Instruction::Mul; break;
		default:	cerr << "Illegal operation '" << m_op << "'" << endl; return 0;
	}
	
	return BinaryOperator::Create( instr, m_lhs.codeGen( context ),
		m_rhs.codeGen( context ), "", context.currentBlock( ) );
}