#pragma once

#include "alloc.hpp"
#include "arch.hpp"
#include "json.hpp"
#include "value.hpp"
#include <cstddef>
#include <cstdint>
#include <format>
#include <iostream>
#include <string>
#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[this->current].ic_start;
        this->nodes[this->current].acc += diff;
        this->nodes[this->current].ic_start = ic;
        if (node_index == 0)
            return;
        calculate_node_midway_result(ic, this->nodes[this->current].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:
    /// 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_stack_debug;
};

class VM {
public:
    VM(const std::vector<uint32_t>& program, VMOpts opts)
        : opts(opts)
        , program(program.data())
        , program_size(program.size())
    {
    }

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

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

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

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

    inline auto stack_repr_string(size_t max_items) const -> std::string
    {
        auto result = std::string();
        result += "→";
        const auto& stack = view_stack();
        for (size_t i = 0; i < stack.size() and i < max_items; ++i) {
            if (i != 0) {
                result += " ";
            }
            result += std::format(
                "{:<11}", stack[stack.size() - i - 1].to_repr_string());
        }
        if (stack.size() > max_items) {
            result += std::format(" ... + {}", stack.size() - max_items);
        }
        return result;
    }

private:
    void run_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 done() const -> bool
    {
        return not this->halt and this->pc >= this->program_size;
    }

    inline auto fn_stack_at(size_t idx) -> Value&
    {
        return this->stack.at(this->bp + idx);
    }
    inline void assert_fn_stack_has(size_t count)
    {
        if (this->stack.size() - this->bp < count) {
            std::cerr << std::format("stack underflow, pc = {}\n", this->pc);
            std::exit(1);
        }
    }
    inline void assert_stack_has(size_t count)
    {
        if (this->stack.size() < count) {
            std::cerr << std::format("stack underflow, pc = {}\n", this->pc);
            std::exit(1);
        }
    }
    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;
    }

    inline auto assert_program_has(size_t count)
    {
        if (this->pc + count > program_size) {
            std::cerr << std::format("malformed program, pc = {}", this->pc);
            std::exit(1);
        }
    }

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

    bool halt = false;

    FlameGraphBuilder flame_graph;
    CodeCoverageBuilder code_coverage;
};

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

}