#include "parser.hpp"
#include <sstream>
#include <stdexcept>
#include <unordered_map>
#include <variant>
#include <ostream>

using TT = Solstice::TokenType;

namespace Solstice {

    std::unordered_map<TokenType, NodeType> tokToNodeType = {
        {TT::Math_Add, NodeType::Add},
        {TT::Math_Subtract, NodeType::Subtract},
        {TT::Math_Multiply, NodeType::Multiply},
        {TT::Math_Divide, NodeType::Divide},
        {TT::Comparison_Equal, NodeType::Equal},
        {TT::Comparison_NotEqual, NodeType::NotEqual},
        {TT::Comparison_GreaterThan, NodeType::GreaterThan},
        {TT::Comparison_LesserThan, NodeType::LesserThan},
        {TT::Assign_Type, NodeType::SetType},
        {TT::Assign_Set, NodeType::Set},
        {TT::Assign_Bind, NodeType::Bind},
    };

    std::optional<Literal> Node::getLiteral() const {
        if (std::holds_alternative<Literal>(data)) {
            return std::get<Literal>(data);
        }
        return {};
    }
    std::optional<std::string> Node::getIdentifier() const {
        if (std::holds_alternative<std::string>(data)) {
            return std::get<std::string>(data);
        }
        return {};
    }

    static inline std::string nodeTypeToString(NodeType type) {
        switch (type) {
            case NodeType::Root:
                return "Root";
            case NodeType::Literal:
                return "Literal";
            case NodeType::Identifier:
                return "Identifier";
            case NodeType::Expression:
                return "Expression";
            case NodeType::Tuple:
                return "Tuple";
            case NodeType::CodeBlock:
                return "CodeBlock";
            case NodeType::FunctionBind:
                return "FunctionBind";
            case NodeType::Bind:
                return "Bind";
            case NodeType::Set:
                return "Set";
            case NodeType::SetType:
                return "SetType";
            case NodeType::Lambda:
                return "Lambda";
            case NodeType::FunctionCall:
                return "FunctionCall";
            case NodeType::Add:
                return "Add";
            case NodeType::Subtract:
                return "Subtract";
            case NodeType::Multiply:
                return "Multiply";
            case NodeType::Divide:
                return "Divide";
            case NodeType::Equal:
                return "Equal";
            case NodeType::NotEqual:
                return "NotEqual";
            case NodeType::GreaterThan:
                return "GreaterThan";
            case NodeType::LesserThan:
                return "LesserThan";
        }
        throw std::runtime_error("FIXME Unhandled case in nodeTypeToString");
    }

    std::ostream& operator<<(std::ostream& stream, const Node& node) {
        stream << node.stringify();
        return stream;
    }

    std::optional<Token> Parser::peek(int64_t ahead) {
        if (current + ahead >= input.size() || current + ahead < 0) {
            return {};
        }
        return input[current + ahead];
    }

    std::optional<Token> Parser::consume() {
        if (current >= input.size()) {
            return {};
        }
        return input[current++];
    }

    std::optional<Node> Parser::getPreviousNode(Node& parent, bool popBack) {
        if (parent.children.empty()) {
            return {};
        }
        Node node = parent.children[parent.children.size() - 1];
        if (popBack) {
            parent.children.pop_back();
        }
        return node;
    }
    Node Parser::parseActionFunction() {
        Node node{NodeType::FunctionBind};

        auto previous = getPreviousNode(*context);
        if (!previous.has_value()) {
            throw std::runtime_error("expecting function arguments before '->'");
        }

        if (previous->type == NodeType::FunctionCall) {
            node.type = NodeType::FunctionBind;
            node.children.push_back(previous->children[0]); // name 
            node.children.push_back(previous->children[1]); // args 
        } else if (previous->type == NodeType::Tuple || previous->type == NodeType::Expression) {
            node.type = NodeType::Lambda;
            previous->type = NodeType::Tuple;
            node.children.push_back(*previous);
        } else {
            throw std::runtime_error("expecting function arguments before '->'");
        }


        Node bodyContext{NodeType::Root};
        Node* prevContext = context;
        context = &bodyContext;

        while (auto n = parseOneNode(Precedence::NewLine)) {
            bodyContext.children.push_back(*n);
        }

        context = prevContext;

        if (bodyContext.children.empty()) {
            throw std::runtime_error("expecting function body after '->'");
        }
        node.children.push_back(bodyContext.children.back());

        return node;
    }

