#include "elna/source/semantic.hpp"
#include "elna/source/result.hpp"
#include <cstdlib>

namespace elna::source
{
    name_analysis_visitor::name_analysis_visitor(std::shared_ptr<symbol_table> table,
            const std::filesystem::path& filename, const std::size_t target_pointer_size)
        : table(table), filename(filename), pointer_size(target_pointer_size)
    {
    }

    void name_analysis_visitor::visit(constant_definition *definition)
    {
        auto constant_type = std::make_shared<const class type>(int_type);
        this->table->enter(definition->identifier(),
                std::make_shared<constant_info>(constant_type, definition->body().number()));
    }

    std::shared_ptr<const type> name_analysis_visitor::convert_declaration_type(const type_expression& ast_type) const
    {
        auto variable_type = std::dynamic_pointer_cast<type_info>(table->lookup(ast_type.base()))
            ->type();
        std::shared_ptr<type> declaration_type;

        if (ast_type.is_pointer())
        {
            return std::make_shared<pointer_type>(variable_type, 4);
        }
        else
        {
            return variable_type;
        }
    }

    void name_analysis_visitor::visit(declaration *declaration)
    {
        std::shared_ptr<const type> declaration_type = convert_declaration_type(declaration->type());

        this->table->enter(declaration->identifier(),
                std::make_shared<variable_info>(declaration_type));
    }

    void name_analysis_visitor::visit(program *program)
    {
        class procedure_type main_type{ std::vector<std::shared_ptr<const class type>>(), this->pointer_size };
        this->table->enter("_start", std::make_shared<procedure_info>(main_type, this->table));
        empty_visitor::visit(program);
    }

    void name_analysis_visitor::visit(procedure_definition *procedure)
    {
        std::vector<std::shared_ptr<const type>> arguments;

        for (auto& parameter : procedure->parameters())
        {
            auto declaration_type = convert_declaration_type(parameter->type());
            arguments.push_back(declaration_type);
        }
        procedure_type definition_type{ std::move(arguments), this->pointer_size };
        auto info = std::make_shared<procedure_info>(definition_type, this->table);

        this->table->enter(procedure->identifier(), info);
        this->table = info->scope();

        for (std::size_t i = 0; i < procedure->parameters().size(); ++i)
        {
            this->table->enter(procedure->parameters().at(i)->identifier(),
                    std::make_shared<parameter_info>(definition_type.arguments.at(i)));
        }
        procedure->body().accept(this);

        this->table = info->scope()->scope();
    }

    const std::list<std::unique_ptr<error>>& name_analysis_visitor::errors() const noexcept
    {
        return m_errors;
    }

    allocator_visitor::allocator_visitor(std::shared_ptr<symbol_table> table)
        : table(table)
    {
    }

    void allocator_visitor::visit(declaration *declaration)
    {
        auto declaration_info = this->table->lookup(declaration->identifier());

        if (auto variable = std::dynamic_pointer_cast<variable_info>(declaration_info))
        {
            this->local_offset -= sizeof(std::int32_t);
            variable->offset = this->local_offset;
        }
        else if (auto parameter = std::dynamic_pointer_cast<parameter_info>(declaration_info))
        {
            parameter->offset = this->argument_offset;
            this->argument_offset += sizeof(std::int32_t);
        }
    }

    void allocator_visitor::visit(program *program)
    {
        this->local_offset = 0;
        this->argument_offset = 0;

        empty_visitor::visit(program);
        std::dynamic_pointer_cast<procedure_info>(table->lookup("_start"))->local_stack_size =
            std::abs(this->local_offset);
    }

    void allocator_visitor::visit(procedure_definition *procedure)
    {
        this->local_offset = 0;
        this->argument_offset = 0;
        auto info = std::dynamic_pointer_cast<procedure_info>(this->table->lookup(procedure->identifier()));
        this->table = info->scope();

        empty_visitor::visit(procedure);

        this->table = info->scope()->scope();
        info->local_stack_size = std::abs(this->local_offset);
        info->argument_stack_size = this->argument_offset;
    }

    void allocator_visitor::visit(call_statement *statement)
    {
        auto call_info = std::dynamic_pointer_cast<intrinsic_info>(this->table->lookup(statement->name()));

        this->argument_offset = std::max(static_cast<std::size_t>(this->argument_offset),
                call_info->parameter_stack_size());
    }

    type_analysis_visitor::type_analysis_visitor(std::shared_ptr<symbol_table> table,
            const std::filesystem::path& filename, const std::size_t target_pointer_size)
        : table(table), filename(filename), pointer_size(target_pointer_size)
    {
    }

    void type_analysis_visitor::visit(program *program)
    {
        for (auto& definition : program->definitions())
        {
            if (dynamic_cast<procedure_definition *>(definition.get()) != nullptr)
            {
                definition->accept(this);
            }
        }
        program->body().accept(this);
    }

    void type_analysis_visitor::visit(procedure_definition *procedure)
    {
        auto info = std::dynamic_pointer_cast<procedure_info>(this->table->lookup(procedure->identifier()));
        this->table = info->scope();

        procedure->body().accept(this);

        this->table = info->scope()->scope();
    }

    void type_analysis_visitor::visit(integer_literal *literal)
    {
        literal->data_type = std::dynamic_pointer_cast<type_info>(table->lookup("Int"))->type();
    }

    void type_analysis_visitor::visit(boolean_literal *literal)
    {
        literal->data_type = std::dynamic_pointer_cast<type_info>(table->lookup("Boolean"))->type();
    }

