/*
 * sockbus.cpp — SockBus v2.0.0
 * Pure C++17 TCP socket IPC library.
 * Copyright (c) 2025-2026 Are Bjørby <are.bjorby@langsyn.org>
 * SPDX-License-Identifier: MIT
 */

#include "sockbus.hpp"

#include <algorithm>
#include <atomic>
#include <cerrno>
#include <cstring>
#include <fstream>
#include <mutex>
#include <string>
#include <system_error>
#include <thread>

#include <arpa/inet.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>

namespace sockbus {

namespace {

/* ── Wire frame ──────────────────────────────────────────────────────── *
 * [4 bytes: uint32_t payload length, network byte order][payload]       *
 * Preserves write atomicity through the broker.                         */
constexpr size_t FRAME_HDR_SIZE = 4;

/* ── Limits ──────────────────────────────────────────────────────────── */
constexpr int MAX_BROKERS = 16;
constexpr char CONF_ENV[] = "SOCKBUS_CONF";
constexpr char CONF_PATH[] = "/etc/sockbus/buses.conf";

/* ── Config resolution ───────────────────────────────────────────────── */

struct Address {
    std::string host;
    uint16_t    port;
};

Address parse_addr(std::string_view addr) {
    auto pos = addr.rfind(':');
    if (pos == std::string_view::npos)
        throw std::system_error(EINVAL, std::system_category());

    std::string host(addr.substr(0, pos));
    int port = 0;
    for (char c : addr.substr(pos + 1)) {
        if (c < '0' || c > '9')
            throw std::system_error(EINVAL, std::system_category());
        port = port * 10 + (c - '0');
    }
    if (port <= 0 || port > 65535)
        throw std::system_error(EINVAL, std::system_category());

    return {std::move(host), static_cast<uint16_t>(port)};
}

Address resolve_name(std::string_view name) {
    /* Literal host:port — skip config */
    if (name.find(':') != std::string_view::npos)
        return parse_addr(name);

    const char* conf_env = std::getenv(CONF_ENV);
    std::string conf_file = conf_env ? conf_env : CONF_PATH;

    std::ifstream f(conf_file);
    if (!f.is_open())
        throw std::system_error(errno, std::system_category());

    std::string line;
    while (std::getline(f, line)) {
        auto p = line.find_first_not_of(" \t");
        if (p == std::string::npos || line[p] == '#') continue;

        auto end = line.find_first_of(" \t", p);
        if (end == std::string::npos) continue;
        std::string_view ename(line.data() + p, end - p);

        auto astart = line.find_first_not_of(" \t", end);
        if (astart == std::string::npos) continue;
        auto aend = line.find_first_of(" \t\n\r", astart);
        std::string_view eaddr(line.data() + astart,
                               (aend == std::string::npos ? line.size() : aend) - astart);

        if (ename == name) return parse_addr(eaddr);
    }

    throw std::system_error(ENOENT, std::system_category());
}

/* ── Framed send helpers ─────────────────────────────────────────────── */

ssize_t send_all(int fd, const uint8_t* data, size_t len) {
    size_t rem = len;
    while (rem > 0) {
        ssize_t n = ::send(fd, data, rem, MSG_NOSIGNAL);
        if (n < 0) {
            if (errno == EINTR) continue;
            return -1;
        }
        data += n;
        rem -= static_cast<size_t>(n);
    }
    return static_cast<ssize_t>(len);
}

ssize_t send_frame(int fd, const uint8_t* data, size_t len) {
    size_t total = FRAME_HDR_SIZE + len;
    uint8_t sbuf[4096];
    std::unique_ptr<uint8_t[]> heap;
    uint8_t* buf = sbuf;
    if (total > sizeof(sbuf)) {
        heap = std::make_unique<uint8_t[]>(total);
        buf = heap.get();
    }

    uint32_t hdr = htonl(static_cast<uint32_t>(len));
    std::memcpy(buf, &hdr, FRAME_HDR_SIZE);
    if (len > 0) std::memcpy(buf + FRAME_HDR_SIZE, data, len);

    ssize_t ret = send_all(fd, buf, total);
    return (ret < 0) ? -1 : static_cast<ssize_t>(len);
}

size_t extract_frame(uint8_t* buf, size_t buf_len,
                     uint8_t** payload, size_t* consumed) {
    if (buf_len < FRAME_HDR_SIZE) return 0;
    uint32_t plen;
    std::memcpy(&plen, buf, sizeof(plen));
    plen = ntohl(plen);
    if (buf_len < FRAME_HDR_SIZE + plen) return 0;
    *payload = buf + FRAME_HDR_SIZE;
    *consumed = FRAME_HDR_SIZE + plen;
    return plen;
}

/* ── Broker ──────────────────────────────────────────────────────────── */

struct Broker {
    std::string        name;
    int                listen_fd = -1;
    std::thread        thread;
    std::atomic<bool>  running{false};
    int                pipe_fd[2] = {-1, -1};
    size_t             buf_size = 0;

