#include <lexer_base.h>
#include <parser.h>
#include <unordered_map>
#include "mcalc.h"

void State::to_json(nlohmann::json &j) const {
    for (const auto &entry : vars)
        j["vars"].push_back(entry.second.expr);
    j["history"] = history;
}

void State::from_json(const nlohmann::json &j) {
    clear();
    for (const auto &expr : j["vars"]) {
        input_buf = expr;
        handle_input();
    }
    history.clear();
    j.at("history").get_to(history);
    eval();
}


void yy::Parser::error(const location_type &l, const std::string &err) {
    state->error = std::format("Error at {}: {}", l.begin.column, err);
}

void Node::used_vars(std::unordered_set<std::string> *out) const {
    switch (ty) {
        case type::NUM:
            break;
        case type::VAR: {
            out->emplace(var);
            break;
        }
        case type::NEG: {
            a->used_vars(out);
            break;
        }
        case type::PLUS: case type::MINUS: case type::MUL: case type::DIV: {
            a->used_vars(out);
            b->used_vars(out);
            break;
        }
        case type::INVALID: {
            assert(false);
            exit(1);
        }
    }
}

void State::handle_input() {
    std::istringstream input_stream{input_buf};

    error.reset();

    std::pair<std::string, std::unique_ptr<Node>> result;
    Scanner scanner{&input_stream};
    yy::Parser parser{scanner, &result, this};

    if (parser.parse() != 0)
        return;

    Variable new_var_temp {
        .node = std::move(result.second),
        .name = result.first,
        .expr = input_buf
    };
    new_var_temp.node->used_vars(&new_var_temp.depends_on);

    bool old_exists = false;
    auto old_node_it = vars.find(new_var_temp.name);
    Variable* new_var = nullptr;
    if (old_node_it != vars.end()) {
        old_exists = true;
        std::swap(old_node_it->second, new_var_temp);
        new_var = &old_node_it->second;
    } else {
        new_var = &vars.emplace(new_var_temp.name, std::move(new_var_temp)).first->second;
    }

    std::unordered_map<Variable*, Variable*> seen = { { new_var, nullptr } };
    std::vector<Variable*> todo = { new_var };

    while (!todo.empty()) {
        auto var = todo.back();
        todo.pop_back();

        for (const auto& v : var->depends_on) {
            if (v == new_var->name) {
                std::string error_msg = new_var->name;
                auto temp = seen.find(var);
                while (temp != seen.end()) {
                    error_msg = temp->first->name + " -> " + error_msg;
                    temp = seen.find(temp->second);
                }
                error = "Found dependency on self: " + error_msg;
                if (old_exists)
                    std::swap(old_node_it->second, new_var_temp);
                else
                    vars.erase(new_var->name);
                return;
            }
            if (auto child_var = vars.find(v); child_var != vars.end()) {
                if (seen.emplace(&child_var->second, var).second) {
                    todo.push_back(&child_var->second);
                }
            }
        }
    }
    history.push_back(input_buf);
    history_entry = history.size();
    input_buf = "";
}

template<Node::type ty>
std::optional<double> eval_ab(State* s, Node* a, Node* b) {
    if (auto av = a->eval(s)) {
        if (auto bv = b->eval(s)) {
            if constexpr (ty == Node::type::PLUS)
                return (*av) + (*bv);
            else if constexpr (ty == Node::type::MINUS)
                return (*av) - (*bv);
            else if constexpr (ty == Node::type::MUL)
                return (*av) * (*bv);
            else if constexpr (ty == Node::type::DIV)
                return (*av) / (*bv);
            else
                static_assert(false, "Invalid type used for eval_ab");
        }
    }
    return std::nullopt;
}
std::optional<double> Node::eval(State *s) const {
    switch (ty) {
        case type::NUM:
            return num;
        case type::VAR: {
            auto v = s->vars.find(var);
            if (v != s->vars.end())
                return v->second.value;
            return std::nullopt;
        }
        case type::NEG:
            if (auto v = a->eval(s))
                return -(*v);
            return std::nullopt;
        case type::PLUS:
            return eval_ab<type::PLUS>(s, a.get(), b.get());
        case type::MINUS:
            return eval_ab<type::MINUS>(s, a.get(), b.get());
        case type::MUL:
            return eval_ab<type::MUL>(s, a.get(), b.get());
        case type::DIV:
            return eval_ab<type::DIV>(s, a.get(), b.get());
        case type::INVALID:
            return std::nullopt;
    }
    return std::nullopt;
}


void State::eval() {
    std::unordered_set<std::string> todo;
    clock::time_point start = clock::now();

    for (auto &v : vars) {
        v.second.value.reset();

        todo.emplace(v.second.name);
    }

    auto it = todo.begin();
    while (!todo.empty()) {
        if (it == todo.end())
            it = todo.begin();
        auto &var = vars.at(*it);

        bool waiting_for_dep = false;
        for (const auto& dep : var.depends_on) {
            if (todo.contains(dep)) {
                waiting_for_dep = true;
                break;
            }
        }

        if (waiting_for_dep) {
            ++it;
            continue;
        }

        var.value = var.node->eval(this);
        it = todo.erase(it);
    }

    last_eval_duration = clock::now() - start;
}

void State::clear() {
    history.clear();
    vars.clear();
    input_buf.clear();
    error.reset();
    last_eval_duration.reset();
}
