#include "elna/source/parser.hpp"
#include <stdexcept>

namespace elna::source
{
    void empty_visitor::visit(declaration *declaration)
    {
    }

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

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

    void empty_visitor::visit(call_statement *statement)
    {
        statement->arguments().accept(this);
    }

    void empty_visitor::visit(compound_statement *statement)
    {
        for (auto& nested_statement : statement->statements())
        {
            nested_statement->accept(this);
        }
    }

    void empty_visitor::visit(assign_statement *statement)
    {
        statement->rvalue().accept(this);
    }

    void empty_visitor::visit(if_statement *statement)
    {
        statement->prerequisite().accept(this);
        statement->body().accept(this);
    }

    void empty_visitor::visit(while_statement *statement)
    {
        statement->prerequisite().accept(this);
        statement->body().accept(this);
    }

    void empty_visitor::visit(block *block)
    {
        for (const auto& block_definition : block->definitions())
        {
            block_definition->accept(this);
        }
        for (const auto& block_declaration : block->declarations())
        {
            block_declaration->accept(this);
        }
        block->body().accept(this);
    }

    void empty_visitor::visit(program *program)
    {
        visit(dynamic_cast<block *>(program));
    }

    void empty_visitor::visit(binary_expression *expression)
    {
        expression->lhs().accept(this);
        expression->rhs().accept(this);
    }

    void empty_visitor::visit(variable_expression *variable)
    {
    }

    void empty_visitor::visit(integer_literal *number)
    {
    }

    void empty_visitor::visit(boolean_literal *boolean)
    {
    }

    /**
     * AST node.
     */
    void node::accept(parser_visitor *)
    {
    }

    declaration::declaration(const std::string& identifier)
        : m_identifier(identifier)
    {
    }

    void declaration::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    std::string& declaration::identifier() noexcept
    {
        return m_identifier;
    }

    definition::definition(const std::string& identifier)
        : m_identifier(identifier)
    {
    }

    std::string& definition::identifier() noexcept
    {
        return m_identifier;
    }

    constant_definition::constant_definition(const std::string& identifier, std::unique_ptr<integer_literal>&& body)
        : definition(identifier), m_body(std::move(body))
    {
    }

    void constant_definition::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    integer_literal& constant_definition::body()
    {
        return *m_body;
    }

    procedure_definition::procedure_definition(const std::string& identifier, std::unique_ptr<block>&& body)
        : definition(identifier), m_body(std::move(body))
    {
    }

    void procedure_definition::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    block& procedure_definition::body()
    {
        return *m_body;
    }

    block::block(std::vector<std::unique_ptr<definition>>&& definitions,
            std::vector<std::unique_ptr<declaration>>&& declarations,
            std::unique_ptr<statement>&& body)
        : m_definitions(std::move(definitions)),
        m_declarations(std::move(declarations)), m_body(std::move(body)),
        m_table(std::make_shared<symbol_table>())
    {
    }

    void block::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    statement& block::body()
    {
        return *m_body;
    }

    std::vector<std::unique_ptr<definition>>& block::definitions() noexcept
    {
        return m_definitions;
    }

    std::vector<std::unique_ptr<declaration>>& block::declarations() noexcept
    {
        return m_declarations;
    }

    std::shared_ptr<symbol_table> block::table()
    {
        return m_table;
    }

    program::program(std::vector<std::unique_ptr<definition>>&& definitions,
            std::vector<std::unique_ptr<declaration>>&& declarations,
            std::unique_ptr<statement>&& body)
        : block(std::move(definitions), std::move(declarations), std::move(body))
    {
    }

    void program::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    integer_literal::integer_literal(const std::int32_t value)
        : m_number(value)
    {
    }

    void integer_literal::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    std::int32_t integer_literal::number() const noexcept
    {
        return m_number;
    }

    boolean_literal::boolean_literal(const bool value)
        : m_boolean(value)
    {
    }

    void boolean_literal::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    bool boolean_literal::boolean() const noexcept
    {
        return m_boolean;
    }

    variable_expression::variable_expression(const std::string& name)
        : m_name(name)
    {
    }

    void variable_expression::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    const std::string& variable_expression::name() const noexcept
    {
        return m_name;
    }

    binary_expression::binary_expression(std::unique_ptr<expression>&& lhs,
            std::unique_ptr<expression>&& rhs, const unsigned char operation)
        : m_lhs(std::move(lhs)), m_rhs(std::move(rhs))
    {
        switch (operation)
        {
            case '+':
                this->m_operator = binary_operator::sum;
                break;
            case '-': 
                this->m_operator = binary_operator::subtraction;
                break;
            case '*':
                this->m_operator = binary_operator::multiplication;
                break;
            case '/':
                this->m_operator = binary_operator::division;
                break;
            default:
                throw std::logic_error("Invalid binary operator");
        }
    }

    void binary_expression::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    expression& binary_expression::lhs()
    {
        return *m_lhs;
    }

    expression& binary_expression::rhs()
    {
        return *m_rhs;
    }

    binary_operator binary_expression::operation() const noexcept
    {
        return m_operator;
    }

