#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#include <cstdarg>
#include <cstdint>
#include <cstdlib>
#include <string>
#include <unordered_map>
#include <vector>

namespace {

constexpr const char* kRelayVersion = "1.4";

constexpr uint32_t kMagic    = 0x4B534450u;
constexpr uint32_t kRelayVer = 1;

constexpr uint32_t kCodeLen = 16;
constexpr uint32_t kNameLen = 32;
constexpr uint32_t kIdLen   = 16;

constexpr uint16_t MSG_HELLO        = 0x0001;
constexpr uint16_t MSG_RELAY_HOST   = 0x0030;
constexpr uint16_t MSG_RELAY_READY  = 0x0031;
constexpr uint16_t MSG_RELAY_OFFER  = 0x0032;
constexpr uint16_t MSG_RELAY_ACCEPT = 0x0033;
constexpr uint16_t MSG_RELAY_ERROR  = 0x0034;
constexpr uint16_t MSG_RELAY_STATUS = 0x0035;

#pragma pack(push, 1)
struct MsgHeader {
    uint32_t magic;
    uint16_t type;
    uint16_t flags;
    uint32_t length;
};
static_assert(sizeof(MsgHeader) == 12, "MsgHeader must stay 12 bytes on the wire");

struct HelloPayload {
    uint32_t version;
    uint8_t  channel;
    uint8_t  reserved[3];
    uint64_t session;
    char     code[kCodeLen];
    char     name[kNameLen];
    char     id[kIdLen];
};

struct RelayHostPayload {
    uint32_t version;
    char     id[kIdLen];
    char     name[kNameLen];
};

struct RelayTicketPayload {
    uint64_t ticket;
};

struct RelayStatusPayload {
    uint32_t version;
    uint32_t hosts;
    uint32_t sessions;
    uint32_t maxHosts;
};
#pragma pack(pop)

constexpr uint32_t kMaxIntroFrame  = 4096;
constexpr int      kIntroTimeoutSec = 20;
constexpr size_t   kMaxParkedBytes = 256u * 1024u;
constexpr size_t   kMaxOutBuffer   = 512u * 1024u;

enum class Role { Intro, Host, Waiting, Piped };

struct Conn {
    int         fd    = -1;
    Role        role  = Role::Intro;
    time_t      since = 0;
    std::string peer;

    std::vector<uint8_t> intro;
    size_t               introNeed = sizeof(MsgHeader);
    bool                 introHeaderDone = false;

    std::string id;
    uint64_t    ticket = 0;

