#include "compilepass.h"
#include "clasplayer.h"
#include <ast.h>
#include <iostream>
#include "query/containers.h"
#include "query/context.h"
#include "compilation/instr-containers.h"
#include "ExternLayer.h"
#include "pass/adhocpass.h"

#include <boost/optional.hpp>

using namespace std;
using namespace xreate;
using namespace llvm;

//SECTIONTAG types/convert implementation
llvm::Value*
doAutomaticTypeConversion(llvm::Value* source, llvm::Type* tyTarget, LLVMLayer* llvm){
	if (tyTarget->isIntegerTy() && source->getType()->isIntegerTy())
	{
		llvm::IntegerType* tyTargetInt = llvm::dyn_cast<IntegerType>(tyTarget);
		llvm::IntegerType* tySourceInt = llvm::dyn_cast<IntegerType>(source->getType());

		if (tyTargetInt->getBitWidth() < tySourceInt->getBitWidth()){
			return llvm->builder.CreateCast(llvm::Instruction::Trunc, source, tyTarget);
		}

		if (tyTargetInt->getBitWidth() > tySourceInt->getBitWidth()){
			return llvm->builder.CreateCast(llvm::Instruction::SExt, source, tyTarget);
		}
	}

	return source;
}

CompilePass::CodeScopeUnit::CodeScopeUnit(CodeScope* codeScope, FunctionUnit* f, CompilePass* compilePass)
    : scope(codeScope), pass(compilePass), function(f)
{}

void
CompilePass::CodeScopeUnit::bindArg(llvm::Value* var, std::string&& name)
{
    assert(scope->__vartable.count(name));
    VID id = scope->__vartable.at(name);
    __rawVars[id] = var;
}

//SECTIONTAG context-based function specialization
CompilePass::FunctionUnit*
CompilePass::CodeScopeUnit::findFunctionUnit(const std::string f){
	AST* ast = pass->man->root;

	const ScopePacked& scopeId = pass->man->clasp->pack(scope);

	const std::list<Expression>& context = pass->queryContext->getContext(scopeId);
	std::unordered_set<string> contextIndex;
	for (const Expression& c: context){
		if (c.__state != Expression::IDENT) continue;
		contextIndex.insert(c.getValueString());
	}

	ManagedFnPtr function = pass->man->root->findFunctionVariant(f, FunctionSpecializationQuery{contextIndex});
	if (!function.isValid()) return nullptr;

	return pass->getFunctionUnit(function);
}