    /* Client connections (broker thread mutates after lock) */
    int                              fds[MAX_READERS];
    std::unique_ptr<uint8_t[]>       recv_bufs[MAX_READERS];
    size_t                           recv_lens[MAX_READERS] = {};
    size_t                           recv_caps[MAX_READERS] = {};
    int                              count = 0;
    std::mutex                       mtx;

    Broker() { for (auto& fd : fds) fd = -1; }
};

void broker_remove_client(Broker& b, int idx) {
    ::close(b.fds[idx]);
    b.recv_bufs[idx].reset();
    b.count--;
    if (idx < b.count) {
        b.fds[idx]       = b.fds[b.count];
        b.recv_bufs[idx] = std::move(b.recv_bufs[b.count]);
        b.recv_lens[idx] = b.recv_lens[b.count];
        b.recv_caps[idx] = b.recv_caps[b.count];
    }
    b.fds[b.count] = -1;
    b.recv_bufs[b.count].reset();
    b.recv_lens[b.count] = 0;
}

void broker_loop(Broker* b) {
    auto relay = std::make_unique<uint8_t[]>(b->buf_size);

    while (b->running.load()) {
        pollfd pfds[2 + MAX_READERS];
        pfds[0] = {b->pipe_fd[0], POLLIN, 0};
        pfds[1] = {b->listen_fd, POLLIN, 0};

        int nc;
        {
            std::lock_guard lock(b->mtx);
            nc = b->count;
            for (int i = 0; i < nc; i++)
                pfds[2 + i] = {b->fds[i], POLLIN, 0};
        }

        int ret = ::poll(pfds, static_cast<nfds_t>(2 + nc), 100);
        if (ret < 0) { if (errno == EINTR) continue; break; }
        if (ret == 0) continue;

        /* Shutdown signal */
        if (pfds[0].revents & POLLIN) break;

        std::lock_guard lock(b->mtx);

        /* Accept new connections */
        if (pfds[1].revents & POLLIN) {
            int cfd = ::accept4(b->listen_fd, nullptr, nullptr,
                                SOCK_NONBLOCK | SOCK_CLOEXEC);
            if (cfd >= 0) {
                if (b->count < static_cast<int>(MAX_READERS)) {
                    int one = 1;
                    ::setsockopt(cfd, IPPROTO_TCP, TCP_NODELAY,
                                 &one, sizeof(one));
                    int idx = b->count++;
                    b->fds[idx]       = cfd;
                    b->recv_bufs[idx] = std::make_unique<uint8_t[]>(b->buf_size);
                    b->recv_lens[idx] = 0;
                    b->recv_caps[idx] = b->buf_size;
                } else {
                    ::close(cfd);
                }
            }
        }

        /* Phase 1: Receive data, mark dead clients */
        bool dead[MAX_READERS] = {};

        for (int i = 0; i < nc && i < b->count; i++) {
            if (!(pfds[2 + i].revents & (POLLIN | POLLHUP | POLLERR)))
                continue;
            if (b->recv_caps[i] <= b->recv_lens[i]) continue;

            ssize_t n = ::recv(b->fds[i],
                               b->recv_bufs[i].get() + b->recv_lens[i],
                               b->recv_caps[i] - b->recv_lens[i], 0);
            if (n <= 0) { dead[i] = true; continue; }
            b->recv_lens[i] += static_cast<size_t>(n);
        }

        /* Phase 2: Extract frames, broadcast to all clients */
        for (int i = 0; i < b->count; i++) {
            if (dead[i]) continue;
            for (;;) {
                uint8_t* payload;
                size_t consumed;
                size_t plen = extract_frame(b->recv_bufs[i].get(),
                                            b->recv_lens[i],
                                            &payload, &consumed);
                if (plen == 0) break;

                if (consumed <= b->buf_size) {
                    std::memcpy(relay.get(), b->recv_bufs[i].get(), consumed);

                    /* Broadcast to ALL clients (including sender) */
                    for (int j = 0; j < b->count; j++) {
                        if (dead[j]) continue;
                        ssize_t s = ::send(b->fds[j], relay.get(), consumed,
                                           MSG_DONTWAIT | MSG_NOSIGNAL);
                        if (s < 0) {
                            if (errno != EAGAIN && errno != EWOULDBLOCK)
                                dead[j] = true;
                        } else if (static_cast<size_t>(s) < consumed) {
                            dead[j] = true;
                        }
                    }
                }

                b->recv_lens[i] -= consumed;
                if (b->recv_lens[i] > 0)
                    std::memmove(b->recv_bufs[i].get(),
                                 b->recv_bufs[i].get() + consumed,
                                 b->recv_lens[i]);
            }
        }

        /* Phase 3: Remove dead clients (reverse order for stable indices) */
        for (int i = b->count - 1; i >= 0; i--) {
            if (dead[i]) broker_remove_client(*b, i);
        }
    }

    /* Final cleanup: close all client sockets */
    std::lock_guard lock(b->mtx);
    for (int i = 0; i < b->count; i++) {
        ::close(b->fds[i]);
        b->recv_bufs[i].reset();
    }
    b->count = 0;
}

/* ── Broker registry (process-global) ────────────────────────────────── */

std::unique_ptr<Broker> g_brokers[MAX_BROKERS];
int                     g_broker_count = 0;
std::mutex              g_lock;

} // anonymous namespace

/* ═══════════════════════════════════════════════════════════════════════ */
/* PUBLIC API                                                             */
/* ═══════════════════════════════════════════════════════════════════════ */

void create(std::string_view name, size_t size) {
    if (name.empty())
        throw std::system_error(EINVAL, std::system_category());
    if (size == 0) size = DEFAULT_SIZE;

    auto addr = resolve_name(name);

    std::lock_guard lock(g_lock);

    for (int i = 0; i < g_broker_count; i++) {
        if (g_brokers[i]->name == name)
            throw std::system_error(EEXIST, std::system_category());
    }
    if (g_broker_count >= MAX_BROKERS)
        throw std::system_error(ENOMEM, std::system_category());

    int lfd = ::socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
    if (lfd < 0)
        throw std::system_error(errno, std::system_category());

    int opt = 1;
    ::setsockopt(lfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    sockaddr_in sa{};
    sa.sin_family = AF_INET;
    sa.sin_port   = htons(addr.port);
    if (::inet_pton(AF_INET, addr.host.c_str(), &sa.sin_addr) != 1) {
        ::close(lfd);
        throw std::system_error(EINVAL, std::system_category());
    }

    if (::bind(lfd, reinterpret_cast<sockaddr*>(&sa), sizeof(sa)) < 0) {
        int e = errno;
        ::close(lfd);
        throw std::system_error(e, std::system_category());
    }
    if (::listen(lfd, static_cast<int>(MAX_READERS)) < 0) {
        int e = errno;
        ::close(lfd);
        throw std::system_error(e, std::system_category());
    }

    auto b = std::make_unique<Broker>();
    b->name      = std::string(name);
    b->listen_fd = lfd;
    b->buf_size  = size;
    b->running.store(true);

    if (::pipe2(b->pipe_fd, O_CLOEXEC) < 0) {
        int e = errno;
        ::close(lfd);
        throw std::system_error(e, std::system_category());
    }

    Broker* bp = b.get();
    b->thread = std::thread(broker_loop, bp);

    g_brokers[g_broker_count++] = std::move(b);
}

void destroy(std::string_view name) {
    if (name.empty())
        throw std::system_error(EINVAL, std::system_category());

    std::unique_lock lock(g_lock);

    for (int i = 0; i < g_broker_count; i++) {
        if (g_brokers[i]->name != name) continue;

        auto& b = g_brokers[i];
        b->running.store(false);
        (void)::write(b->pipe_fd[1], "x", 1);

        /* Unlock during join to avoid holding global lock */
        lock.unlock();
        if (b->thread.joinable()) b->thread.join();
        lock.lock();

        ::close(b->listen_fd);
        ::close(b->pipe_fd[0]);
        ::close(b->pipe_fd[1]);

        g_broker_count--;
        if (i < g_broker_count)
            g_brokers[i] = std::move(g_brokers[g_broker_count]);
        g_brokers[g_broker_count].reset();

        return;
    }

    throw std::system_error(ENOENT, std::system_category());
}

/* ── Bus::Impl ───────────────────────────────────────────────────────── */

struct Bus::Impl {
    std::string                bus_name;
    int                        fd = -1;
    char                       reader_name[READER_NAME_LEN] = {};
    std::mutex                 mtx;
    std::unique_ptr<uint8_t[]> recv_buf;
    size_t                     recv_len = 0;
    size_t                     recv_cap = 0;

    ~Impl() {
        if (fd >= 0) ::close(fd);
    }
};

Bus::Bus(std::string_view name) : impl_(std::make_unique<Impl>()) {
    if (name.empty())
        throw std::system_error(EINVAL, std::system_category());

    auto addr = resolve_name(name);

    int fd = ::socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
    if (fd < 0)
        throw std::system_error(errno, std::system_category());

    int one = 1;
    ::setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));