    void type_analysis_visitor::visit(variable_expression *expression)
    {
        expression->data_type = std::dynamic_pointer_cast<typed_info>(table->lookup(expression->name()))->type();
    }

    void type_analysis_visitor::visit(unary_expression *expression)
    {
        empty_visitor::visit(expression);

        switch (expression->operation())
        {
            case unary_operator::reference:
                expression->data_type = std::make_shared<const pointer_type>(expression->operand().data_type,
                        this->pointer_size);
                break;
            case unary_operator::dereference:
                auto operand_type = expression->operand().data_type;

                if (auto referenced_type = std::dynamic_pointer_cast<const pointer_type>(operand_type))
                {
                    expression->data_type = referenced_type;
                }
                else if (operand_type != nullptr)
                {
                    auto new_error = std::make_unique<type_mismatch>(operand_type,
                            type_mismatch::operation::dereference, this->filename, expression->position());
                    m_errors.push_back(std::move(new_error));
                }
                break;
        }
    }

    void type_analysis_visitor::visit(binary_expression *expression)
    {
        empty_visitor::visit(expression);

        switch (expression->operation())
        {
            case binary_operator::sum:
            case binary_operator::subtraction:
            case binary_operator::multiplication:
            case binary_operator::division:
            case binary_operator::less:
            case binary_operator::greater:
            case binary_operator::less_equal:
            case binary_operator::greater_equal:
                if (expression->lhs().data_type != nullptr && expression->rhs().data_type != nullptr)
                {
                    auto lhs_type = std::dynamic_pointer_cast<const primitive_type>(expression->lhs().data_type);
                    auto rhs_type = std::dynamic_pointer_cast<const primitive_type>(expression->rhs().data_type);

                    std::unique_ptr<type_mismatch> new_error;
                    if (lhs_type == nullptr || *lhs_type != int_type)
                    {
                        new_error = std::make_unique<type_mismatch>(lhs_type,
                                type_mismatch::operation::arithmetic, this->filename, expression->lhs().position());
                    }
                    if (rhs_type == nullptr || *rhs_type != int_type)
                    {
                        new_error = std::make_unique<type_mismatch>(rhs_type,
                                type_mismatch::operation::arithmetic, this->filename, expression->rhs().position());
                    }
                    if (new_error != nullptr)
                    {
                        m_errors.push_back(std::move(new_error));
                    }
                }
            break;
            case binary_operator::equals:
            case binary_operator::not_equals:
                if (expression->lhs().data_type != nullptr && expression->rhs().data_type != nullptr)
                {
                    if (expression->lhs().data_type != expression->rhs().data_type)
                    {
                        auto new_error = std::make_unique<type_mismatch>(expression->rhs().data_type,
                                type_mismatch::operation::comparison, this->filename, expression->rhs().position());
                    }
                }
            break;
        }
    }

    void type_analysis_visitor::visit(call_statement *statement)
    {
        auto call_info = std::dynamic_pointer_cast<intrinsic_info>(this->table->lookup(statement->name()));

        std::size_t i{ 0 };
        for (const auto& argument : statement->arguments())
        {
            argument->accept(this);

            if (argument->data_type != nullptr && i < call_info->type()->arguments.size()
                    && call_info->type()->arguments[i] != argument->data_type)
            {
                auto new_error = std::make_unique<type_mismatch>(argument->data_type,
                        type_mismatch::operation::argument, this->filename, argument->position());
                m_errors.push_back(std::move(new_error));
            }

            ++i;
        }
    }

    void type_analysis_visitor::visit(constant_definition *definition)
    {
        definition->body().accept(this);
    }

    void type_analysis_visitor::visit(while_statement *statement)
    {
        statement->prerequisite().accept(this);
        auto condition_type = std::dynamic_pointer_cast<const primitive_type>(statement->prerequisite().data_type);

        if (condition_type != nullptr && *condition_type != boolean_type)
        {
            auto new_error = std::make_unique<type_mismatch>(condition_type,
                    type_mismatch::operation::condition, this->filename, statement->prerequisite().position());
            m_errors.push_back(std::move(new_error));
        }
        statement->body().accept(this);
    }

    void type_analysis_visitor::visit(if_statement *statement)
    {
        statement->prerequisite().accept(this);
        auto condition_type = std::dynamic_pointer_cast<const primitive_type>(statement->prerequisite().data_type);

        if (condition_type != nullptr && *condition_type != boolean_type)
        {
            auto new_error = std::make_unique<type_mismatch>(condition_type,
                    type_mismatch::operation::condition, this->filename, statement->prerequisite().position());
            m_errors.push_back(std::move(new_error));
        }
        statement->body().accept(this);
    }

    void type_analysis_visitor::visit(assign_statement *statement)
    {
        statement->rvalue().accept(this);
        auto lvalue_info = std::dynamic_pointer_cast<typed_info>(this->table->lookup(statement->lvalue()));

        if (statement->rvalue().data_type != nullptr && lvalue_info->type() == statement->rvalue().data_type)
        {
            auto new_error = std::make_unique<type_mismatch>(statement->rvalue().data_type,
                    type_mismatch::operation::assignment, this->filename, statement->position());
            m_errors.push_back(std::move(new_error));
        }
    }

    const std::list<std::unique_ptr<error>>& type_analysis_visitor::errors() const noexcept
    {
        return m_errors;
    }
}