llvm::Value*
CompilePass::CodeScopeUnit::process(const Expression& expr, const std::string& hintVarDecl){
#define VARNAME(x) (hintVarDecl.empty()? x: hintVarDecl)
    llvm::Value *left; llvm::Value *right;
    LLVMLayer& l = *pass->man->llvm;
    containers::Instructions instructions = containers::Instructions({function, this, pass});

    switch (expr.op) {
        case Operator::ADD:
        case Operator::SUB:
        case Operator::MUL:
        case Operator::DIV:
        case Operator::EQU:
        case Operator::LSS:
        case Operator::GTR:
            assert(expr.__state == Expression::COMPOUND);
            assert(expr.operands.size() == 2);

            left = process(expr.operands[0]);
            right = process(expr.operands[1]);

            //SECTIONTAG types/convert binary operation
           	right =	doAutomaticTypeConversion(right, left->getType(), &l);
            break;

        default:;
    }

    switch (expr.op) {
        case Operator::ADD:
            return l.builder.CreateAdd(left, right, VARNAME("tmp_add"));
            break;

        case Operator::SUB:
            return l.builder.CreateSub(left, right, VARNAME("tmp_sub"));
            break;

        case Operator::MUL:
            return l.builder.CreateMul(left, right, VARNAME("tmp_mul"));
            break;

        case Operator::DIV:
            return l.builder.CreateSDiv(left, right, VARNAME("tmp_div"));
            break;

        case Operator::EQU:
            left->dump();
            right->dump();
            return l.builder.CreateICmpEQ(left, right, VARNAME("tmp_equ"));
            break;

        case Operator::LSS:
            return l.builder.CreateICmpSLT(left, right, VARNAME("tmp_lss"));
            break;

        case Operator::GTR:
            return l.builder.CreateICmpSGT(left, right, VARNAME("tmp_gtr"));
            break;

        case Operator::NEG:
            left = process(expr.operands[0]);
            return l.builder.CreateNeg(left, VARNAME("tmp_neg"));
            break;

        case Operator::CALL: {
            assert(expr.__state == Expression::COMPOUND);

            llvm::BasicBlock* blockCurrent = pass->man->llvm->builder.GetInsertBlock();

            std::string fname = expr.getValueString();
            std::vector<llvm::Value *> args;
            args.reserve(expr.operands.size());

            std::transform(expr.operands.begin(), expr.operands.end(), std::inserter(args, args.end()),
				[this](const Expression &operand) {
					return process(operand);
				}
            );

            ScopePacked outerScopeId = pass->man->clasp->pack(this->scope);
            llvm::Value* callee = this->function->lateContext.findFunction(fname, outerScopeId, function->raw);

            	// external function
            if (!callee) {
            	llvm::Function* external = pass->man->llvm->layerExtern->lookupFunction(fname);
            	return l.builder.CreateCall(external, args, hintVarDecl);
            }

            //SECTIONTAG late-context propagation arg
            size_t calleeDomSize = pass->queryContext->getFunctionDomain(fname).size();
            if (calleeDomSize){
            	llvm::Value* argLateContext = function->lateContext.compileArgument(fname, outerScopeId);
            	args.push_back(argLateContext);
            }

            //TODO implement inline functions
            /*
            if (calleeUnit->isInline()) {
                return calleeUnit->compileInline(move(args), this->function);
            }
            */

            llvm::BasicBlock* blockPrev = pass->man->llvm->builder.GetInsertBlock();
            //llvm::Value* callee = calleeUnit->compile();
            pass->man->llvm->builder.SetInsertPoint(blockCurrent);
            return l.builder.CreateCall(callee, args, hintVarDecl);
        }

        case Operator::IF:
        {
        	return instructions.compileIf(expr, hintVarDecl);
        }

        case Operator::SWITCH:
        {
        	return instructions.compileSwitch(expr, hintVarDecl);
        }

        case Operator::LOOP_CONTEXT:
        {
        	return instructions.compileLoopContext(expr, hintVarDecl);
        }

        case Operator::LOGIC_AND: {
        	assert(expr.operands.size() == 1);
        	return process (expr.operands[0]);
        }

        case Operator::LIST:
        {
           return instructions.compileConstantArray(expr, hintVarDecl);
        };

        case Operator::LIST_RANGE:
        {
            assert(false); //no compilation phase for a range list
          //  return InstructionList(this).compileConstantArray(expr, l, hintRetVar);
        };

        case Operator::LIST_NAMED:
        {
            typedef Expanded<TypeAnnotation> ExpandedType;

            ExpandedType tyRaw = l.ast->expandType(expr.type);

            const std::vector<string> fields = (tyRaw.get().__operator == TypeOperator::CUSTOM)?
                l.layerExtern->getStructFields(l.layerExtern->lookupType(tyRaw.get().__valueCustom))
                : tyRaw.get().fields;

            std::map<std::string, size_t> indexFields;
            for(size_t i=0, size = fields.size(); i<size; ++i){
                indexFields.emplace(fields[i], i);
            }

            llvm::StructType* tyRecord = llvm::cast<llvm::StructType>(l.toLLVMType(tyRaw));
            llvm::Value* record = llvm::UndefValue::get(tyRecord);

            for (size_t i=0; i<expr.operands.size(); ++i){
                const Expression& operand = expr.operands.at(i);
                unsigned int fieldId = indexFields.at(expr.bindings.at(i));

                llvm::Value* result = 0;

//TODO Null ad hoc Llvm implementation
//                if (operand.isNone()){
//                    llvm::Type* tyNullField = tyRecord->getElementType(fieldId);
//                    result = llvm::UndefValue::get(tyNullField);
//
//                } else {
                    result = process(operand);
//                }

                assert (result);
                record = l.builder.CreateInsertValue(record, result, llvm::ArrayRef<unsigned>({fieldId}));
            }

            return record;
        };



        case Operator::MAP:
        {
            assert(expr.blocks.size());
            return instructions.compileMapSolid(expr, VARNAME("map"));
        };

        case Operator::FOLD:
        {
            return instructions.compileFold(expr, VARNAME("fold"));
        };

        case Operator::INDEX:
        {
                //TODO allow multiindex
            assert(expr.operands.size()==1);
            const std::string &ident = expr.getValueString();
            Symbol s = scope->findSymbol(ident);
            const TypeAnnotation& t = s.scope->findDefinition(s);
            const ExpandedType& t2 = pass->man->root->expandType(t);

            switch (t2.get().__operator)
            {
                case TypeOperator::STRUCT: case TypeOperator::CUSTOM:
                {
                    Expression idx = expr.operands.at(0);
                    assert(idx.__state == Expression::STRING);
                    std::string idxField = idx.getValueString();

                    llvm::Value* aggr  = compileSymbol(s, ident);
                    return instructions.compileStructIndex(aggr, t2, idxField);
                };

                case TypeOperator::ARRAY: {
                    std::vector<llvm::Value*> indexes;
                    std::transform(++expr.operands.begin(), expr.operands.end(), std::inserter(indexes, indexes.end()),
                                   [this] (const Expression& op){return process(op);}
                    );

                    return instructions.compileArrayIndex(s, indexes, VARNAME(string("el_") + ident));
                };

                default:
                    assert(false);
            }
        };

        	//SECTIONTAG adhoc actual compilation
        case Operator::ADHOC: {
        	assert(function->adhocImplementation && "Adhoc implementation not found");
        	string comm = expr.operands[0].getValueString();

        	CodeScope* scope = function->adhocImplementation->getImplementationForCommand(comm);
        	CodeScopeUnit* unitScope = function->getScopeUnit(scope);
        	return unitScope->compile();
        };

        case Operator::NONE:
            assert(expr.__state != Expression::COMPOUND);

            switch (expr.__state) {
                case Expression::IDENT: {
                    const std::string &ident = expr.getValueString();
                    Symbol s = scope->findSymbol(ident);
                    return compileSymbol(s, ident);
                }

                case Expression::NUMBER: {
                    int literal = expr.getValueDouble();
                    return llvm::ConstantInt::get(llvm::Type::getInt32Ty(llvm::getGlobalContext()), literal);
                }

                case Expression::STRING: {
                    return instructions.compileConstantStringAsPChar(expr.getValueString(), hintVarDecl);
                };

                case Expression::VARIANT: {
                	const ExpandedType& typVariant = pass->man->root->expandType(expr.type);
                	llvm::Type* typRaw = l.toLLVMType(typVariant);
                	int value = expr.getValueDouble();
                    return llvm::ConstantInt::get(typRaw, value);
                }

                default: {
                    break;
                }
            };

            break;

        default: break;

    }

    assert(false);
    return 0;
}