    sockaddr_in sa{};
    sa.sin_family = AF_INET;
    sa.sin_port   = htons(addr.port);
    if (::inet_pton(AF_INET, addr.host.c_str(), &sa.sin_addr) != 1) {
        ::close(fd);
        throw std::system_error(EINVAL, std::system_category());
    }
    if (::connect(fd, reinterpret_cast<sockaddr*>(&sa), sizeof(sa)) < 0) {
        int e = errno;
        ::close(fd);
        throw std::system_error(e, std::system_category());
    }

    impl_->bus_name = std::string(name);
    impl_->fd       = fd;
    impl_->recv_cap = DEFAULT_SIZE;
    impl_->recv_buf = std::make_unique<uint8_t[]>(impl_->recv_cap);
}

Bus::~Bus() = default;

Bus::Bus(Bus&& o) noexcept = default;
Bus& Bus::operator=(Bus&& o) noexcept = default;

size_t Bus::write(const uint8_t* data, size_t len) {
    if (!data)
        throw std::system_error(EINVAL, std::system_category());
    if (len > MAX_MSG_SIZE)
        throw std::system_error(EMSGSIZE, std::system_category());

    std::lock_guard lock(impl_->mtx);
    ssize_t ret = send_frame(impl_->fd, data, len);
    if (ret < 0)
        throw std::system_error(errno, std::system_category());
    return static_cast<size_t>(ret);
}

