#pragma once

#include "alloc.hpp"
#include "arch.hpp"
#include "json.hpp"
#include "value.hpp"
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <string>
#include <utility>
#include <vector>

namespace sliger {

struct SourcePos {
    int index;
    int line;
    int col;
};

struct FGNode {
    FGNode(uint32_t fn, size_t parent)
        : fn(fn)
        , parent(parent)
    {
    }

    uint32_t fn;
    int64_t acc = 0;
    int64_t ic_start = 0;
    size_t parent;
    // the vector's data may be placed all over the heap. this really really
    // sucks. expect cachemisses when calling infrequently called functions. we
    // might be lucky, that the current function and frequent call paths will be
    // cached, but we have no way to assure this.
    //
    // maybe to fix this, a many-to-many relation table, using one single
    // vector could be used. that would reduce cache misses, in exchange for
    // table lookups.
    std::vector<size_t> children = {};
};

class FlameGraphBuilder : public json::ToAndFromJson {
public:
    inline void report_call(uint32_t fn, int64_t ic_start)
    {
        size_t found = find_or_create_child(fn);
        this->nodes[found].ic_start = ic_start;
        this->current = found;
    }

    inline void report_return(int64_t ic_end)
    {
        int64_t diff = ic_end - this->nodes[this->current].ic_start;
        this->nodes[this->current].acc += diff;
        this->current = this->nodes[this->current].parent;
    }

    inline void calculate_midway_result(int64_t ic)
    {
        calculate_node_midway_result(ic, this->current);
    }

    void to_json(json::Writer& writer) const override;

private:
    inline auto find_or_create_child(uint32_t fn) -> size_t
    {
        auto found_child_index = this->find_child(fn);
        if (found_child_index.has_value())
            return found_child_index.value();
        size_t new_child_index = this->nodes.size();
        this->nodes.push_back(FGNode(fn, this->current));
        this->nodes[this->current].children.push_back(new_child_index);
        return new_child_index;
    }

    inline auto find_child(uint32_t fn) const -> std::optional<size_t>
    {
        for (auto child_idx : this->nodes[this->current].children) {
            if (fn == this->nodes[child_idx].fn) {
                return child_idx;
            }
        }
        return {};
    }

    inline void calculate_node_midway_result(int64_t ic, size_t node_index)
    {
        int64_t diff = ic - this->nodes[node_index].ic_start;
        this->nodes[node_index].acc += diff;
        this->nodes[node_index].ic_start = ic;
        if (node_index == 0)
            return;
        calculate_node_midway_result(ic, this->nodes[node_index].parent);
    }

    void fg_node_to_json(json::Writer& writer, size_t node_index) const;

    std::vector<FGNode> nodes = { FGNode(0, 0) };
    size_t current = 0;
};

struct CCPosEntry {
    SourcePos pos;
    int64_t covers = 0;
};

class CodeCoverageBuilder : public json::ToAndFromJson {
public:
    inline void make_sure_entry_exists(SourcePos pos)
    {
        find_or_create_entry(pos);
    }

    /// call when leaving a source location
    inline void report_cover(SourcePos pos)
    {
        size_t entry_index = find_or_create_entry(pos);
        this->entries[entry_index].covers += 1;
    }

    void to_json(json::Writer& writer) const override;

private:
    inline size_t find_or_create_entry(SourcePos pos)
    {
        if (auto found_index = find_pos_entry(pos); found_index.has_value())
            return found_index.value();
        size_t new_index = this->entries.size();
        this->entries.push_back({ .pos = pos });
        return new_index;
    }

    inline std::optional<size_t> find_pos_entry(SourcePos pos) const
    {
        for (size_t i = 0; i < this->entries.size(); ++i)
            if (this->entries[i].pos.index == pos.index)
                return i;
        return {};
    }

    std::vector<CCPosEntry> entries = {};
};

struct VMOpts {
    bool flame_graph;
    bool code_coverage;
    bool print_debug;
};

class VM {
public:
    VM(std::vector<uint32_t> program, VMOpts opts)
        : opts(opts)
        , program(std::move(program))
    {
    }

    void run_until_done();
    void run_n_instructions(size_t amount);
    void run_instruction();

    inline auto done() const -> bool
    {
        return this->pc >= this->program.size();
    }

    inline auto flame_graph_json() const -> std::string
    {
        return json::to_json(this->flame_graph);
    }

    inline auto code_coverage_json() -> std::string
    {
        for (size_t i = 0; i < this->program.size(); ++i) {
            if (this->program.at(i) == std::to_underlying(Op::SourceMap)
                && this->program.size() - 1 - i >= 3) {
                auto index = static_cast<int32_t>(this->program.at(i + 1));
                auto line = static_cast<int32_t>(this->program.at(i + 2));
                auto col = static_cast<int32_t>(this->program.at(i + 3));
                this->code_coverage.make_sure_entry_exists(
                    { index, line, col });
            }
            i += instruction_size(i);
        }
        return json::to_json(this->code_coverage);
    }

    inline auto view_stack() const -> const std::vector<Value>&
    {
        return this->stack;
    }

    auto stack_repr_string(size_t max_items) const -> std::string;

private:
    void run_builtin(Builtin builtin_id);
    void run_string_builtin(Builtin builtin_id);
    void run_array_builtin(Builtin builtin_id);
    void run_file_builtin(Builtin builtin_id);

    inline void step() { this->pc += 1; }

    inline auto eat_op() -> Op
    {
        auto value = curr_as_op();
        step();
        return value;
    }
    inline auto curr_as_op() const -> Op
    {
        return static_cast<Op>(this->program[this->pc]);
    }

    inline auto eat_int32() -> int32_t
    {
        auto value = curr_as_int32();
        step();
        return value;
    }
    inline auto curr_as_int32() const -> int32_t
    {
        return static_cast<int32_t>(this->program[this->pc]);
    }

    inline auto eat_uint32() -> uint32_t
    {
        auto value = curr_as_uint32();
        step();
        return value;
    }
    inline auto curr_as_uint32() const -> uint32_t
    {
        return this->program[this->pc];
    }

    inline auto fn_stack_at(size_t idx) -> Value&
    {
        return this->stack.at(this->bp + idx);
    }
    void assert_program_has(size_t count);
    void assert_fn_stack_has(size_t count);
    void assert_stack_has(size_t count);
    inline void stack_push(Value&& value) { this->stack.push_back(value); }
    inline void stack_push(Value& value) { this->stack.push_back(value); }
    inline auto stack_pop() -> Value
    {
        auto value = this->stack.at(this->stack.size() - 1);
        this->stack.pop_back();
        return value;
    }
    size_t instruction_size(size_t i) const;

    VMOpts opts;
    uint32_t pc = 0;
    uint32_t bp = 0;
    std::vector<uint32_t> program;
    std::vector<Value> stack;
    std::vector<Value> statics;
    heap::Heap heap;
    SourcePos current_pos = { 0, 1, 1 };
    int64_t instruction_counter = 0;

    int32_t file_id_counter = 3;
    std::unordered_map<int32_t, FILE*> open_files {
        { 0, stdin },
        { 1, stdout },
        { 2, stderr },
    };

    FlameGraphBuilder flame_graph;
    CodeCoverageBuilder code_coverage;
};

auto maybe_op_to_string(uint32_t value) -> std::string;
auto maybe_builtin_to_string(uint32_t value) -> std::string;

}
