#include "types.h"
#include "parsing.h"


namespace parsing {
    namespace {
        Result<std::unique_ptr<AST::Expression>, std::string> parse_expression(token::TokenStream& stream);

        Result<std::unique_ptr<types::Type>, std::string> parse_type(token::TokenStream& stream) {
            token::TokenStream inner{ stream };
            try {
                auto token = inner.expect(token::Type::Ident);

                stream.m_position = inner.m_position;

                auto ty = new types::FundamentalType{ types::FundamentalTypeKind::Int };
                return new std::unique_ptr<types::Type>{ ty };
            }
            catch (std::runtime_error error) {
                return new std::string{ error.what() };
            }
        }

        Result<std::unique_ptr<AST::Expression>, std::string> parse_plain_expression(token::TokenStream& stream) {
            token::TokenStream inner{ stream };
            try {
                auto token = inner.next();
                if (token.type == token::Type::LiteralInt) {
                    stream.m_position = inner.m_position;

                    auto expr = new AST::IntLiteralExpression{ std::stoi(token.content) };
                    return new std::unique_ptr<AST::Expression>{ expr };
                }
                else if (token.type == token::Type::Ident) {
                    stream.m_position = inner.m_position;

                    auto expr = new AST::ValueReferenceExpression{ token.content };
                    return new std::unique_ptr<AST::Expression>{ expr };
                }
                else {
                    throw std::runtime_error("Expected expression");
                }
            }
            catch (std::runtime_error error) {
                return new std::string{ error.what() };
            }
        }

        Result<std::unique_ptr<AST::Expression>, std::string> parse_primary_expression(token::TokenStream& stream) {
            token::TokenStream inner{ stream };
            try {
                auto plain_expr = parse_plain_expression(inner);
                while (inner.peek().content == "(") {
                    inner.next();

                    std::vector<std::unique_ptr<AST::Expression>> args{};

                    int counter = 0;
                    while (inner.peek().content != ")") {
                        if (counter++ > 0)
                            inner.expect(token::Type::Symbol, ",");
                        args.push_back(parse_expression(inner).unwrap());
                    }

                    inner.expect(token::Type::Symbol, ")");

                    auto fn_call = new AST::FunctionCallExpression{ std::move(plain_expr.unwrap()), std::move(args) };
                    plain_expr = new std::unique_ptr<AST::Expression>{ fn_call };
                }


                stream.m_position = inner.m_position;
                return plain_expr;
            }
            catch (std::runtime_error error) {
                return new std::string{ error.what() };
            }
        }

        Result<types::BinOp, std::string> parse_binop(token::TokenStream& stream) {
            token::TokenStream inner{ stream };
            try {
                auto token = inner.next();
                if (token.type != token::Type::Symbol) {
                    throw std::runtime_error("Expected binop");
                }
                else if (token.content == "=") {
                    stream.m_position = inner.m_position;
                    return new types::BinOp{ types::BinOp::Assignment };
                }
                else if (token.content == "+") {
                    stream.m_position = inner.m_position;
                    return new types::BinOp{ types::BinOp::Add };
                }
                else if (token.content == "-") {
                    stream.m_position = inner.m_position;
                    return new types::BinOp{ types::BinOp::Sub };
                }
                else if (token.content == "<") {
                    stream.m_position = inner.m_position;
                    return new types::BinOp{ types::BinOp::LessThan };
                }
                else if (token.content == ">") {
                    stream.m_position = inner.m_position;
                    return new types::BinOp{ types::BinOp::GreaterThan };
                }

                throw std::runtime_error("Expected binop");
            }
            catch (std::runtime_error error) {
                return new std::string{ error.what() };
            }
        }

        std::unique_ptr<AST::Expression> parse_rhs(
            token::TokenStream& stream, std::unique_ptr<AST::Expression> lhs, int prev_precedence) {

            auto binop_res = parse_binop(stream);
            while (binop_res.ok()) {
                auto binop = binop_res.unwrap();
                auto rhs = parse_primary_expression(stream).unwrap();

                if (types::operator_precedence(binop) > prev_precedence) {
                    rhs = parse_rhs(stream, std::move(rhs), types::operator_precedence(binop));
                }

                auto binop_expr = new AST::BinaryOperationExpression{ std::move(lhs), binop, std::move(rhs) };
                lhs = std::unique_ptr<AST::Expression>{ binop_expr };

                binop_res = parse_binop(stream);
            }

            return lhs;
        }

