#include "shm_ring.hpp"

#include <atomic>
#include <chrono>
#include <cerrno>
#include <cstring>
#include <fcntl.h>
#include <new>
#include <stdexcept>
#include <string>
#include <string_view>
#include <sys/mman.h>
#include <sys/stat.h>
#include <thread>
#include <unistd.h>
#include <utility>

namespace radar::ipc {

struct alignas(64) ShmRing::Header {
    char magic[8]{};
    std::uint32_t version = 0;
    std::uint32_t capacity = 0;
    std::uint32_t slot_size_bytes = 0;
    std::uint32_t reserved = 0;
    std::atomic<std::uint64_t> write_seq{};
    std::atomic<std::uint64_t> read_seq{};
    std::atomic<std::uint64_t> dropped{};
};

struct alignas(16) ShmRing::SlotHeader {
    std::uint32_t payload_size = 0;
    std::uint32_t reserved = 0;
    std::uint64_t sequence = 0;
};

namespace {

constexpr std::string_view kRingMagic = "RDRRING2";
constexpr std::uint32_t kRingVersion = 1;
constexpr auto kHeaderInitWaitTimeout = std::chrono::milliseconds(1000);
constexpr auto kHeaderInitPollInterval = std::chrono::milliseconds(2);

[[nodiscard]] auto checked_u64_diff(std::uint64_t left, std::uint64_t right) -> std::uint64_t {
    return left >= right ? left - right : 0;
}

[[nodiscard]] auto errno_message(const std::string& action, const std::string& name) -> std::runtime_error {
    return std::runtime_error(action + " " + name + ": " + std::strerror(errno));
}

class ScopedFd {
  public:
    explicit ScopedFd(int fd) : fd_(fd) {}

    ~ScopedFd() {
        if (fd_ >= 0) {
            ::close(fd_);
        }
    }

    ScopedFd(const ScopedFd&) = delete;
    auto operator=(const ScopedFd&) -> ScopedFd& = delete;

    ScopedFd(ScopedFd&& other) noexcept : fd_(std::exchange(other.fd_, -1)) {}

    auto release() -> int {
        return std::exchange(fd_, -1);
    }

  private:
    int fd_ = -1;
};

class ScopedMmap {
  public:
    ScopedMmap(void* address, std::size_t size) : address_(address), size_(size) {}

    ~ScopedMmap() {
        if (address_ != nullptr && address_ != MAP_FAILED) {
            munmap(address_, size_);
        }
    }

    ScopedMmap(const ScopedMmap&) = delete;
    auto operator=(const ScopedMmap&) -> ScopedMmap& = delete;

    ScopedMmap(ScopedMmap&& other) noexcept
        : address_(std::exchange(other.address_, nullptr)),
          size_(std::exchange(other.size_, 0U)) {}

    auto release() -> std::pair<void*, std::size_t> {
        return {std::exchange(address_, nullptr), std::exchange(size_, 0U)};
    }

