#include "dram_tracer.h"
#include "build_options.h"
#include "dram_tracer_create.h"
#include "drmemtrace/trace_entry.h"

#include <cstdint>
#include <filesystem>
#include <fmt/core.h>
#include <iostream>
#include <tuple>
#include <utility>

analysis_tool_t *dram_tracer_create(const std::string &output_dir)
{
    return new dram_tracer(output_dir);
}

dram_tracer::dram_tracer(std::string_view output_dir) : output_dir(output_dir)
{}

bool dram_tracer::process_memref(const memref_t &memref)
{
    // For online tracing we do not know beforehand how many threads there will be,
    // so we need to keep track of them dynamically.
    if (memref.marker.type == TRACE_TYPE_MARKER && memref.marker.marker_type == TRACE_MARKER_TYPE_TIMESTAMP &&
        memref.marker.tid != 0) {

        tid_t tid = memref.marker.tid;

        if (thread_ids.find(tid) == thread_ids.end()) {
            // Store all thread ids of the process in a set.
            thread_ids.emplace(memref.marker.tid);

            // Create a new trace file for the thread.
            if (output_dir.empty())
                output_dir = fmt::format("dramsys.trace.{}.dir", memref.marker.pid);

            std::filesystem::create_directory(output_dir);

            if constexpr (build_options.binary_output) {
                trace_files.try_emplace(
                    tid, output_dir / fmt::format("dramsys.trace.{}.{}.bin", memref.marker.pid, tid), std::ios::binary);
            } else {
                trace_files.try_emplace(
                    tid, output_dir / fmt::format("dramsys.trace.{}.{}.trace", memref.marker.pid, tid), std::ios::out);

                // Print a header as comment
                trace_files[tid] << "# instruction count,read/write,data size,data address\n"
                                 << "# <timestamp>\n";
            }
        }
    }

    if (type_is_instr(memref.instr.type)) {
        ++instruction_counts[memref.instr.tid];
    } else if (memref.data.type == TRACE_TYPE_READ || memref.data.type == TRACE_TYPE_WRITE) {
        addr_t address = memref.data.addr;
        size_t size = memref.data.size;
        tid_t tid = memref.data.tid;

        // One instruction can issue multiple data accesses, for example a read and a write.
        // Do not substract 1 in this case.
        instr_count_t instruction_count = instruction_counts[memref.data.tid] == 0
                                              ? instruction_counts[memref.data.tid]
                                              : instruction_counts[memref.data.tid] - 1;

        if constexpr (build_options.debug_output) {
            std::string_view access_string = memref.data.type == TRACE_TYPE_READ ? "Read" : "Write";
            fmt::print("[{}] {} at address 0x{:x} with size {} after {} instructions.\n", memref.data.tid,
                       access_string, address, size, instruction_count);
        }

        if constexpr (build_options.binary_output) {
            trace_files[tid] << BinaryTraceEntry(
                BinaryTraceEntry::DataRef(instruction_count, memref.data.type == TRACE_TYPE_WRITE, size, address));
        } else {
            std::string_view access_type = memref.data.type == TRACE_TYPE_READ ? "r" : "w";
            trace_files[tid] << fmt::format("{},{},{},{:x}\n", instruction_count, access_type, size, address);
        }

        instruction_counts[memref.data.tid] = 0;
        ++data_references;
    } else if (memref.marker.type == TRACE_TYPE_MARKER && memref.marker.marker_type == TRACE_MARKER_TYPE_TIMESTAMP) {
        // Each 32KB block of thread data has a timestamp, allowing reconstructing
        // the thread interleaving at that granularity.
        // We place the thread specific timestamp into the trace so that the thread player
        // can wait on other threads to progress at least to this point in time.

        uint64_t timestamp = memref.marker.marker_value;
        tid_t tid = memref.marker.tid;

        if constexpr (build_options.binary_output) {
            trace_files[tid] << BinaryTraceEntry(BinaryTraceEntry::Timestamp(timestamp));
        } else {
            trace_files[tid] << fmt::format("<{}>\n", timestamp);
        }
    }

    return true;
}

bool dram_tracer::print_results()
{
    fmt::print("Traced {} data references in {} threads.\n", data_references, thread_ids.size());
    return true;
}