/**
 * lisp-compiler: Compiles LISP using LLVM.
 * Copyright (C) 2022  Clyne Sullivan
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program 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 General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * @file ast.cpp
 * @brief Abstract Syntax Tree (AST) implementation.
 */

#include "ast.hpp"
#include "state.hpp"

#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/Optional.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Type.h"

#include <algorithm>
#include <any>
#include <cctype>
#include <cstdio>
#include <deque>
#include <iostream>
#include <map>
#include <optional>
#include <string>
#include <typeinfo>
#include <variant>
#include <vector>

AST::Value AST::Identifier::codegen(CompilerState& state)
{
    if (state.namedValues.contains(name))
        return state.namedValues[name];
    else
        return nullptr;
}

AST::Value AST::Literal::codegen(CompilerState& state)
{
    if (type == AST::Literal::Type::Number) {
        if (std::holds_alternative<int>(value))
            return llvm::ConstantInt::get(state.context, llvm::APInt(sizeof(int) * 8, std::get<int>(value)));
        else
            return llvm::ConstantFP::get(state.context, llvm::APFloat(std::get<double>(value)));
    } else {
        return nullptr;
    }
}

AST::Value AST::ProcedureCall::codegen(CompilerState& state)
{
    if (auto id = dynamic_cast<AST::Identifier *>(callee); id) {
        if (state.namedValues.contains(id->name)) {
            auto ptrToFuncPtr = state.namedValues[id->name];
            auto funcPtr = state.builder.CreateLoad(
                ptrToFuncPtr->getType()->getContainedType(0),
                ptrToFuncPtr);
            auto funcType = (llvm::FunctionType *)funcPtr->getType()->getContainedType(0);
            return state.builder.CreateCall(funcType, funcPtr, llvm::None, "calltmp");
        } else {
            std::vector<llvm::Value *> args;
            std::vector<llvm::Type *> argtypes;

            for (auto& a : operands) {
                auto gen = a->codegen(state);
                args.push_back(gen);
                argtypes.push_back(gen ? gen->getType() : llvm::Type::getVoidTy(state.context));
            }


            auto func = state.module.getOrInsertFunction(id->name,
                llvm::FunctionType::get(
                    llvm::Type::getVoidTy(state.context),
                    argtypes,
                    false));
            return state.builder.CreateCall(func, args);
        }

        // work off of id's name
        // builtin?
        // named value? (i.e. is defined)
        return nullptr;
    } else if (auto v = callee->codegen(state); v) {
        // v needs to be a callable procedure
        //std::vector<Node *> operands;
        return nullptr;
    } else {
        return nullptr;
    }
}

AST::Value AST::LambdaExpression::codegen(CompilerState& state)
{
    std::vector<llvm::Type *> args;
    std::vector<std::string> argnames;

    for (auto& op : operands) {
        args.push_back(llvm::Type::getDoubleTy(state.context));
        argnames.push_back(op->name);
    }

    auto ftype = llvm::FunctionType::get(
        llvm::Type::getDoubleTy(state.context), args, false);
    auto func = llvm::Function::Create(
        ftype, llvm::Function::ExternalLinkage, "lambda", &state.module);

    auto n = argnames.cbegin();
    for (auto& a : func->args()) {
        a.setName(*n);
        state.namedValues[*n] = &a;
        ++n;
    }

    auto block = llvm::BasicBlock::Create(state.context, "entry", func);

    auto ip = state.builder.saveIP();
    state.builder.SetInsertPoint(block);
    ++state.scope;

    llvm::Value *ret;
    for (auto& b : body)
        ret = b->codegen(state);

    if (ret)
        state.builder.CreateRet(ret);
    else
        state.builder.CreateRetVoid();

    --state.scope;
    state.builder.restoreIP(ip);

    for (auto& a : argnames)
        state.namedValues.erase(a);

    return func;
}

AST::Value AST::Conditional::codegen(CompilerState& state)
{
    auto cval = state.builder.CreateFCmpONE(
        cond->codegen(state),
        llvm::ConstantFP::get(state.context, llvm::APFloat(0.0)),
        "ifcond");

    auto func = state.builder.GetInsertBlock()->getParent();

    auto bthen = llvm::BasicBlock::Create(state.context, "then", func);
    auto belse = llvm::BasicBlock::Create(state.context, "else", func);
    auto bcont = llvm::BasicBlock::Create(state.context, "cont", func);
    state.builder.CreateCondBr(cval, bthen, belse);

    state.builder.SetInsertPoint(bthen);
    auto vthen = iftrue->codegen(state);
    if (!vthen || vthen->getType() != llvm::Type::getDoubleTy(state.context))
        vthen = llvm::ConstantFP::get(state.context, llvm::APFloat(0.0));
    state.builder.CreateBr(bcont);

    state.builder.SetInsertPoint(belse);
    auto velse = iffalse->codegen(state);
    if (!velse || velse->getType() != llvm::Type::getDoubleTy(state.context))
        velse = llvm::ConstantFP::get(state.context, llvm::APFloat(0.0));
    state.builder.CreateBr(bcont);

    state.builder.SetInsertPoint(bcont);
    auto PN = state.builder.CreatePHI(
        llvm::Type::getDoubleTy(state.context), 2, "iftmp");

    PN->addIncoming(vthen, bthen);
    PN->addIncoming(velse, belse);
    return PN;
}

AST::Value AST::Definition::codegen(CompilerState& state)
{
    if (!state.namedValues.contains(ident->name)) {
        if (state.scope == 0) {
            auto val = (llvm::Constant *)value->codegen(state);

            state.module.getOrInsertGlobal(ident->name, val->getType());

            auto var = state.module.getNamedGlobal(ident->name);
            var->setLinkage(llvm::Function::ExternalLinkage);
            var->setInitializer(val);
            state.namedValues[ident->name] = var;
            return var;
        } else {
            auto alloc = state.builder.CreateAlloca(
                llvm::Type::getDoubleTy(state.context),
                nullptr,
                ident->name);
            state.builder.CreateStore(value->codegen(state), alloc);
            state.namedValues[ident->name] = alloc;
            return alloc;
        }
    } else {
        return nullptr;
    }
}

AST::Value AST::Assignment::codegen(CompilerState& state)
{
    if (state.scope > 0) {
        if (state.namedValues.contains(ident->name)) {
            return state.builder.CreateStore(value->codegen(state), state.namedValues[ident->name]);
        }
    }

    return nullptr;
}