    call_statement::call_statement(const std::string& name, std::unique_ptr<expression>&& body)
        : m_name(name), m_body(std::move(body))
    {
    }

    void call_statement::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    std::string& call_statement::name() noexcept
    {
        return m_name;
    }

    expression& call_statement::arguments()
    {
        return *m_body;
    }

    compound_statement::compound_statement(std::vector<std::unique_ptr<statement>>&& statements)
        : m_statements(std::move(statements))
    {
    }

    void compound_statement::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    std::vector<std::unique_ptr<statement>>& compound_statement::statements()
    {
        return m_statements;
    }

    void assign_statement::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    assign_statement::assign_statement(const std::string& lvalue, std::unique_ptr<expression>&& rvalue)
        : m_lvalue(lvalue), m_rvalue(std::move(rvalue))
    {
    }

    std::string& assign_statement::lvalue() noexcept
    {
        return m_lvalue;
    }

    expression& assign_statement::rvalue()
    {
        return *m_rvalue;
    }

    if_statement::if_statement(std::unique_ptr<expression>&& prerequisite, std::unique_ptr<statement>&& body)
        : m_prerequisite(std::move(prerequisite)), m_body(std::move(body))
    {
    }

    void if_statement::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    expression& if_statement::prerequisite()
    {
        return *m_prerequisite;
    }

    statement& if_statement::body()
    {
        return *m_body;
    }

    while_statement::while_statement(std::unique_ptr<expression>&& prerequisite, std::unique_ptr<statement>&& body)
        : m_prerequisite(std::move(prerequisite)), m_body(std::move(body))
    {
    }

    void while_statement::accept(parser_visitor *visitor)
    {
        visitor->visit(this);
    }

    expression& while_statement::prerequisite()
    {
        return *m_prerequisite;
    }

    statement& while_statement::body()
    {
        return *m_body;
    }

    parser::parser(lexer&& tokens)
        : iterator(std::move(tokens))
    {
    }

    std::unique_ptr<program> parser::parse()
    {
        auto definitions = parse_definitions();
        auto declarations = parse_declarations();
        auto parsed_statement = parse_statement();

        if (parsed_statement == nullptr)
        {
            return nullptr;
        }
        return std::make_unique<program>(std::move(definitions),
                std::move(declarations), std::move(parsed_statement));
    }

    const std::list<std::unique_ptr<error>>& parser::errors() const noexcept
    {
        return iterator.errors();
    }

    std::unique_ptr<expression> parser::parse_factor()
    {
        if (iterator->of() == source::token::type::identifier)
        {
            auto result = std::make_unique<variable_expression>(iterator->identifier());
            ++iterator;
            return result;
        }
        else if (iterator->of() == source::token::token::type::number)
        {
            auto result = std::make_unique<integer_literal>(iterator->number());
            ++iterator;
            return result;
        }
        else if (iterator->of() == source::token::token::type::boolean)
        {
            auto result = std::make_unique<boolean_literal>(iterator->number());
            ++iterator;
            return result;
        }
        else if (iterator->of() == source::token::type::left_paren)
        {
            ++iterator;

            auto expression = parse_expression();

            ++iterator;

            return expression;
        }
        return nullptr;
    }

    std::unique_ptr<expression> parser::parse_term()
    {
        auto lhs = parse_factor();
        if (lhs == nullptr || iterator.current().of() != source::token::type::factor_operator)
        {
            return lhs;
        }
        while (iterator->of() == source::token::type::factor_operator)
        {
            auto _operator = iterator->identifier()[0];
            ++iterator;

            auto rhs = parse_factor();
            lhs = std::make_unique<binary_expression>(std::move(lhs),
                    std::move(rhs), _operator);
        }
        return lhs;
    }

    std::unique_ptr<expression> parser::parse_expression()
    {
        auto term = parse_term();
        if (term == nullptr || iterator.current().of() != source::token::type::term_operator)
        {
            return term;
        }
        while (iterator->of() == source::token::type::term_operator)
        {
            auto _operator = iterator->identifier()[0];
            ++iterator;

            auto rhs = parse_term();
            term = std::make_unique<binary_expression>(std::move(term),
                    std::move(rhs), _operator);
        }
        return term;
    }

    std::unique_ptr<constant_definition> parser::parse_constant_definition()
    {
        auto definition_identifier = iterator.advance(token::type::identifier);

        if (!definition_identifier.has_value())
        {
            return nullptr;
        }
        if (!iterator.skip(token::type::equals))
        {
            return nullptr;
        }

        if (iterator->of() == source::token::type::number)
        {
            auto result = std::make_unique<constant_definition>(definition_identifier.value().get().identifier(),
                    std::make_unique<integer_literal>(iterator->number()));
            ++iterator;
            return result;
        }
        return nullptr;
    }

    std::unique_ptr<declaration> parser::parse_declaration()
    {
        auto declaration_identifier = iterator.advance(token::type::identifier);

        if (!declaration_identifier.has_value() || !iterator.skip(token::type::colon))
        {
            return nullptr;
        }
        auto type_identifier = iterator.advance(token::type::identifier);

        if (!type_identifier.has_value())
        {
            return nullptr;
        }
        return std::make_unique<declaration>(declaration_identifier.value().get().identifier());
    }

