#include <cstring>
#include <drmemtrace/analysis_tool.h>
#include <filesystem>
#include <fstream>
#include <set>
#include <unordered_map>

class dram_tracer : public analysis_tool_t
{
public:
    dram_tracer(std::string_view output_dir = "");

    bool process_memref(const memref_t &memref) override;
    bool print_results() override;

private:
    using tid_t = uint64_t;
    using instr_count_t = unsigned int;

    // Used for the binary format.
    union BinaryTraceEntry
    {
        enum class Type
        {
            DataRef,
            Timestamp
        };

        struct DataRef
        {
            const Type type = Type::DataRef;
            instr_count_t instruction_count;
            bool write;
            size_t data_size;
            uintptr_t data_address;

            DataRef(instr_count_t instruction_count, bool write, size_t data_size, uintptr_t data_address)
                : instruction_count(instruction_count), write(write), data_size(data_size), data_address(data_address)
            {}
        } data_ref;

        struct Timestamp
        {
            const Type type = Type::Timestamp;
            uint64_t timestamp;

            Timestamp(uint64_t timestamp) : timestamp(timestamp)
            {}
        } timestamp;

        friend std::ofstream &operator<<(std::ofstream &out, const BinaryTraceEntry &entry)
        {
            out.write((char *)&entry, sizeof entry);
            return out;
        }

        BinaryTraceEntry(DataRef data_ref) : data_ref(data_ref)
        {}
        BinaryTraceEntry(Timestamp timestamp) : timestamp(timestamp)
        {}
    };

    uint64_t data_references = 0;
    std::set<tid_t> thread_ids;
    std::unordered_map<tid_t, instr_count_t> instruction_counts;
    std::unordered_map<tid_t, std::ofstream> trace_files;

    std::filesystem::path output_dir;
};