/**
 * forspll - LLVM-based Forsp compiler
 * Copyright (C) 2024  Clyne Sullivan <clyne@bitgloo.com>
 *
 * 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/>.
 */

#include <algorithm>
#include <cctype>
#include <iostream>
#include <list>
#include <map>
#include <stack>
#include <string>
#include <string_view>
#include <tuple>

#include "ast.hpp"
#include "llvm.hpp"
#include "parser.hpp"
#include "var.hpp"

llvm::Function *curFunc;
llvm::Value *curEnv;

static LLVMState llvmState;
static std::list<ThunkAST> scope;
static std::stack<std::reference_wrapper<ThunkAST>> callstack;

static bool parseString(std::string_view sv);
static std::list<ThunkAST>::iterator buildThunk(std::list<ThunkAST>::iterator it);

int main()
{
    // 1. Parse code into ThunkASTs, i.e. functions with AST lists for bodies.
    std::string line;
    for (unsigned lineno = 1; std::cin.good(); ++lineno) {
        std::getline(std::cin, line);
        if (!parseString(line)) {
            std::cerr << "    at line " << lineno << std::endl;
            return -1;
        }
    }

    // 2. Traverse through thunks in the order they are given in the file and build IR.
    Var::pushScope();
    for (auto it = scope.begin(); it != scope.end();) {
        it = buildThunk(it);
    }

    // 3. Create main() which initiate global vars and calls the top-level thunk.
    auto envtype = llvm::ArrayType::get(llvmState.inttype, llvmState.envidx);
    auto zerovec = llvm::ConstantVector::get(llvm::ArrayRef(llvmState.zero));
    auto env = new llvm::GlobalVariable(llvmState.modul, envtype, false,
        llvm::GlobalValue::InternalLinkage, zerovec, "env");

    auto func = llvmState.createFunction("main");
    auto entry = llvmState.createEntry(func);
    llvmState.builder.SetInsertPoint(entry);

    auto [t0, _] = Var::lookup("__t0");
    llvmState.builder.CreateCall(llvmState.ftype, t0, llvm::ArrayRef<llvm::Value *> {env});
    llvmState.builder.CreateRetVoid();

    // 4. Output the IR to stdout.
    llvmState.output();
    std::cerr << "envidx: " << llvmState.envidx << std::endl;

    return 0;
}

bool parseString(std::string_view sv)
{
    do {
        const auto [nsv, tok] = nextToken(sv);

        if (tok == Token::none && !nsv.empty()) {
            std::cerr << "error: unknown term: " << nsv << std::endl;
            return false;
        } else {
            std::unique_ptr<BaseAST> expr;

            switch (tok) {
            case Token::ThunkOpen:
                callstack.push(scope.emplace_back());
                break;
            case Token::ThunkClose:
                expr.reset(new PushAST {callstack.top().get().name});
                callstack.pop();
                break;
            case Token::Quote:
                std::cerr << "error: quoting is not supported!" << std::endl;
                return false;
            case Token::PopVar:
                expr.reset(new PopAST {name});
                break;
            case Token::PushVar:
                expr.reset(new PushAST {name});
                break;
            case Token::Var:
                expr.reset(new CallAST {name});
                break;
            case Token::Number:
                expr.reset(new NumberAST {name});
                break;
            case Token::none:
                break;
            }

            if (expr) {
                if (!callstack.empty()) {
                    callstack.top().get().ast.emplace_back().swap(expr);
                } else if (tok != Token::ThunkClose) {
                    std::cerr << "error: non-thunk at top level!" << std::endl;
                    return false;
                }
            }
        }

        sv = nsv;
    } while (!sv.empty());

    return true;
}

std::list<ThunkAST>::iterator buildThunk(std::list<ThunkAST>::iterator it)
{
    auto next = it;
    ++next;

    it->beginGen(llvmState);
    Var::pushScope();

    curFunc = it->func;
    curEnv = it->env;
    llvmState.lastSp = it->lastSp;

    for (auto ait = it->ast.begin(); ait != it->ast.end(); ++ait) {
        auto& a = *ait;

        if (a->name.starts_with("__t")) {
            it->lastSp = llvmState.lastSp;
            next = buildThunk(next);
            curFunc = it->func;
            curEnv = it->env;
            llvmState.lastSp = it->lastSp;
        }

        auto next = ait;
        ++next;
        if (next != it->ast.end() && dynamic_cast<PushAST *>(ait->get()) && dynamic_cast<PopAST *>(next->get())) {
            // Value (e.g. thunk) is immediately pushed then popped into a variable.
            // Shortcut and simply rename the value:
            Var::rename(a->name, (*next)->name);
            it->ast.erase(next);
        } else {
            // Otherwise, just do the regular codegen.
            if (a->codegen(llvmState) == nullptr)
                return scope.end();
        }
    }

    it->endGen(llvmState);

    auto gen = it->codegen(llvmState);
    if (!gen)
        return scope.end();

    Var::popScope();
    Var::addLocal(it->name, Var {gen, true});
    scope.erase(it);
    return next;
}