    std::unique_ptr<statement> parser::parse_statement()
    {
        if (iterator.look_ahead(token::type::assignment))
        {
            return parse_assign_statement();
        }
        else if (iterator.current(token::type::identifier) && iterator.look_ahead(token::type::left_paren))
        {
            return parse_call_statement();
        }
        else if (iterator.current(token::type::begin))
        {
            return parse_compound_statement();
        }
        else if (iterator.current(token::type::when))
        {
            return parse_if_statement();
        }
        else if (iterator.current(token::type::_while))
        {
            return parse_while_statement();
        }
        iterator.add_error(*iterator);
        return nullptr;
    }

    std::unique_ptr<call_statement> parser::parse_call_statement()
    {
        auto function_name = iterator.advance(token::type::identifier);
        if (function_name.has_value() && !iterator.skip(token::type::left_paren))
        {
            return nullptr;
        }
        auto bang_body = parse_expression();

        if (bang_body != nullptr && iterator.skip(token::type::right_paren))
        {
            return std::make_unique<call_statement>(function_name->get().identifier(), std::move(bang_body));
        }
        return nullptr;
    }

    std::unique_ptr<compound_statement> parser::parse_compound_statement()
    {
        if (!iterator.advance(token::type::begin))
        {
            return nullptr;
        }
        auto result = std::make_unique<compound_statement>();
        std::unique_ptr<statement> next_statement;

        while ((next_statement = parse_statement()) != nullptr)
        {
            result->statements().push_back(std::move(next_statement));

            if (iterator->of() == token::type::semicolon)
            {
                ++iterator;
            }
            else if (iterator->of() == token::type::end)
            {
                ++iterator;
                break;
            }
            else
            {
                iterator.add_error(*iterator);
                break;
            }
        }

        return result;
    }

    std::unique_ptr<assign_statement> parser::parse_assign_statement()
    {
        auto name = iterator.advance(token::type::identifier);
        if (!name.has_value() || !iterator.skip(token::type::assignment))
        {
            return nullptr;
        }
        auto rvalue = parse_expression();

        if (rvalue == nullptr)
        {
            return nullptr;
        }
        return std::make_unique<assign_statement>(name.value().get().identifier(), std::move(rvalue));
    }

    std::unique_ptr<if_statement> parser::parse_if_statement()
    {
        if (!iterator.skip(token::type::when))
        {
            return nullptr;
        }
        auto condition = parse_expression();

        if (condition == nullptr || !iterator.skip(token::type::then))
        {
            return nullptr;
        }
        auto body = parse_statement();

        if (body == nullptr)
        {
            return nullptr;
        }
        return std::make_unique<if_statement>(std::move(condition), std::move(body));
    }

    std::unique_ptr<while_statement> parser::parse_while_statement()
    {
        if (!iterator.skip(token::type::_while))
        {
            return nullptr;
        }
        auto condition = parse_expression();

        if (condition == nullptr || !iterator.skip(token::type::_do))
        {
            return nullptr;
        }
        auto body = parse_statement();

        if (body == nullptr)
        {
            return nullptr;
        }
        return std::make_unique<while_statement>(std::move(condition), std::move(body));
    }

    std::vector<std::unique_ptr<definition>> parser::parse_definitions()
    {
        std::vector<std::unique_ptr<definition>> definitions;

        if (iterator->of() != token::type::let)
        {
            return definitions;
        }
        ++iterator; // Skip const.

        std::unique_ptr<definition> parsed_definition;
        while ((parsed_definition = parse_constant_definition()) != nullptr)
        {
            definitions.push_back(std::move(parsed_definition));

            if (iterator->of() == source::token::type::comma)
            {
                ++iterator;
            }
            else if (iterator->of() == source::token::type::semicolon)
            {
                ++iterator;
                break;
            }
            else
            {
                iterator.add_error(*iterator);
                break;
            }
        }
        return definitions;
    }

    std::vector<std::unique_ptr<declaration>> parser::parse_declarations()
    {
        std::vector<std::unique_ptr<declaration>> declarations;

        if (iterator->of() != token::type::var)
        {
            return declarations;
        }
        ++iterator; // Skip var.

        std::unique_ptr<declaration> parsed_declaration;
        while ((parsed_declaration = parse_declaration()) != nullptr)
        {
            declarations.push_back(std::move(parsed_declaration));

            if (iterator->of() == token::type::comma)
            {
                ++iterator;
            }
            else if (iterator->of() == token::type::semicolon)
            {
                ++iterator;
                break;
            }
            else
            {
                iterator.add_error(*iterator);
                break;
            }
        }
        return declarations;
    }

    std::unique_ptr<block> parser::parse_block()
    {
        auto definitions = parse_definitions();
        auto declarations = parse_declarations();
        auto parsed_statement = parse_statement();

        if (parsed_statement == nullptr)
        {
            return nullptr;
        }
        return std::make_unique<block>(std::move(definitions),
                std::move(declarations), std::move(parsed_statement));
    }
}