llvm::Value*
CompilePass::CodeScopeUnit::compile(const std::string& hintBlockDecl){
    if (raw != nullptr) return raw;


    if (!hintBlockDecl.empty()) {
        llvm::BasicBlock *block = llvm::BasicBlock::Create(llvm::getGlobalContext(), hintBlockDecl, function->raw);
        pass->man->llvm->builder.SetInsertPoint(block);
    }

    raw = process(scope->__body);
    return raw;
}

llvm::Value*
CompilePass::CodeScopeUnit::compileSymbol(const Symbol& s, std::string hintRetVar)
{
    CodeScope* scope = s.scope;
    CodeScopeUnit* self = function->getScopeUnit(scope);

    if (self->__rawVars.count(s.identifier))     {
        return self->__rawVars[s.identifier];
    }

    return self->__rawVars[s.identifier] = self->process(scope->findDeclaration(s), hintRetVar);
}

bool
CompilePass::FunctionUnit::isInline(){
    Symbol ret = Symbol{0, function->__entry};
    bool flagOnTheFly = SymbolAttachments::get<IsImplementationOnTheFly>(ret, false);

    return flagOnTheFly;
}

llvm::Function*
CompilePass::FunctionUnit::compile(){
    if (raw != nullptr) return raw;

    std::vector<llvm::Type *> types;
    LLVMLayer* llvm = pass->man->llvm;
    CodeScope* entry = function->__entry;
    AST* ast = pass->man->root;

    std::transform(entry->__args.begin(), entry->__args.end(), std::inserter(types, types.end()),
            [this, llvm, ast, entry](const std::string &arg)->llvm::Type* {
                assert(entry->__vartable.count(arg));
                VID argid = entry->__vartable.at(arg);
                assert(entry->__definitions.count(argid));
                return llvm->toLLVMType(ast->expandType(entry->__definitions.at(argid)));
            });

    	//SECTIONTAG late-context signature type
    size_t sizeDom = lateContext.domain->size();
    if (sizeDom) {
    	llvm::Type* tyBool = llvm::Type::getInt1Ty(llvm::getGlobalContext());
    	llvm::Type* tyDom = llvm::ArrayType::get(tyBool, sizeDom);
    	types.push_back(tyDom);
    }

    	//SECTIONTAG adhoc func signature determination
    llvm::Type* funcResultType;
    if (function->isPrefunction){
    	AdhocPass* adhocpass = reinterpret_cast<AdhocPass*>(pass->man->getPassById(PassId::AdhocPass));

    	adhocImplementation = adhocpass->determineForScope(entry);
    	funcResultType = llvm->toLLVMType(ast->expandType(adhocImplementation->getResultType()));
    } else {
    	funcResultType = llvm->toLLVMType(ast->expandType(entry->__definitions[0]));
    }

    llvm::FunctionType *ft = llvm::FunctionType::get(funcResultType, types, false);

    raw = llvm::cast<llvm::Function>(llvm->module->getOrInsertFunction(function->__name + std::to_string(function.id()), ft));

    CodeScopeUnit* entryCompilation = getScopeUnit(entry);
    llvm::Function::arg_iterator fargsI = raw->arg_begin();
    for (std::string &arg : entry->__args) {
        VID argid = entry->__vartable[arg];

        entryCompilation->__rawVars[argid] = fargsI;
        fargsI->setName(arg);
        ++fargsI;
    }

    if (sizeDom){
    	fargsI->setName("latecontext");
    	lateContext.raw = fargsI;
    }

    const std::string&blockName =  "entry";

    //SECTIONTAG types/convert function ret value
    llvm->builder.CreateRet(
    		doAutomaticTypeConversion(
    				entryCompilation->compile(blockName), funcResultType, llvm));

    llvm->moveToGarbage(ft);
    return raw;
}