    int  peerFd = -1;
    std::vector<uint8_t> out;
    bool wantWrite = false;
    bool dying     = false;
};

int                                  g_epoll = -1;
std::unordered_map<int, Conn>        g_conns;
std::unordered_map<std::string, int> g_hosts;
std::unordered_map<uint64_t, int>    g_tickets;
uint64_t                             g_nextTicket  = 1;
int                                  g_maxSessions = 64;
volatile sig_atomic_t                g_stop = 0;

void OnSignal(int) { g_stop = 1; }

void Log(const char* fmt, ...) __attribute__((format(printf, 1, 2)));
void Log(const char* fmt, ...) {
    char when[32];
    const time_t now = time(nullptr);
    struct tm tm {};
    localtime_r(&now, &tm);
    strftime(when, sizeof(when), "%Y-%m-%d %H:%M:%S", &tm);

    va_list ap;
    va_start(ap, fmt);
    printf("[%s] ", when);
    vprintf(fmt, ap);
    printf("\n");
    va_end(ap);
    fflush(stdout);
}

bool SetNonBlocking(int fd) {
    const int f = fcntl(fd, F_GETFL, 0);
    return f != -1 && fcntl(fd, F_SETFL, f | O_NONBLOCK) == 0;
}

void Arm(Conn& c) {
    epoll_event ev{};
    ev.events  = static_cast<uint32_t>(EPOLLIN) | static_cast<uint32_t>(EPOLLRDHUP) |
                 (c.wantWrite ? static_cast<uint32_t>(EPOLLOUT) : 0u);
    ev.data.fd = c.fd;
    epoll_ctl(g_epoll, EPOLL_CTL_MOD, c.fd, &ev);
}

std::string PeerName(int fd) {
    sockaddr_storage ss{};
    socklen_t len = sizeof(ss);
    if (getpeername(fd, reinterpret_cast<sockaddr*>(&ss), &len) != 0) return "?";

    char host[INET6_ADDRSTRLEN] = {};
    if (ss.ss_family == AF_INET6) {
        auto* a = reinterpret_cast<sockaddr_in6*>(&ss);
        inet_ntop(AF_INET6, &a->sin6_addr, host, sizeof(host));
        std::string s = host;
        if (s.rfind("::ffff:", 0) == 0) s.erase(0, 7);
        return s;
    }
    auto* a = reinterpret_cast<sockaddr_in*>(&ss);
    inet_ntop(AF_INET, &a->sin_addr, host, sizeof(host));
    return host;
}

bool Queue(Conn& c, const void* data, size_t len) {
    if (c.dying) return false;
    const auto* p = static_cast<const uint8_t*>(data);

    if (c.out.empty()) {
        while (len > 0) {
            const ssize_t n = send(c.fd, p, len, MSG_NOSIGNAL);
            if (n > 0) { p += n; len -= static_cast<size_t>(n); continue; }
            if (n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) break;
            if (n < 0 && errno == EINTR) continue;
            return false;
        }
        if (len == 0) return true;
    }

    if (c.out.size() + len > kMaxOutBuffer) return false;
    c.out.insert(c.out.end(), p, p + len);
    if (!c.wantWrite) { c.wantWrite = true; Arm(c); }
    return true;
}

bool SendFrame(Conn& c, uint16_t type, const void* payload, uint32_t len) {
    MsgHeader h{ kMagic, type, 0, len };
    if (!Queue(c, &h, sizeof(h))) return false;
    return len == 0 || Queue(c, payload, len);
}

void SendError(Conn& c, const char* text) {
    SendFrame(c, MSG_RELAY_ERROR, text, static_cast<uint32_t>(strlen(text)));
}

void Close(int fd);

void Unlink(Conn& c) {
    if (c.role == Role::Host) {
        auto it = g_hosts.find(c.id);
        if (it != g_hosts.end() && it->second == c.fd) {
            g_hosts.erase(it);
            Log("host %s gone (%s)", c.id.c_str(), c.peer.c_str());
        }
    }
    if (c.role == Role::Waiting) g_tickets.erase(c.ticket);
}

void Close(int fd) {
    auto it = g_conns.find(fd);
    if (it == g_conns.end()) return;
    Conn& c = it->second;
    if (c.dying) return;
    c.dying = true;

    Unlink(c);

    const int peer = c.peerFd;
    if (peer >= 0) {
        auto pit = g_conns.find(peer);
        if (pit != g_conns.end()) pit->second.peerFd = -1;
        Close(peer);
    }

    epoll_ctl(g_epoll, EPOLL_CTL_DEL, fd, nullptr);
    ::close(fd);
    g_conns.erase(fd);
}

std::string CodeFrom(const char* raw, size_t max) {
    size_t n = 0;
    while (n < max && raw[n] != '\0') ++n;
    return std::string(raw, n);
}

void OnRelayHost(Conn& c, const std::vector<uint8_t>& body) {
    if (body.size() != sizeof(RelayHostPayload)) { SendError(c, "bad host payload"); Close(c.fd); return; }

    RelayHostPayload p{};
    memcpy(&p, body.data(), sizeof(p));
    if (p.version != kRelayVer) {
        SendError(c, "this relay speaks a different registration version");
        Close(c.fd);
        return;
    }

    const std::string id = CodeFrom(p.id, kIdLen);
    if (id.empty()) { SendError(c, "an ID is required"); Close(c.fd); return; }

    if (g_hosts.size() >= static_cast<size_t>(g_maxSessions) && !g_hosts.count(id)) {
        SendError(c, "relay is full");
        Close(c.fd);
        return;
    }

    auto old = g_hosts.find(id);
    if (old != g_hosts.end() && old->second != c.fd) Close(old->second);

    c.role = Role::Host;
    c.id   = id;
    c.intro.clear();
    g_hosts[id] = c.fd;

    SendFrame(c, MSG_RELAY_READY, nullptr, 0);
    Log("host %s registered from %s (%s)", id.c_str(), c.peer.c_str(),
        CodeFrom(p.name, kNameLen).c_str());
}

void OnRelayStatus(Conn& c) {
    uint32_t piped = 0;
    for (const auto& kv : g_conns)
        if (kv.second.role == Role::Piped) ++piped;

    RelayStatusPayload s{};
    s.version  = kRelayVer;
    s.hosts    = static_cast<uint32_t>(g_hosts.size());
    s.sessions = piped / 2;
    s.maxHosts = static_cast<uint32_t>(g_maxSessions);
    if (!SendFrame(c, MSG_RELAY_STATUS, &s, sizeof(s))) { Close(c.fd); return; }

    c.intro.clear();
    c.introHeaderDone = false;
    c.introNeed = sizeof(MsgHeader);
    c.since = time(nullptr) - kIntroTimeoutSec + 2;
}

void OnViewerHello(Conn& c, const MsgHeader& hdr, const std::vector<uint8_t>& body) {
    if (body.size() != sizeof(HelloPayload)) { SendError(c, "bad hello"); Close(c.fd); return; }

    HelloPayload h{};
    memcpy(&h, body.data(), sizeof(h));

    const std::string id = CodeFrom(h.id, kIdLen);
    auto host = g_hosts.find(id);
    if (id.empty() || host == g_hosts.end()) {
        SendError(c, "nobody is sharing with that ID");
        Log("viewer %s asked for ID %s: nobody there", c.peer.c_str(),
            id.empty() ? "(none)" : id.c_str());
        Close(c.fd);
        return;
    }

    auto hit = g_conns.find(host->second);
    if (hit == g_conns.end()) { SendError(c, "host went away"); Close(c.fd); return; }

    c.intro.clear();
    c.intro.resize(sizeof(MsgHeader) + body.size());
    memcpy(c.intro.data(), &hdr, sizeof(MsgHeader));
    if (!body.empty()) memcpy(c.intro.data() + sizeof(MsgHeader), body.data(), body.size());

    c.role   = Role::Waiting;
    c.ticket = g_nextTicket++;
    c.since  = time(nullptr);
    g_tickets[c.ticket] = c.fd;

    RelayTicketPayload t{ c.ticket };
    if (!SendFrame(hit->second, MSG_RELAY_OFFER, &t, sizeof(t))) {
        Close(hit->second.fd);
        Close(c.fd);
        return;
    }
    Log("viewer %s waiting on ID %s, ticket %llu, channel %u", c.peer.c_str(),
        id.c_str(), static_cast<unsigned long long>(c.ticket), h.channel);
}

void OnRelayAccept(Conn& c, const std::vector<uint8_t>& body) {
    if (body.size() != sizeof(RelayTicketPayload)) { SendError(c, "bad ticket"); Close(c.fd); return; }

    RelayTicketPayload t{};
    memcpy(&t, body.data(), sizeof(t));

    auto tit = g_tickets.find(t.ticket);
    if (tit == g_tickets.end()) { SendError(c, "unknown or expired ticket"); Close(c.fd); return; }

    const int viewerFd = tit->second;
    g_tickets.erase(tit);

    auto vit = g_conns.find(viewerFd);
    if (vit == g_conns.end()) { SendError(c, "viewer went away"); Close(c.fd); return; }
    Conn& v = vit->second;

    if (!c.intro.empty()) c.intro.clear();
    const size_t held = v.intro.size();
    if (!v.intro.empty() && !Queue(c, v.intro.data(), v.intro.size())) {
        Close(c.fd);
        Close(viewerFd);
        return;
    }
    v.intro.clear();

    c.role   = Role::Piped;
    v.role   = Role::Piped;
    c.peerFd = viewerFd;
    v.peerFd = c.fd;

    Log("paired ticket %llu: %s <-> %s (%zu bytes held)",
        static_cast<unsigned long long>(t.ticket), v.peer.c_str(), c.peer.c_str(), held);
}

void OnReadable(int fd) {
    auto it = g_conns.find(fd);
    if (it == g_conns.end()) return;
    Conn& c = it->second;

    if (c.role == Role::Piped) {
        auto pit = g_conns.find(c.peerFd);
        if (pit == g_conns.end()) { Close(fd); return; }
        Conn& peer = pit->second;

        if (peer.out.size() >= kMaxOutBuffer / 2) return;

        uint8_t buf[64 * 1024];
        for (;;) {
            const ssize_t n = recv(fd, buf, sizeof(buf), 0);
            if (n > 0) {
                if (!Queue(peer, buf, static_cast<size_t>(n))) { Close(fd); return; }
                if (peer.out.size() >= kMaxOutBuffer / 2) return;
                continue;
            }
            if (n == 0) { Close(fd); return; }
            if (errno == EINTR) continue;
            if (errno == EAGAIN || errno == EWOULDBLOCK) return;
            Close(fd);
            return;
        }
    }

    if (c.role == Role::Waiting) {
        for (;;) {
            uint8_t buf[8192];
            const ssize_t n = recv(fd, buf, sizeof(buf), 0);
            if (n > 0) {
                if (c.intro.size() + static_cast<size_t>(n) > kMaxParkedBytes) {
                    Log("dropping %s: sent %zu bytes while still waiting for its host",
                        c.peer.c_str(), c.intro.size());
                    Close(fd);
                    return;
                }
                c.intro.insert(c.intro.end(), buf, buf + n);
                continue;
            }
            if (n == 0) { Close(fd); return; }
            if (errno == EINTR) continue;
            if (errno == EAGAIN || errno == EWOULDBLOCK) return;
            Close(fd);
            return;
        }
    }

    if (c.role == Role::Host) {
        uint8_t buf[512];
        for (;;) {
            const ssize_t n = recv(fd, buf, sizeof(buf), 0);
            if (n > 0) continue;
            if (n == 0) { Close(fd); return; }
            if (errno == EINTR) continue;
            if (errno == EAGAIN || errno == EWOULDBLOCK) return;
            Close(fd);
            return;
        }
    }

    for (;;) {
        const size_t have = c.intro.size();
        if (have >= c.introNeed) break;

        uint8_t buf[kMaxIntroFrame];
        const size_t want = c.introNeed - have;
        const ssize_t n = recv(fd, buf, want > sizeof(buf) ? sizeof(buf) : want, 0);
        if (n > 0) { c.intro.insert(c.intro.end(), buf, buf + n); continue; }
        if (n == 0) { Close(fd); return; }
        if (errno == EINTR) continue;
        if (errno == EAGAIN || errno == EWOULDBLOCK) return;
        Close(fd);
        return;
    }

    if (!c.introHeaderDone) {
        MsgHeader hdr{};
        memcpy(&hdr, c.intro.data(), sizeof(hdr));
        if (hdr.magic != kMagic || hdr.length > kMaxIntroFrame) { Close(fd); return; }
        c.introHeaderDone = true;
        c.introNeed = sizeof(MsgHeader) + hdr.length;
        if (c.intro.size() < c.introNeed) { OnReadable(fd); return; }
    }

    MsgHeader hdr{};
    memcpy(&hdr, c.intro.data(), sizeof(hdr));
    std::vector<uint8_t> body(c.intro.begin() + sizeof(MsgHeader), c.intro.end());

    switch (hdr.type) {
        case MSG_RELAY_HOST:   OnRelayHost(c, body);        break;
        case MSG_HELLO:        OnViewerHello(c, hdr, body); break;
        case MSG_RELAY_ACCEPT: OnRelayAccept(c, body);      break;
        case MSG_RELAY_STATUS: OnRelayStatus(c);            break;
        default:
            SendError(c, "unexpected first message");
            Close(fd);
            break;
    }
}

void OnWritable(int fd) {
    auto it = g_conns.find(fd);
    if (it == g_conns.end()) return;
    Conn& c = it->second;

    size_t sent = 0;
    while (sent < c.out.size()) {
        const ssize_t n = send(c.fd, c.out.data() + sent, c.out.size() - sent, MSG_NOSIGNAL);
        if (n > 0) { sent += static_cast<size_t>(n); continue; }
        if (n < 0 && errno == EINTR) continue;
        if (n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) break;
        Close(fd);
        return;
    }
    c.out.erase(c.out.begin(), c.out.begin() + static_cast<long>(sent));

    if (c.out.empty() && c.wantWrite) { c.wantWrite = false; Arm(c); }
}

int Listen(uint16_t port) {
    const int fd = socket(AF_INET6, SOCK_STREAM, 0);
    if (fd < 0) { perror("socket"); return -1; }

    int on = 1;
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
    int off = 0;
    setsockopt(fd, IPPROTO_IPV6, IPV6_V6ONLY, &off, sizeof(off));

    sockaddr_in6 addr{};
    addr.sin6_family = AF_INET6;
    addr.sin6_addr   = in6addr_any;
    addr.sin6_port   = htons(port);
    if (bind(fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) != 0) {
        perror("bind");
        ::close(fd);
        return -1;
    }
    if (listen(fd, 64) != 0) { perror("listen"); ::close(fd); return -1; }
    SetNonBlocking(fd);
    return fd;
}

void OnAccept(int listenFd) {
    for (;;) {
        const int fd = accept(listenFd, nullptr, nullptr);
        if (fd < 0) return;

        SetNonBlocking(fd);
        int on = 1;
        setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on));