size_t Bus::write(std::string_view sv) {
    return write(reinterpret_cast<const uint8_t*>(sv.data()), sv.size());
}

size_t Bus::read(uint8_t* buf, size_t len) {
    if (!buf)
        throw std::system_error(EINVAL, std::system_category());

    std::lock_guard lock(impl_->mtx);

    /* Try to recv more data (non-blocking) */
    if (impl_->recv_cap > impl_->recv_len) {
        ssize_t n = ::recv(impl_->fd,
                           impl_->recv_buf.get() + impl_->recv_len,
                           impl_->recv_cap - impl_->recv_len, MSG_DONTWAIT);
        if (n > 0)
            impl_->recv_len += static_cast<size_t>(n);
        else if (n == 0)
            throw std::system_error(ECONNRESET, std::system_category());
        /* EAGAIN is fine — no data yet */
    }

    /* Extract one frame */
    uint8_t* payload;
    size_t consumed;
    size_t plen = extract_frame(impl_->recv_buf.get(), impl_->recv_len,
                                &payload, &consumed);
    if (plen == 0) return 0;

    size_t copy = std::min(plen, len);
    std::memcpy(buf, payload, copy);
    impl_->recv_len -= consumed;
    if (impl_->recv_len > 0)
        std::memmove(impl_->recv_buf.get(),
                     impl_->recv_buf.get() + consumed,
                     impl_->recv_len);

    return copy;
}

