#pragma once

#include <array>
#include <cstdint>
#include <cstring>
#include <istream>
#include <map>
#include <string>
#include <utility>

#include "memory/line.h"

namespace memory {
    struct dram : public infra::sim {
        static constexpr std::uint64_t PAGE_LINES_LOG2 = 20 - LINE_BYTES_LOG2;
        static constexpr std::uint64_t PAGE_LINES = 1 << PAGE_LINES_LOG2;
        static constexpr std::uint64_t PAGE_LINE_OFFSET_MASK = PAGE_LINES - 1;
        static constexpr std::uint64_t PAGE_BYTES_LOG2 = PAGE_LINES_LOG2 + LINE_BYTES_LOG2;
        static constexpr std::uint64_t PAGE_BYTES = 1 << PAGE_BYTES_LOG2;
        static constexpr std::uint64_t PAGE_BYTE_OFFSET_MASK = PAGE_BYTES - 1;

        typedef std::array<line, PAGE_LINES> page;

        const unsigned int latency;

        std::map<std::uint64_t, page> image;

        struct response {
            infra::transaction transaction;
            std::uint64_t line_address;
            line data;
        };

        struct command {
            infra::transaction transaction;
            std::uint64_t line_address;
            line data;
            std::array<bool, LINE_BYTES> mask;
            bool write = false;
            infra::port<response> *responsep = nullptr;
        };

        infra::port<command> *commandp = nullptr;

        unsigned int progress = 0;
        std::optional<command> current_command;

        dram(unsigned int latency=0)
            : latency(latency)
        { }

        void clock() {
            if (!current_command.has_value() && commandp->can_read()) {
                current_command = commandp->read();
                progress = latency;
            }
            if (current_command.has_value()) {
                if (progress) {
                    --progress;
                    return;
                }
                const auto &c = *current_command;
                if (!c.responsep || c.responsep->can_write()) {
                    auto page_address = c.line_address >> PAGE_LINES_LOG2;
                    auto page_line = c.line_address & PAGE_LINE_OFFSET_MASK;
                    if (c.write) {
                        pte(c.transaction, "s", fmt::format("store {:x}-{:x}", page_address, page_line));
                        if (c.responsep) {
                            response r;
                            r.transaction = c.transaction;
                            r.line_address = c.line_address;
                            r.data = c.data;
                            c.responsep->write(std::move(r));
                        }
                        auto [p, emplaced] = image.try_emplace(page_address);
                        if (emplaced)
                            for (unsigned int i = 0; i < PAGE_LINES; ++i)
                                p->second[i].fill(0);
                        auto &l = p->second[page_line];
                        for (unsigned int i = 0; i < LINE_BYTES; ++i)
                            if (c.mask[i])
                                l[i] = c.data[i];
                    } else {
                        pte(c.transaction, "f", fmt::format("fill {:x}-{:x}", page_address, page_line));
                        if (c.responsep) {
                            response r;
                            r.transaction = c.transaction;
                            r.line_address = c.line_address;
                            if (auto p = image.find(page_address); p != image.end())
                                r.data = p->second[page_line];
                            else
                                r.data.fill(0);
                            c.responsep->write(std::move(r));
                        }
                    }
                    current_command.reset();
                }
            }
        }

        void load(std::istream &fh) {
            for (unsigned int page = 0; ; ++page) {
                auto [p, emplaced] = image.try_emplace(page);
                if (emplaced)
                    for (unsigned int i = 0; i < PAGE_LINES; ++i)
                        p->second[i].fill(0);
                for (unsigned int line = 0; line < PAGE_LINES; ++line)
                    if (!fh.read(reinterpret_cast<char *>(p->second[line].data()), LINE_BYTES))
                        return;
            }
        }

        unsigned int fetch(std::uint64_t address) const {
            auto page_address = address >> PAGE_LINES_LOG2;
            if (auto i = image.find(page_address); i != image.end()) {
                auto &page = i->second;
                auto line_address = (address & PAGE_LINE_OFFSET_MASK) >> LINE_BYTES_LOG2;
                auto &line = page[line_address];
                auto word_address = address & LINE_BYTE_OFFSET_MASK;
                return line[word_address];
            }
            return 0;
        }

        void store(std::uint64_t address, unsigned int value) {
            auto page_address = address >> PAGE_LINES_LOG2;
            auto [p, emplaced] = image.try_emplace(page_address);
            auto &page = p->second;
            if (emplaced)
                for (unsigned int i = 0; i < PAGE_LINES; ++i)
                    page[i].fill(0);
            auto line_address = (address & PAGE_LINE_OFFSET_MASK) >> LINE_BYTES_LOG2;
            auto &line = page[line_address];
            auto word_address = address & LINE_BYTE_OFFSET_MASK;
            line[word_address] = value;
        }
    };
}