    Node Parser::parseExpr(TT type) {
        auto left = getPreviousNode(*context);
        if (!left.has_value()) {
            throw std::runtime_error("Expecting expression on left of expression");
        }
        Precedence leftPrec = getNodePrecedence(*left);
        Precedence precedence = getTokenPrecedence(Token(type));
        auto right = parseOneNode(precedence);
        if (!right.has_value()) {
            throw std::runtime_error("Expecting expression on right of expression");
        }

        NodeType nodeType;
        if (tokToNodeType.find(type) == tokToNodeType.end()) {
            throw std::runtime_error("FIXME couldn't map token type to node type");
        } else {
            nodeType = tokToNodeType[type];
        }

        if (leftPrec >= precedence) {
            // enclose the left node inside this node 
            return Node(nodeType, {*left, *right});    
        } else {
            // steal the right child of the left node and use it as our left value,
            // then enclose ourself inside the left node
            auto newLeft = left->children[1];
            Node us{nodeType, {newLeft, *right}};

            left->children[1] = us;
            return *left;
        }

    }

    Node Parser::parseOpenParen() {
        Node node{NodeType::Expression};

        for (;;) {
            auto next = peek();
            if (!next.has_value()) {
                throw std::runtime_error("unclosed parens");
            }
            if (next->type == TT::CloseParen) {
                consume();
                if (node.children.size() != 1) {
                    node.type = NodeType::Tuple;
                }
                break;
            }

            Node elem{NodeType::Root};
            Node* prevContext = context;
            context = &elem;

            for (;;) {
                auto t = peek();
                if (!t.has_value() || t->type == TT::CloseParen || t->type == TT::Comma) {
                    break;
                }
                auto n = parseOneNode(Precedence::NewLine);
                if (!n.has_value()) break;
                elem.children.push_back(*n);
            }

            context = prevContext;

            if (elem.children.empty()) {
                throw std::runtime_error("expecting expression in parens");
            }
            node.children.push_back(elem.children.back());

            auto sep = peek();
            if (!sep.has_value()) {
                throw std::runtime_error("unclosed parens");
            }
            if (sep->type == TT::Comma) {
                consume();
            } else if (sep->type != TT::CloseParen) {
                throw std::runtime_error("expecting comma between expressions in parens");
            }
        }

        return node;
    }

    Node Parser::parseOpenCurly() {
        Node node{NodeType::CodeBlock};

        for (;;) {
            while (peek().has_value() && peek()->type == TT::NewLine) {
                consume();
            }
            auto next = peek();
            if (!next.has_value()) {
                throw std::runtime_error("unclosed curly bracket");
            }
            if (next->type == TT::CloseCurly) {
                consume();
                break;
            }

            Node elem{NodeType::Root};
            Node* prevContext = context;
            context = &elem;

            for (;;) {
                auto t = peek();
                if (!t.has_value() || t->type == TT::CloseCurly) {
                    break;
                }
                auto n = parseOneNode(Precedence::NewLine);
                if (!n.has_value()) break;
                elem.children.push_back(*n);
            }

            context = prevContext;

            if (elem.children.empty()) {
                throw std::runtime_error("expecting expression in code block");
            }
            node.children.push_back(elem.children.back());

            auto sep = peek();
            if (!sep.has_value()) {
                throw std::runtime_error("unclosed parens");
            }
        }

        return node;
    }

    Node Parser::parseLiteral() {
        auto current = peek(-1);
        if (!current.has_value()) {
            throw std::runtime_error("FIXME couldn't get current token");
        }
        auto literal = current->getLiteral();
        if (!literal.has_value()) {
            throw std::runtime_error("FIXME token with type literal does not hold a literal");
        }

        return Node(*literal);
    }

    Node Parser::parseIdentifier() {
        auto current = peek(-1);
        if (!current.has_value()) {
            throw std::runtime_error("FIXME couldn't get current token");
        }
        auto id = current->getIdentifier();
        if (!id.has_value()) {
            throw std::runtime_error("FIXME token with type identifier does not hold an identifier");
        }

        // Check if it's a function call
        auto next = peek();
        if (next.has_value() && next->type == TT::OpenParen) {
            consume();
            Node argsNode = parseOpenParen();
            argsNode.type = NodeType::Tuple;

            Node callNode = {NodeType::FunctionCall, {*id, argsNode}};
            return callNode;
        }

        return Node(*id);
    }