size_t Bus::read_wait(uint8_t* buf, size_t len, int timeout_ms) {
    if (!buf)
        throw std::system_error(EINVAL, std::system_category());
    if (timeout_ms == 0) return read(buf, len);

    timespec deadline{};
    if (timeout_ms > 0) {
        clock_gettime(CLOCK_MONOTONIC, &deadline);
        deadline.tv_sec  += timeout_ms / 1000;
        deadline.tv_nsec += static_cast<long>(timeout_ms % 1000) * 1000000L;
        if (deadline.tv_nsec >= 1000000000L) {
            deadline.tv_sec++;
            deadline.tv_nsec -= 1000000000L;
        }
    }

    for (;;) {
        {
            std::lock_guard lock(impl_->mtx);

            /* Check for complete frame in buffer */
            uint8_t* payload;
            size_t consumed;
            size_t plen = extract_frame(impl_->recv_buf.get(), impl_->recv_len,
                                        &payload, &consumed);
            if (plen > 0) {
                size_t copy = std::min(plen, len);
                std::memcpy(buf, payload, copy);
                impl_->recv_len -= consumed;
                if (impl_->recv_len > 0)
                    std::memmove(impl_->recv_buf.get(),
                                 impl_->recv_buf.get() + consumed,
                                 impl_->recv_len);
                return copy;
            }
        }

        /* Calculate remaining timeout */
        int poll_ms = -1;
        if (timeout_ms > 0) {
            timespec now{};
            clock_gettime(CLOCK_MONOTONIC, &now);
            long rem = (deadline.tv_sec - now.tv_sec) * 1000
                     + (deadline.tv_nsec - now.tv_nsec) / 1000000;
            if (rem <= 0) return 0;
            poll_ms = static_cast<int>(rem);
        }

        pollfd pfd = {impl_->fd, POLLIN, 0};
        int ret = ::poll(&pfd, 1, poll_ms);
        if (ret < 0) {
            if (errno == EINTR) continue;
            throw std::system_error(errno, std::system_category());
        }
        if (ret == 0) return 0;

        /* Recv more data */
        std::lock_guard lock(impl_->mtx);
        if (impl_->recv_cap > impl_->recv_len) {
            ssize_t n = ::recv(impl_->fd,
                               impl_->recv_buf.get() + impl_->recv_len,
                               impl_->recv_cap - impl_->recv_len, 0);
            if (n > 0)
                impl_->recv_len += static_cast<size_t>(n);
            else if (n == 0)
                throw std::system_error(ECONNRESET, std::system_category());
        }
    }
}

void Bus::set_reader_name(std::string_view label) {
    size_t copy = std::min(label.size(), READER_NAME_LEN - 1);
    std::memcpy(impl_->reader_name, label.data(), copy);
    impl_->reader_name[copy] = '\0';
}

std::string_view Bus::name() const {
    return impl_->bus_name;
}

} // namespace sockbus