#pragma once
#include <unordered_map>
#include "parser/nodes.hpp"
#include "prelude/error.hpp"
#include "prelude/linkedlist.hpp"
#include "prelude/string.hpp"
#include "ir/value.hpp"
#include "ir/ops.hpp"

namespace IR
{

    class IRBuilder
    {
    public:
        IRBuilder(const StringView &filename, const Node *root)
            : m_root(root), m_ops(new DoubleLinkedList<Op*>), m_filename(filename) {}

    public:
        // TODO: support other literals
        ValueHandle *ParseIntLiteral(const IntLiteralNode *literal)
        {
            auto dst = AllocateUnnamed<ConstantInt>(literal->integer());
            return dst;
        }

        void ParseVarDecl(const VarDeclNode *varDecl)
        {
            auto value = ParseExpression(varDecl->value());
            // TODO: gather type information from var decl signature, aka local <int> v = 0;
            auto dst = AllocateNamed<Pointer>(value->GetType());
            m_ops->Append(m_ops->New(new AllocateOp(dst, value->GetType())));
            m_ops->Append(m_ops->New(new StoreOp(value, reinterpret_cast<Pointer *>(dst))));
            m_locals.insert(std::make_pair(varDecl->name(), reinterpret_cast<Pointer *>(dst)));
        }

        ValueHandle *ParseVariable(const VariableNode *var)
        {
            if (m_locals.find(var->name()) == m_locals.end())
            {
                // TODO: pass line:offset when Node will have them
                ErrorLogger::Raise(Error::CompileError(m_filename, StringView::FromFormat("use of undefined variable '%s'", var->name().c_str())));
                assert(false);
            }
            auto dst = AllocateNamed<Instruction>(m_locals[var->name()]->GetValueType());
            m_ops->Append(m_ops->New(new LoadOp(dst, m_locals[var->name()])));
            return reinterpret_cast<ValueHandle *>(dst);
        }

        ValueHandle *ParseFnCall(const FnCallNode *fn)
        {
            auto args = ValueBuilder();
            for (size_t i = 0; i < fn->args().size; ++i)
            {
                auto arg = ParseExpression(fn->args().data[i]);
                args.Push(arg);
            }
            // TODO: gather return type of the function
            auto dst = AllocateNamed<Instruction>(new ValueHandle::Type {ValueHandle::Type::Kind::Int});
            m_ops->Append(m_ops->New(new CallOp(dst, fn->name(), args.view())));
            return dst;
        }

        ValueHandle *ParseFactor(const Node *factor)
        {
            switch (factor->GetType())
            {
            case NodeType::IntLiteral:
                return ParseIntLiteral(reinterpret_cast<const IntLiteralNode *>(factor));
            case NodeType::Variable:
                return ParseVariable(reinterpret_cast<const VariableNode *>(factor));
            case NodeType::FnCall:
                return ParseFnCall(reinterpret_cast<const FnCallNode *>(factor));
            default:
                assert(0 && "some factor may not be handled");
                break;
            }

            assert(0 && "unreachable");
            return nullptr;
        }

        ValueHandle *ParseExpression(const Node *expression)
        {
            if (expression->GetType() == NodeType::Expression)
            {
                auto expr = reinterpret_cast<const ExpressionNode *>(expression);
                auto lhs = ParseExpression(expr->left());
                auto rhs = ParseExpression(expr->right());
                auto dst = AllocateNamed<Instruction>(lhs->GetType());

                assert(4 == static_cast<int>(ExpressionNode::Operator::COUNT_OPERATORS) && "some operators may not be handled");
                switch (expr->op())
                {
                case ExpressionNode::Operator::Plus:
                    m_ops->Append(m_ops->New(new MathOp(dst, lhs, rhs, OpType::ADD)));
                    break;
                case ExpressionNode::Operator::Multiply:
                    m_ops->Append(m_ops->New(new MathOp(dst, lhs, rhs, OpType::MUL)));
                    break;
                case ExpressionNode::Operator::Minus:
                    m_ops->Append(m_ops->New(new MathOp(dst, lhs, rhs, OpType::SUB)));
                    break;
                case ExpressionNode::Operator::Divide:
                    m_ops->Append(m_ops->New(new MathOp(dst, lhs, rhs, OpType::DIV)));
                    break;
                default:
                    assert(0 && "unreachable");
                    break;
                }

                return dst;
            }

            return ParseFactor(expression);
        }

        Block ParseBlock(const CompoundNode *compound)
        {
            StartBlock();
            for (auto &statement : *compound)
            {
                switch (statement->GetType())
                {
                case NodeType::VarDecl:
                    ParseVarDecl(reinterpret_cast<VarDeclNode *>(statement));
                    continue;
                default:
                    ParseExpression(statement);
                    continue;
                }
            }
            auto ops = EndBlock();
            auto block = Block(m_block_counter++, *ops);
            operator delete(ops);
            return block;
        }

        DoubleLinkedList<Op*>* Build()
        {
            assert(m_root->GetType() == NodeType::Program && "root should be a program");
            auto program = reinterpret_cast<const ProgramNode *>(m_root);

            // Externs
            for (auto &extrn : program->externs())
            {
                m_ops->Append(m_ops->New(new ExternOp(extrn->symbol())));
            }

            // Functions
            for (auto &fn : program->funcs())
            {
                auto block = ParseBlock(fn->body());
                m_ops->Append(m_ops->New(new FnOp(fn->name(), fn->params(), std::move(block))));
            }

            return m_ops;
        }

    public:
        const DoubleLinkedList<Op*>* ops() const { return m_ops; }

    private:
        void StartBlock()
        {
            m_containers.Push(m_ops);
            m_ops = new DoubleLinkedList<Op*>;
        }

        DoubleLinkedList<Op*> *EndBlock()
        {
            assert(m_containers.size > 0 && "containers stack is empty");
            auto current = m_ops;
            m_ops = m_containers.data[m_containers.size - 1];
            m_containers.size--;
            return current;
        }

    private:
        template <typename V, typename... Args>
        ValueHandle *AllocateNamed(Args &&...args)
        {
            return new V(++m_value_counter, std::forward<Args>(args)...);
        }

        template <typename V, typename... Args>
        ValueHandle *AllocateUnnamed(Args &&...args)
        {
            return new V(ValueHandle::kNoId, std::forward<Args>(args)...);
        }

    private:
        const Node *m_root = nullptr;
        StringView m_filename;

        DoubleLinkedList<Op*> *m_ops = nullptr;
        unsigned int m_value_counter = 0;
        unsigned int m_block_counter = 0;

        std::unordered_map<StringView, Pointer *> m_locals;

        Builder<DoubleLinkedList<Op*>*> m_containers;
    };

} // namespace IR