#include "ast.h"
#include <stdexcept>
using namespace std;

TypeAnnotation::TypeAnnotation()
{}

TypeAnnotation::TypeAnnotation(const Atom<Type_t> &typ)
{
    __value = typ.get();
}

TypeAnnotation::TypeAnnotation(const TypeOperator &op, const TypeAnnotation &typ)
{}

llvm::Type*
TypeAnnotation::toLLVMType()
{
    switch (__value)
    {
        case Bool:
         return llvm::Type::getInt1Ty(llvm::getGlobalContext());

        case Int:
        case Float:
        case Num:
         return llvm::Type::getDoubleTy(llvm::getGlobalContext());

        default:
            assert(false);
    }

    return NULL;
}

Expression::Expression()
    : __op(Operator::NONE), state(INVALID)
{}

Expression::Expression(const Atom<Number_t>& number)
: state(NUMBER), __op(Operator::NONE), __valueD(number.get())
{
}

Expression::Expression(const Atom<Identifier_t> &ident)
    : state(IDENT), __op(Operator::NONE), __valueS(ident.get())
{
}

Expression::Expression(const Operator &op, const Expression &arg)
    : state(COMPOUND), __op(op)
{
    operands.push_back(arg);
}

void
Expression::addArg(const Expression &arg)
{
    assert(state == COMPOUND);
    assert(op!=Operator::NONE);

    operands.push_back(arg);
}

llvm::Value*
Expression::compile(LLVMLayer& l, Function* f,  std::string* retVarHint)
{
    std::string var;
    if (retVarHint && f->__vartable.count(*retVarHint))
    {
        var = *retVarHint;
    }

    #define VARNAME(x) (var.size()? var : x)
    llvm::Value* left; llvm::Value* right;


    switch (__op)
    {
        case Operator::ADD: case Operator::SUB: case Operator::MUL: case Operator::DIV: case Operator::EQU: case Operator::LSS: case Operator::GTR:
            left = operands[0].compile(l, f);
            right = operands[1].compile(l, f);
            break;

        default: ;
    }

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

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

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

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

        case Operator::EQU:
        return  l.builder.CreateFCmpOEQ(left, right, VARNAME("tmp_eq"));
        break;

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

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

        case Operator::NEG:
        left = operands[0].compile(l, f);
        return  l.builder.CreateFNeg(left, VARNAME("tmp_eq"));
        break;

        case Operator::CALL:
    {
            std::string fname = operands[0].__valueS;
            assert(f->root->__rawFunctions.count(fname));

            llvm::Function* callee = f->root->__rawFunctions[fname];

            std::vector<llvm::Value*> args;
            args.reserve(operands.size()-1);
            for(int i=1; i<operands.size(); ++i)
            {
                args.push_back(operands[i].compile(l, f));
            }

            return  l.builder.CreateCall(callee, args, VARNAME("tmp_call"));
    }
        break;


        case Operator::NONE:
            switch (state)
            {
                case IDENT:
            {
                std::string vname = operands[0].__valueS;
                assert(f->__vartable.count(vname));

                vid vId  =f->__vartable[vname];

                assert(var.size()==0);
                if (f->__rawVars.count(vId))
                {
                    return f->__rawVars[vId];
                }

                Expression& e = f->__declarations[vId];
                llvm::Value* result = e.compile(l, f);
                f->__rawVars[vId] = result;
                return result;
            break;
            }

            case NUMBER:
                double literal = operands[0].__valueD;
                return llvm::ConstantFP::get(llvm::getGlobalContext(), llvm::APFloat(literal));
            };

        break;

    }

    assert(false);
    return 0;
}

AST::AST()
{
}

std::string
AST::getModuleName()
{
    const std::string name = "moduleTest";

    return name;
}

void
AST::compile(LLVMLayer &layer)
{
    layer.module = new llvm::Module(getModuleName(),llvm::getGlobalContext());

    for(Function& f: __functions)
    {
        llvm::Function* rawf = f.compile(layer);
        __rawFunctions[f.getName()] = rawf;
    }
}

void
AST::run(LLVMLayer &l)
{
    llvm::PassManager PM;
    PM.add(llvm::createPrintModulePass(&llvm::outs()));
    PM.run(*l.module);
}

Function::Function(const wstring &name)
{
    char buffer[1000];
    wcstombs(buffer, name.c_str(), 1000);
    __name = buffer;
}

void
Function::addArg(const wstring &vname, const TypeAnnotation &typ)
{
    char buffer[1000];
    wcstombs(buffer, vname.c_str(), 1000);
    std::string name(buffer);

    vid id = registerVar(name);
    __definitions[id] = typ;
    __args.push_back(name);
}

void
Function::setBody(const Expression &body)
{
    __body = body;
}

void
Function::setReturnType(const TypeAnnotation &rtyp)
{
    __retType = rtyp;
}

vid
Function::registerVar(const std::string& vname)
{
    return  ++__vCounter;
}

const std::string&
Function::getName() const
{
    return __name;
}

void
Function::addListDeclaration(const wstring &vname, const TypeAnnotation &typ, const Expression &e)
{}

void
Function::addDeclaration(const wstring &vname, const TypeAnnotation &typ, const Expression &e)
{
    char buffer[1000];
    wcstombs(buffer, vname.c_str(), 1000);

    std::string v(buffer);

    vid id  = registerVar(v);
    __definitions[id] = typ;
    __declarations[id] = e;
}

llvm::Function*
Function::compile(LLVMLayer &l)
{

    std::vector<llvm::Type*> types;
    std::transform(__args.begin(), __args.end(), std::inserter(types, types.end()),
         [this](std::string& arg)     {
           assert(__vartable.count(arg));

           vid argid = __vartable[arg];
           assert(__definitions.count(argid));

           return __definitions[argid].toLLVMType();
    } );

    llvm::FunctionType* ft  = llvm::FunctionType::get(__retType.toLLVMType(), types,false);
    __raw = llvm::cast<llvm::Function>(l.module->getOrInsertFunction(__name,ft));

    llvm::Function::arg_iterator fargsI = __raw->arg_begin();
    for(std::string& arg : __args)
    {
     vid argid = __vartable[arg];

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

    llvm::BasicBlock *block = llvm::BasicBlock::Create(llvm::getGlobalContext(), "entry", __raw);
    l.builder.SetInsertPoint(block);

    l.builder.CreateRet(__body.compile(l, this, 0));
    l.moveToGarbage(ft);

    return __raw;
};