#include "Parser.hpp"
#include "lib/Node.hpp"
#include <memory>

namespace roza
{
  /*explicit*/ Parser::Parser(Lexer& lexer, StatusLog& log)
    : m_lexer { lexer }
    , m_log { log }
  {
  }

  /*virtual*/ Parser::~Parser()
  {
  }

  std::shared_ptr<Node> Parser::parse()
  {
    m_cursor = 0;
    auto root = parse_prog();

    if (m_cursor < m_lexer.size())
      {
        m_log.fatal(m_lexer.loc(), "unexpected end");
      }

    return root;
  }

  std::shared_ptr<Node> Parser::node(size_t offset) const
  {
    return m_lexer.get_or_nullptr(m_cursor + offset);
  }

  NodeType Parser::type(size_t offset) const
  {
    auto tok = m_lexer.get_or_nullptr(m_cursor + offset);
    assert(tok);
    return tok->type();
  }

  bool Parser::type_is(NodeType type, size_t offset) const
  {
    auto tok = m_lexer.get_or_nullptr(m_cursor + offset);
    if (!tok) { return false; }

    return tok->type() == type;
  }

  std::shared_ptr<Node> Parser::consume_or_nullptr(NodeType ty)
  {
    if (m_cursor >= m_lexer.size())
      {
        return nullptr;
      }

    return consume(ty);
  }

  std::shared_ptr<Node> Parser::consume(NodeType ty)
  {
    if (!type_is(ty))
      {
        m_log.fatal(node()->loc(),
                    std::string()
                    + "syntax error, expected '"
                    + NodeTypeStr[ty]
                    + "', got '"
                    + NodeTypeStr[type()]
                    + "'");

        return nullptr;
      }

    auto ret = node();
    next();

    return ret;
  }

  void Parser::consume_all(NodeType ty)
  {
    while (type_is(ty))
      {
        next();
      }
  }

  std::shared_ptr<Node> Parser::consume()
  {
    auto ret = node();
    next();

    return ret;
  }

  void Parser::ensure(NodeType type)
  {
    consume_or_nullptr(type);
    consume_all(type);
  }

  void Parser::next()
  {
    m_cursor++;
  }