  private:
    void* address_ = nullptr;
    std::size_t size_ = 0;
};

}  // namespace

ShmRing::~ShmRing() {
    close();
}

ShmRing::ShmRing(ShmRing&& other) noexcept {
    *this = std::move(other);
}

auto ShmRing::operator=(ShmRing&& other) noexcept -> ShmRing& {
    if (this != &other) {
        close();

        fd_ = std::exchange(other.fd_, -1);
        mapped_size_ = std::exchange(other.mapped_size_, 0U);
        mapped_ = std::exchange(other.mapped_, nullptr);
        header_ = std::exchange(other.header_, nullptr);
    }
    return *this;
}

auto ShmRing::open_or_create(
    const std::string& name,
    std::uint32_t capacity,
    std::uint32_t slot_size_bytes
) -> ShmRing {
    validate_name(name);
    if (capacity == 0U) {
        throw std::runtime_error("Ring capacity must be > 0");
    }
    if (slot_size_bytes == 0U) {
        throw std::runtime_error("Ring slot size must be > 0");
    }

    bool created = false;
    int fd = shm_open(name.c_str(), O_RDWR | O_CREAT | O_EXCL, 0660);
    if (fd >= 0) {
        created = true;
    } else if (errno == EEXIST) {
        fd = shm_open(name.c_str(), O_RDWR, 0660);
    }

    if (fd < 0) {
        throw errno_message("Failed to open shared memory ring", name);
    }

    ScopedFd scoped_fd(fd);

    const auto slot_stride = sizeof(SlotHeader) + static_cast<std::size_t>(slot_size_bytes);
    const auto mapped_size = sizeof(Header) + slot_stride * capacity;

    if (created) {
        if (ftruncate(fd, static_cast<off_t>(mapped_size)) != 0) {
            throw errno_message("Failed to resize shared memory ring", name);
        }
    }

    void* mapped = mmap(nullptr, mapped_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if (mapped == MAP_FAILED) {
        throw errno_message("Failed to mmap shared memory ring", name);
    }

    ScopedMmap scoped_mmap(mapped, mapped_size);
    auto* header = static_cast<Header*>(mapped);

    if (created) {
        std::memset(mapped, 0, mapped_size);
        std::memcpy(header->magic, kRingMagic.data(), kRingMagic.size());
        header->version = kRingVersion;
        header->capacity = capacity;
        header->slot_size_bytes = slot_size_bytes;

        new (&header->write_seq) std::atomic<std::uint64_t>(0);
        new (&header->read_seq) std::atomic<std::uint64_t>(0);
        new (&header->dropped) std::atomic<std::uint64_t>(0);
        std::atomic_thread_fence(std::memory_order_release);
    } else {
        const auto deadline = std::chrono::steady_clock::now() + kHeaderInitWaitTimeout;

        while (true) {
            std::atomic_thread_fence(std::memory_order_acquire);

            const bool magic_ok = std::memcmp(header->magic, kRingMagic.data(), kRingMagic.size()) == 0;
            const bool version_ok = header->version == kRingVersion;
            const bool geometry_ok = header->capacity == capacity && header->slot_size_bytes == slot_size_bytes;

            if (magic_ok && version_ok && geometry_ok) {
                break;
            }

            if (std::chrono::steady_clock::now() >= deadline) {
                if (!magic_ok) {
                    throw std::runtime_error("Shared memory ring magic mismatch for " + name);
                }
                if (!version_ok) {
                    throw std::runtime_error("Shared memory ring version mismatch for " + name);
                }
                throw std::runtime_error("Shared memory ring geometry mismatch for " + name);
            }

            std::this_thread::sleep_for(kHeaderInitPollInterval);
        }
    }

    ShmRing ring{};
    ring.fd_ = scoped_fd.release();

    const auto [released_mapped, released_size] = scoped_mmap.release();
    ring.mapped_ = released_mapped;
    ring.mapped_size_ = released_size;
    ring.header_ = static_cast<Header*>(released_mapped);
    return ring;
}

auto ShmRing::open_existing(const std::string& name) -> ShmRing {
    validate_name(name);

    const int fd = shm_open(name.c_str(), O_RDWR, 0660);
    if (fd < 0) {
        throw errno_message("Failed to open shared memory ring", name);
    }

    ScopedFd scoped_fd(fd);

    struct stat info {};
    if (fstat(fd, &info) != 0) {
        throw errno_message("Failed to stat shared memory ring", name);
    }

    if (info.st_size < static_cast<off_t>(sizeof(Header))) {
        throw std::runtime_error("Shared memory ring size is too small for " + name);
    }

    const auto mapped_size = static_cast<std::size_t>(info.st_size);
    void* mapped = mmap(nullptr, mapped_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if (mapped == MAP_FAILED) {
        throw errno_message("Failed to mmap shared memory ring", name);
    }

    ScopedMmap scoped_mmap(mapped, mapped_size);
    auto* header = static_cast<Header*>(mapped);

    if (std::memcmp(header->magic, kRingMagic.data(), kRingMagic.size()) != 0) {
        throw std::runtime_error("Shared memory ring magic mismatch for " + name);
    }
    if (header->version != kRingVersion) {
        throw std::runtime_error("Shared memory ring version mismatch for " + name);
    }

    ShmRing ring{};
    ring.fd_ = scoped_fd.release();

    const auto [released_mapped, released_size] = scoped_mmap.release();
    ring.mapped_ = released_mapped;
    ring.mapped_size_ = released_size;
    ring.header_ = static_cast<Header*>(released_mapped);
    return ring;
}

void ShmRing::unlink_ring(const std::string& name) {
    validate_name(name);
    if (shm_unlink(name.c_str()) != 0 && errno != ENOENT) {
        throw errno_message("Failed to unlink shared memory ring", name);
    }
}

auto ShmRing::is_open() const -> bool {
    return header_ != nullptr;
}

auto ShmRing::capacity() const -> std::uint32_t {
    return header_ != nullptr ? header_->capacity : 0;
}

auto ShmRing::slot_size_bytes() const -> std::uint32_t {
    return header_ != nullptr ? header_->slot_size_bytes : 0;
}

auto ShmRing::dropped_count() const -> std::uint64_t {
    if (header_ == nullptr) {
        return 0;
    }
    return header_->dropped.load(std::memory_order_acquire);
}

auto ShmRing::push(std::span<const std::uint8_t> payload) -> bool {
    if (header_ == nullptr) {
        throw std::runtime_error("Ring is not open");
    }
    if (payload.size() > header_->slot_size_bytes) {
        return false;
    }

    const auto write_seq = header_->write_seq.load(std::memory_order_relaxed);
    const auto read_seq = header_->read_seq.load(std::memory_order_acquire);

    if (checked_u64_diff(write_seq, read_seq) >= header_->capacity) {
        header_->read_seq.store(read_seq + 1U, std::memory_order_release);
        header_->dropped.fetch_add(1U, std::memory_order_relaxed);
    }

    auto* slot = slot_header(write_seq);
    slot->payload_size = static_cast<std::uint32_t>(payload.size());
    slot->sequence = write_seq + 1U;
    std::memcpy(slot_payload(slot), payload.data(), payload.size());

    std::atomic_thread_fence(std::memory_order_release);
    header_->write_seq.store(write_seq + 1U, std::memory_order_release);
    return true;
}

auto ShmRing::pop(std::vector<std::uint8_t>& payload) -> bool {
    if (header_ == nullptr) {
        throw std::runtime_error("Ring is not open");
    }

    const auto read_seq = header_->read_seq.load(std::memory_order_relaxed);
    const auto write_seq = header_->write_seq.load(std::memory_order_acquire);
    if (read_seq >= write_seq) {
        return false;
    }

    auto* slot = slot_header(read_seq);
    if (slot->sequence != read_seq + 1U) {
        // Producer overwrote this slot before consumer read it. Resync to latest.
        header_->read_seq.store(write_seq, std::memory_order_release);
        return false;
    }

    const auto payload_size = slot->payload_size;
    if (payload_size > header_->slot_size_bytes) {
        header_->read_seq.store(write_seq, std::memory_order_release);
        throw std::runtime_error("Invalid payload size in shared memory slot");
    }

    payload.resize(payload_size);
    std::memcpy(payload.data(), slot_payload(slot), payload_size);

    std::atomic_thread_fence(std::memory_order_acquire);
    header_->read_seq.store(read_seq + 1U, std::memory_order_release);
    return true;
}

auto ShmRing::slot_stride_bytes() const -> std::size_t {
    return sizeof(SlotHeader) + header_->slot_size_bytes;
}

auto ShmRing::slot_header(std::uint64_t sequence) const -> SlotHeader* {
    const auto index = sequence % header_->capacity;
    auto* slots_begin = static_cast<std::uint8_t*>(mapped_) + sizeof(Header);
    return reinterpret_cast<SlotHeader*>(slots_begin + index * slot_stride_bytes());
}

auto ShmRing::slot_payload(SlotHeader* slot) const -> std::uint8_t* {
    return reinterpret_cast<std::uint8_t*>(slot) + sizeof(SlotHeader);
}

void ShmRing::close() {
    if (mapped_ != nullptr) {
        munmap(mapped_, mapped_size_);
        mapped_ = nullptr;
    }

    if (fd_ >= 0) {
        ::close(fd_);
        fd_ = -1;
    }

    mapped_size_ = 0;
    header_ = nullptr;
}

void ShmRing::validate_name(const std::string& name) {
    if (name.empty() || name.front() != '/') {
        throw std::runtime_error("POSIX shm name must start with '/'");
    }
}

}  // namespace radar::ipc