        Result<std::unique_ptr<AST::Expression>, std::string> parse_expression(token::TokenStream& stream) {
            try {
                auto lhs = parse_primary_expression(stream).unwrap();
                return new std::unique_ptr{ parse_rhs(stream, std::move(lhs), 0) };
            }
            catch (std::runtime_error error) {
                return new std::string{ error.what() };
            }
        }



        Result<std::unique_ptr<AST::InitializationStatement>, std::string> parse_init_statement(token::TokenStream& stream) {
            token::TokenStream inner{ stream };
            try {
                auto ty = parse_type(inner).unwrap();
                auto name = inner.expect(token::Type::Ident);

                std::optional<std::unique_ptr<AST::Expression>> expr = {};
                if (inner.peek().type == token::Type::Symbol && inner.peek().content == "=") {
                    inner.expect(token::Type::Symbol, "=");
                    expr = parse_expression(inner).unwrap();
                }

                inner.expect(token::Type::Symbol, ";");

                stream.m_position = inner.m_position;
                auto init = new AST::InitializationStatement{ std::move(ty), name.content, std::move(expr) };
                return new std::unique_ptr<AST::InitializationStatement>{ init };
            }
            catch (std::runtime_error error) {
                return new std::string{ error.what() };
            }
        }

        Result<std::unique_ptr<AST::Statement>, std::string> parse_statement(token::TokenStream& stream) {
            token::TokenStream inner{ stream };
            try {
                if (inner.peek().type == token::Type::ReturnKeyword) {
                    inner.next();
                    auto expression = parse_expression(inner).unwrap();
                    inner.expect(token::Type::Symbol, ";");

                    stream.m_position = inner.m_position;

                    auto ret = new AST::ReturnStatement{ std::move(expression) };
                    return new std::unique_ptr<AST::Statement>{ ret };
                }
                else if (inner.peek().type == token::Type::IfKeyword) {
                    inner.next();
                    inner.expect(token::Type::Symbol, "(");
                    auto expression = parse_expression(inner).unwrap();
                    inner.expect(token::Type::Symbol, ")");

                    auto then_statement = parse_statement(inner).unwrap();
                    std::optional<std::unique_ptr<AST::Statement>> else_statement{};
                    if (inner.peek().type == token::Type::ElseKeyword) {
                        inner.next();
                        else_statement = parse_statement(inner).unwrap();
                    }

                    stream.m_position = inner.m_position;

                    auto statement = new AST::IfStatement{
                        std::move(expression),
                        std::move(then_statement),
                        std::move(else_statement)
                    };
                    return new std::unique_ptr<AST::Statement>{ statement };
                }
                else if (auto init = parse_init_statement(inner); init.ok()) {
                    stream.m_position = inner.m_position;
                    return new std::unique_ptr<AST::Statement>{ init.unwrap() };
                }
                else if (auto expr = parse_expression(inner); expr.ok()) {
                    stream.m_position = inner.m_position;
                    stream.expect(token::Type::Symbol, ";");
                    auto expr_statement = new AST::ExpressionStatement{ expr.unwrap() };
                    return new std::unique_ptr<AST::Statement>{ expr_statement };
                }
                else {
                    throw std::runtime_error("Expected return-keyword");
                }

            }
            catch (std::runtime_error error) {
                return new std::string{ error.what() };
            }
        }
    }

    Result<std::unique_ptr<AST::TopLevelStatement>, std::string> parse_top_level_statement(token::TokenStream& stream) {
        token::TokenStream inner{ stream };
        try {
            auto type = parse_type(inner).unwrap();
            auto name_token = inner.expect(token::Type::Ident);
            inner.expect(token::Type::Symbol, "(");

            std::vector<std::pair<std::string, std::unique_ptr<types::Type>>> params;
            while (inner.peek().content != ")") {
                if (params.size() > 0) {
                    inner.expect(token::Type::Symbol, ",");
                }
                auto param_ty = parse_type(inner).unwrap();
                auto param_name = inner.expect(token::Type::Ident);
                params.push_back(std::pair(param_name.content, std::move(param_ty)));
            }

            inner.expect(token::Type::Symbol, ")");
            inner.expect(token::Type::Symbol, "{");

            std::vector<std::unique_ptr<AST::Statement>> statements{};

            auto statement = parse_statement(inner);
            while (statement.ok()) {
                statements.push_back(statement.unwrap());
                statement = parse_statement(inner);
            }


            inner.expect(token::Type::Symbol, "}");

            stream.m_position = inner.m_position;

            auto fun = new AST::Function{ std::move(type), std::move(params), name_token.content, std::move(statements) };
            return new std::unique_ptr<AST::TopLevelStatement>{ fun };
        }
        catch (std::runtime_error error) {
            return new std::string(error.what());
        }
    }
}