llvm::Value*
CompilePass::FunctionUnit::compileInline(std::vector<llvm::Value *> &&args, CompilePass::FunctionUnit* outer){
    CodeScopeUnit* entryCompilation = outer->getScopeUnit(function->__entry);
    for(int i=0, size = args.size(); i<size; ++i) {
        entryCompilation->bindArg(args.at(i), string(entryCompilation->scope->__args.at(i)));
    }


    return entryCompilation->compile();
}

CompilePass::CodeScopeUnit*
CompilePass::FunctionUnit::getScopeUnit(CodeScope* scope){
    if (!scopes.count(scope)){
        CodeScopeUnit* unit = new CodeScopeUnit(scope, this, pass);
        scopes.emplace(scope, std::unique_ptr<CodeScopeUnit>(unit));
    }

    return scopes.at(scope).get();
}

CompilePass::CodeScopeUnit*
CompilePass::FunctionUnit::getEntry(){
	return getScopeUnit(function->getEntryScope());
}

CompilePass::CodeScopeUnit*
CompilePass::FunctionUnit::getScopeUnit(ManagedScpPtr scope){
    return getScopeUnit(&*scope);
}

CompilePass::FunctionUnit*
CompilePass::getFunctionUnit(const ManagedFnPtr& function){
	unsigned int id = function.id();

    if (!functions.count(id)){
    	FunctionUnit* unit = new FunctionUnit(function, this);
        functions.emplace(id, std::unique_ptr<FunctionUnit>(unit));
        return unit;
    }

    return functions.at(id).get();
}

void
CompilePass::run(){
	queryContext = reinterpret_cast<ContextQuery*> (man->clasp->getQuery(QueryId::ContextQuery));

    //Find out main function;
    ClaspLayer::ModelFragment model = man->clasp->query(Config::get("function-entry"));
    assert(model && "Error: No entry function found");
    assert(model->first != model->second && "Error: Ambiguous entry function");

    string nameMain = std::get<0>(ClaspLayer::parse<std::string>(model->first->second));
    FunctionUnit* unitMain = getFunctionUnit(man->root->findFunction(nameMain));
    entry = unitMain->compile();
}

llvm::Function*
CompilePass::getEntryFunction(){
	assert(entry);
	return entry;
}

void
CompilePass::prepareQueries(ClaspLayer* clasp){
	clasp->registerQuery(new containers::Query(), QueryId::ContainersQuery);
	clasp->registerQuery(new ContextQuery(), QueryId::ContextQuery);
}

//CODESCOPE COMPILATION PHASE


//FIND SYMBOL(compilation phase):
    //if (!forceCompile)
    //{
    //    return result;
    //}

    //    //search in already compiled vars
    //if (__rawVars.count(vId))
    //{
    //    return result;
    //}

    //if (!__declarations.count(vId)) {
    //    //error: symbol is uncompiled scope arg
    //    assert(false);
    //}

    //const Expression& e = __declarations.at(vId);

    //__rawVars[vId] = process(e, l, name);


//FIND FUNCTION
    //llvm::Function*
    //CompilePass::findFunction(const std::string& name){
    //    ManagedFnPtr calleeFunc = man->root->findFunction(name);
    //    assert(calleeFunc.isValid());

    //    return  nullptr;
    //}