#include "compilepass.h"
#include "clasplayer.h"
#include <ast.h>
#include <iostream>
#include "query/containers.h"
#include "instructions/instr-containers.h"
#include "ExternLayer.h"

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

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;
}

llvm::Value*
CompilePass::CodeScopeUnit::convertType(llvm::Value* source, llvm::Type* tyTarget){
	LLVMLayer* llvm = pass->man->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);
		}
	}

	assert(false && "no automatic type conversion possible");
}

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;
    CompilePass::Context context{function, this, pass};
    containers::Instructions instructions = containers::Instructions(context);

    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]);

            if (left->getType()!= right->getType()) {
            	right =	convertType(right, left->getType());
            }
            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);

            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);
                    }
            );

            FunctionUnit* calleeUnit = pass->getFunctionUnit(string(fname));

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

            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(blockPrev);
            return l.builder.CreateCall(callee, args, hintVarDecl);
        }

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

        case Operator::SWITCH:
        {
        	return nullptr; //instructions.compileSwitch();
        }

        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;

                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);
            }
        };


        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);
                };

                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;

    std::transform(entry->__args.begin(), entry->__args.end(), std::inserter(types, types.end()),
            [this, llvm, 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(pass->man->root->expandType(entry->__definitions.at(argid)));
            });

    llvm::FunctionType *ft = llvm::FunctionType::get(llvm->toLLVMType(pass->man->root->expandType(entry->__definitions[0])), types, false);
    raw = llvm::cast<llvm::Function>(llvm->module->getOrInsertFunction(function->__name, 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;
    }

    const std::string blockName =  "entry";
    llvm->builder.CreateRet(entryCompilation->compile(blockName));

    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 CompilePass::FunctionQuery& f){
    ManagedFnPtr fkey = man->root->findFunction(f.name);

    	//external functions:
    if (!fkey){
    	return nullptr;
    }

    if (!functions.count(&*fkey)){
        functions.emplace(&*fkey, std::unique_ptr<FunctionUnit>(new FunctionUnit(fkey, this)));
    }

    return functions.at(&*fkey).get();
}

void
CompilePass::run(){
    //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(move(nameMain));
    entry = unitMain->compile();
}

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


//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;
    //}