  std::shared_ptr<Node> Parser::parse_prog()
  {
    if (m_lexer.size() == 0)
      {
        return std::make_shared<Node>(NODE_PROG, "", m_lexer.loc());
      }

    auto root = std::make_shared<Node>(NODE_PROG, "", node()->loc());
    root->add_child(parse_instr());

    while (m_cursor < m_lexer.size())
      {
        root->add_child(parse_instr());
      }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_instr()
  {
    consume_all(NODE_EOI);

    if (type_is(NODE_ASSERT))
      {
        auto root = consume();
        root->add_child(parse_expr());

        ensure(NODE_EOI);
        return root;
      }
    else if (type_is(NODE_ASSERT_STATIC_FAIL))
      {
        auto root = consume();
        root->add_child(parse_instr());
        return root;
      }
    else if (type_is(NODE_RETURN))
      {
        auto root = consume();
        root->add_child(parse_expr());
        ensure(NODE_EOI);
        return root;
      }
    else if (type_is(NODE_LET))
      {
        auto root = parse_constdecl();
        ensure(NODE_EOI);
        return root;
      }
    else if (type_is(NODE_LET_MUT))
      {
        auto root = parse_vardecl();
        ensure(NODE_EOI);
        return root;
      }
    else if (type_is(NODE_IDENT) && type_is(NODE_ASSIGN, 1))
      {
        auto root = parse_assign();
        ensure(NODE_EOI);
        return root;
      }
    else if (type_is(NODE_IF))
      {
        auto root = parse_if();
        consume_all(NODE_EOI);
        return root;
      }
    else
      {
        auto root = parse_expr();
        ensure(NODE_EOI);
        return root;
      }

    m_log.fatal(node()->loc(),
                std::string()
                + "unknown instruction '"
                + std::string(NodeTypeStr[node()->type()])
                .substr(std::string("NODE_").size())
                + "'");

    return nullptr;
  }

  std::shared_ptr<Node> Parser::parse_vardecl()
  {
    consume(NODE_LET_MUT);
    auto root = std::make_shared<Node>(NODE_VARDECL, "", node()->loc());
    auto ident = consume(NODE_IDENT);
    consume(NODE_ASSIGN);
    auto expr = parse_expr();
    root->add_child(ident);
    root->add_child(expr);

    return root;
  }

  std::shared_ptr<Node> Parser::parse_constdecl()
  {
    consume(NODE_LET);
    auto root = std::make_shared<Node>(NODE_CONSTDECL, "", node()->loc());
    auto ident = consume(NODE_IDENT);
    consume(NODE_ASSIGN);
    auto expr = parse_expr();
    root->add_child(ident);
    root->add_child(expr);

    return root;
  }

  std::shared_ptr<Node> Parser::parse_assign()
  {
    auto lhs = consume(NODE_IDENT);
    auto root = consume(NODE_ASSIGN);

    root->add_child(lhs);
    root->add_child(parse_expr());

    return root;
  }

  std::shared_ptr<Node> Parser::parse_if()
  {
    auto root = consume(NODE_IF);
    root->add_child(parse_expr());
    root->add_child(parse_then());

    if (type_is(NODE_ELSE)
        && type_is(NODE_IF, 1))
      {
        consume(NODE_ELSE);
        root->add_child(parse_if());
        return root;
      }

    if (type_is(NODE_ELSE))
      {
        root->add_child(parse_else());
      }

    consume(NODE_END);

    return root;
  }

  std::shared_ptr<Node> Parser::parse_then()
  {
    auto root = std::make_shared<Node>(NODE_THEN, "", node()->loc());

    while (!type_is(NODE_ELSE)
           && !type_is(NODE_END))
      {
        root->add_child(parse_instr());
      }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_else()
  {
    auto root = consume(NODE_ELSE);

    while (!type_is(NODE_END))
      {
        root->add_child(parse_instr());
      }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_expr()
  {
    return parse_imp();
  }

  std::shared_ptr<Node> Parser::parse_imp()
  {
    auto lhs = parse_or();

    if (type_is(NODE_IMP))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_or());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_or()
  {
    auto lhs = parse_and();

    while (type_is(NODE_OR))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_and());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_and()
  {
    auto lhs = parse_eq();

    while (type_is(NODE_AND))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_eq());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_eq()
  {
    auto lhs = parse_cmp();

    if (type_is(NODE_EQ)
        || type_is(NODE_NE))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_cmp());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_cmp()
  {
    auto lhs = parse_term();

    if (type_is(NODE_LT)
        || type_is(NODE_LE)
        || type_is(NODE_GT)
        || type_is(NODE_GE))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_term());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_term()
  {
    auto lhs = parse_factor();

    while (type_is(NODE_ADD) || type_is(NODE_SUB))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_factor());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_factor()
  {
    auto lhs = parse_unop();

    while (type_is(NODE_MUL)
           || type_is(NODE_DIV)
           || type_is(NODE_MOD))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_unop());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_unop()
  {
    if (type_is(NODE_ADD))
      {
        auto root = std::make_shared<Node>(NODE_UADD, "", m_lexer.loc());
        next();

        root->add_child(parse_pow());
        return root;
      }

    if (type_is(NODE_SUB))
      {
        auto root = std::make_shared<Node>(NODE_USUB, "", m_lexer.loc());
        next();

        root->add_child(parse_pow());
        return root;
      }

    if (type_is(NODE_NOT))
      {
        auto root = std::make_shared<Node>(NODE_NOT, "", m_lexer.loc());
        next();

        root->add_child(parse_pow());
        return root;
      }

    return parse_pow();
  }

  std::shared_ptr<Node> Parser::parse_pow()
  {
    auto lhs = parse_group();

    if (type_is(NODE_POW))
      {
        auto root = consume();
        root->add_child(lhs);
        root->add_child(parse_group());
        lhs = root;
      }

    return lhs;
  }

  std::shared_ptr<Node> Parser::parse_group()
  {
    if (type_is(NODE_OPAR))
      {
        consume(NODE_OPAR);
        auto root = parse_expr();
        consume(NODE_CPAR);
        return root;
      }

    return parse_base();
  }

  std::shared_ptr<Node> Parser::parse_base()
  {
    if (type_is(NODE_INT)
        || type_is(NODE_BOOL)
        || type_is(NODE_IDENT))
      {
        return consume();
      }

    if (type_is(NODE_FUN))
      {
        return parse_fun();
      }

    if (type_is(NODE_OBRACE))
      {
        return parse_call();
      }

    m_log.fatal(node()->loc(),
                "cannot parse unknown node '"
                + node()->string()
                + "'");

    return nullptr;
  }

    std::shared_ptr<Node> Parser::parse_call()
  {
    auto root = std::make_shared<Node>(NODE_CALL, "", node()->loc());
    consume(NODE_OBRACE);
    root->add_child(consume(NODE_IDENT));
    root->add_child(parse_args());
    consume(NODE_CBRACE);

    return root;
  }

  std::shared_ptr<Node> Parser::parse_args()
  {
    auto root = std::make_shared<Node>(NODE_ARGS, "", node()->loc());

    while (!type_is(NODE_CBRACE))
      {
        root->add_child(parse_expr());
      }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_int()
  {
    auto root = consume(NODE_INT);
    return root;
  }

  std::shared_ptr<Node> Parser::parse_fun()
  {
    auto root = consume(NODE_FUN);
    consume(NODE_OPAR);
    root->add_child(parse_params());
    consume(NODE_CPAR);

    root->add_child(parse_ret());
    root->add_child(parse_body());

    consume(NODE_END);

    return root;
  }

  std::shared_ptr<Node> Parser::parse_params()
  {
    auto root = std::make_shared<Node>(NODE_PARAMS, "", m_lexer.loc());

    if (type_is(NODE_CPAR))
      {
        return root;
      }

    root->add_child(parse_param());

    while (type_is(NODE_COMMA))
      {
        consume();
        root->add_child(parse_param());
      }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_param()
  {
    auto root = std::make_shared<Node>(NODE_PARAM, "", m_lexer.loc());
    root->add_child(consume(NODE_IDENT));

    if(type_is(NODE_AS))
    {
      consume();
      root->add_child(parse_type());
    }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_ret()
  {
    auto root = std::make_shared<Node>(NODE_RET, "", m_lexer.loc());

    if (type_is(NODE_ARROW))
      {
        consume();
        root->add_child(parse_type());
      }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_body()
  {
    auto root = std::make_shared<Node>(NODE_BODY, "", m_lexer.loc());

    while (!type_is(NODE_END))
      {
        root->add_child(parse_instr());
      }

    return root;
  }

  std::shared_ptr<Node> Parser::parse_type()
  {
    if (type_is(NODE_FUN))
      {
        auto root = std::make_shared<Node>(NODE_TYPE, "fun", node()->loc());
        consume();

        consume(NODE_LT);

        while (!type_is(NODE_GT))
          {
            root->add_child(parse_type());

            if (type_is(NODE_ARROW))
              {
                consume();
              }
          }

        consume(NODE_GT);

        return root;
      }

    return consume(NODE_TYPE);
  }
}