        Conn c;
        c.fd    = fd;
        c.since = time(nullptr);
        c.peer  = PeerName(fd);
        g_conns[fd] = std::move(c);

        epoll_event ev{};
        ev.events  = EPOLLIN | EPOLLRDHUP;
        ev.data.fd = fd;
        epoll_ctl(g_epoll, EPOLL_CTL_ADD, fd, &ev);
    }
}

void Reap() {
    const time_t now = time(nullptr);
    std::vector<int> doomed;
    for (auto& [fd, c] : g_conns) {
        if (c.role != Role::Intro && c.role != Role::Waiting) continue;
        if (now - c.since >= kIntroTimeoutSec) doomed.push_back(fd);
    }
    for (int fd : doomed) {
        auto it = g_conns.find(fd);
        if (it != g_conns.end()) {
            Log("dropping %s: %s", it->second.peer.c_str(),
                it->second.role == Role::Waiting ? "host never dialled back"
                                                 : "said nothing");
        }
        Close(fd);
    }
}

}

int main(int argc, char** argv) {
    uint16_t port = 7788;

    for (int i = 1; i < argc; ++i) {
        const std::string a = argv[i];
        if ((a == "--port" || a == "-p") && i + 1 < argc) {
            port = static_cast<uint16_t>(atoi(argv[++i]));
        } else if (a == "--max-sessions" && i + 1 < argc) {
            g_maxSessions = atoi(argv[++i]);
        } else if (a == "--help" || a == "-h") {
            printf("pointdesk-relay [--port 7788] [--max-sessions 64]\n");
            return 0;
        }
    }

    signal(SIGPIPE, SIG_IGN);
    signal(SIGINT, OnSignal);
    signal(SIGTERM, OnSignal);

    const int listenFd = Listen(port);
    if (listenFd < 0) return 1;

    g_epoll = epoll_create1(0);
    if (g_epoll < 0) { perror("epoll_create1"); return 1; }

    epoll_event lev{};
    lev.events  = EPOLLIN;
    lev.data.fd = listenFd;
    epoll_ctl(g_epoll, EPOLL_CTL_ADD, listenFd, &lev);

    Log("pointdesk-relay %s (built %s %s), registration protocol %u",
        kRelayVersion, __DATE__, __TIME__, kRelayVer);
    Log("listening on port %u, up to %d sessions", port, g_maxSessions);

    epoll_event events[256];
    time_t lastReap = time(nullptr);

    while (!g_stop) {
        const int n = epoll_wait(g_epoll, events, 256, 1000);
        if (n < 0) {
            if (errno == EINTR) continue;
            perror("epoll_wait");
            break;
        }

        for (int i = 0; i < n; ++i) {
            const int fd = events[i].data.fd;
            if (fd == listenFd) { OnAccept(listenFd); continue; }

            const uint32_t e = events[i].events;
            if (e & (EPOLLHUP | EPOLLERR)) { Close(fd); continue; }
            if (e & EPOLLOUT) OnWritable(fd);
            if (e & (EPOLLIN | EPOLLRDHUP)) OnReadable(fd);
        }

        const time_t now = time(nullptr);
        if (now - lastReap >= 5) { Reap(); lastReap = now; }
    }

    Log("shutting down");
    std::vector<int> all;
    all.reserve(g_conns.size());
    for (auto& [fd, c] : g_conns) all.push_back(fd);
    for (int fd : all) Close(fd);
    ::close(listenFd);
    ::close(g_epoll);
    return 0;
}