#pragma once

#include <array>
#include <cstdint>

#include "memory/dram.h"
#include "memory/line.h"

namespace memory {
    template<unsigned int SETS_LOG2, unsigned int WAYS> struct inline_cache {
        static constexpr std::uint64_t SETS = 1 << SETS_LOG2;
        static constexpr std::uint64_t LINE_SET_OFFSET_MASK = SETS - 1;

        std::uint64_t last_serial = 0;

        struct tag {
            std::uint64_t serial = 0;
            std::uint64_t line_address;
            infra::transaction transaction;
        };

        typedef std::array<tag, WAYS> set_tags;
        typedef std::array<line, WAYS> set_data;

        std::array<set_tags, SETS> tags;
        std::array<set_data, SETS> data;

        void handle_response(const dram::response &r) {
            auto set_address = r.line_address & LINE_SET_OFFSET_MASK;
            auto &stags = tags[set_address];
            auto &sdata = data[set_address];
            for (unsigned int i = 0; i < WAYS; ++i) {
                auto &stag = stags[i];
                if (stag.serial && stag.line_address == r.line_address) {
                    handle_response(r, stag);
                    return;
                }
            }
            std::uint64_t min_serial = ~(std::uint64_t)0;
            unsigned int victim;
            for (unsigned int i = 0; i < WAYS; ++i) {
                auto &stag = stags[i];
                if (stag.serial < min_serial) {
                    min_serial = stag.serial;
                    victim = i;
                }
            }
            handle_response(r, stags[victim], sdata[victim]);
        }

        void handle_response(const dram::response &r, tag &t) {
            t.serial = ++last_serial;
            t.transaction = r.transaction;
        }

        void handle_response(const dram::response &r, tag &t, line &d) {
            handle_response(r, t);
            t.line_address = r.line_address;
            d = r.data;
        }

        std::optional<tag> probe(std::uint64_t address) {
            auto line_address = address >> LINE_BYTES_LOG2;
            auto set_address = line_address & LINE_SET_OFFSET_MASK;
            auto &stags = tags[set_address];
            for (unsigned int i = 0; i < WAYS; ++i) {
                auto &stag = stags[i];
                if (stag.serial && stag.line_address == line_address)
                    return stag;
            }
            return {};
        }

        std::optional<tag> fetchline(line &dst, std::uint64_t address) {
            auto line_address = address >> LINE_BYTES_LOG2;
            auto set_address = line_address & LINE_SET_OFFSET_MASK;
            auto &stags = tags[set_address];
            auto &sdata = data[set_address];
            for (unsigned int i = 0; i < WAYS; ++i) {
                auto &stag = stags[i];
                if (stag.serial && stag.line_address == line_address) {
                    stag.serial = ++last_serial;
                    dst = sdata[i];
                    return stag;
                }
            }
            return {};
        }

        std::optional<tag> fetch(unsigned int &dst, std::uint64_t address) {
            line data;
            if (auto tag = fetchline(data, address); tag.has_value()) {
                auto line_offset = address & LINE_BYTE_OFFSET_MASK;
                dst = data[line_offset];
                return tag;
            }
            return {};
        }
    };
}