    Precedence Parser::getNodePrecedence(const Node& node) {
        switch (node.type) {
            case NodeType::Root:
                return Precedence::Root;
            case NodeType::Literal:
            case NodeType::Identifier:
            case NodeType::Lambda:
            case NodeType::FunctionCall:
            case NodeType::Expression:
            case NodeType::Tuple:
                return Precedence::Identifier;
            case NodeType::FunctionBind:
            case NodeType::Bind:
            case NodeType::Set:
                return Precedence::Set;
            case NodeType::Add:
            case NodeType::Subtract:
                return Precedence::Add;
            case NodeType::Multiply:
            case NodeType::Divide:
                return Precedence::Multiply;
            case NodeType::Equal:
            case NodeType::NotEqual:
            case NodeType::GreaterThan:
            case NodeType::LesserThan:
                return Precedence::Compare;

        }
        throw std::runtime_error("FIXME unhandled node precedence case");
    }

    Precedence Parser::getTokenPrecedence(const Token& token) {
        switch (token.type) {
            case TT::None:
            case TT::Kw_Type:
            case TT::Hash_CImport: 
            case TT::Hash_Effect:
            case TT::Action_Function:
                return Precedence::Other;
            case TT::Math_Add:
            case TT::Math_Subtract:
                return Precedence::Add;
            case TT::Math_Multiply:
            case TT::Math_Divide:
                return Precedence::Multiply;
            case TT::Comparison_Equal:
            case TT::Comparison_NotEqual:
            case TT::Comparison_GreaterThan:
            case TT::Comparison_LesserThan:
                return Precedence::Compare;
            case TT::Assign_Bind:
            case TT::Assign_Set:
            case TT::Assign_Type:
                return Precedence::Set;
            case TT::NewLine:
                return Precedence::NewLine;
            case TT::Identifier:
                return Precedence::Identifier;
            case TT::Literal:
                return Precedence::Identifier;
            default: 
                return Precedence::Other;
        }
    }

    std::optional<Node> Parser::parseOneNode(Precedence precedence) {
        auto next = peek();
        if (!next.has_value()) {
            return {};
        }
        if (getTokenPrecedence(*next) <= precedence) {
            return {};
        }
        consume();
        switch (next->type) {
            case TT::None: break;
            case TT::Identifier:
                return parseIdentifier();
            case TT::Literal: 
                return parseLiteral();
            case TT::Math_Add:
            case TT::Math_Subtract:
            case TT::Math_Multiply:
            case TT::Math_Divide:
            case TT::Comparison_Equal:
            case TT::Comparison_NotEqual:
            case TT::Comparison_GreaterThan:
            case TT::Comparison_LesserThan:
            case TT::Assign_Set:
            case TT::Assign_Bind:
            case TT::Assign_Type:
                return parseExpr(next->type);
            case TT::OpenParen:
                return parseOpenParen();
            case TT::Action_Function:
                return parseActionFunction();
            case TT::OpenCurly:
                return parseOpenCurly();

            case TT::NewLine:
                // ignore new line
                return parseOneNode(precedence);

            // stuff that causes errors 
            case TT::CloseParen:
                throw std::runtime_error("Extra closing paren");
            case TT::CloseSpiky:
                throw std::runtime_error("Extra closing angle bracket");
            case TT::CloseCurly:
                throw std::runtime_error("Extra closing curly brace");
        }
        throw std::runtime_error("FIXME: unimplemented parsing case");
    }

    const Node& Parser::parse() {
        while (auto node = parseOneNode(Precedence::Root)) {
            output.children.push_back(*node);
        }

        return output;
    }

    std::string Node::stringify(int ident) const {
        std::stringstream stream;
        stream << "Node( Type=" << nodeTypeToString(type);
        if (type == NodeType::Literal) {
            auto lit = getLiteral();
            if (lit.has_value()) {
                stream << ", Literal: " << *lit;
            }
        } else if (type == NodeType::Identifier) {
            auto id = getIdentifier();
            if (id.has_value()) {
                stream << ", Identifier: " << *id;
            }
        }

        if (!children.empty()) {
            stream << ", Children: { ";
            bool first = true;
            for (const auto& node : children) {
                if (first) {
                    first = false;
                } else {
                    stream << ",";
                }
                stream << "\n" << std::string(ident, ' ');
                stream << node.stringify(ident + 2);
            }
            int lessIdent = ident - 2;
            if (lessIdent < 0) {
                lessIdent = 0;
            }
            stream << "\n" << std::string(lessIdent, ' ') << "}";
        }
        stream << ")";
        return stream.str();
    }
}