#include "rpc/uds_server.hpp"

#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <unistd.h>

#include <cerrno>
#include <cstring>
#include <random>
#include <string>

#include <nlohmann/json.hpp>

#include <fcntl.h>

#include "rpc/event_loop.hpp"
#include "rpc/systemd_activation.hpp"
#include "version.hpp"

namespace velox::daemon::rpc {

namespace proto = velox::proto;
using nlohmann::json;

namespace {

std::error_code errc(int e) { return std::error_code(e, std::generic_category()); }

// Largest reply we will buffer for a client that is not reading. Past this the client is
// wedged and the connection is dropped rather than growing the daemon's RSS without bound.
constexpr std::size_t kMaxOutBytes = 16 * 1024 * 1024;

std::string uuid4() {
    std::random_device rd;
    std::uniform_int_distribution<std::uint32_t> d;
    std::uint32_t a = d(rd), b = d(rd), c = d(rd), e = d(rd);
    b = (b & 0xFFFF0FFFu) | 0x00004000u;         // version 4
    c = (c & 0x3FFFFFFFu) | 0x80000000u;         // variant 1
    char buf[37];
    std::snprintf(buf, sizeof(buf), "%08x-%04x-%04x-%04x-%04x%08x", a, (b >> 16), (b & 0xFFFF),
                  (c >> 16), (c & 0xFFFF), e);
    return std::string(buf);
}

int major_of(const std::string& semver) {
    // "1.3.0" -> 1. A missing or non-numeric leading component is treated as major -1 so
    // it can never accidentally match the daemon's.
    try {
        return std::stoi(semver.substr(0, semver.find('.')));
    } catch (...) {
        return -1;
    }
}

json rpc_error(const json& id, proto::ErrorCode code, std::string_view msg, json data = nullptr) {
    return proto::make_error(id, code, msg, std::move(data));
}

}  // namespace

UdsServer::UdsServer(EventLoop& loop, proto::Dispatcher& dispatcher, EventHub& hub,
                    std::string socket_path, TaskActionPort* actions)
    : loop_(loop), dispatcher_(dispatcher), hub_(hub), actions_(actions),
      path_(std::move(socket_path)) {}

UdsServer::~UdsServer() {
    for (auto& [fd, c] : conns_) {
        loop_.del_fd(fd);
        ::close(fd);
    }
    if (listen_fd_ >= 0) {
        loop_.del_fd(listen_fd_);
        ::close(listen_fd_);
    }
    if (bound_) ::unlink(path_.c_str());
}

std::error_code UdsServer::start() {
    // velox.socket (systemd user unit, socket activation): the unit binds this path itself
    // before veloxd ever runs and hands the already-listening fd over at fd 3 — the first
    // connection after boot is queued by the kernel rather than refused, and there is no
    // window where a client sees ECONNREFUSED while the daemon is still starting. Skips
    // create/bind/chmod/listen entirely; the socket file's lifecycle (including removal on
    // stop) belongs to the unit, not to us, so bound_ stays false.
    if (const int activated = systemd_activated_fd(); activated >= 0) {
        ::fcntl(activated, F_SETFL, O_NONBLOCK);
        ::fcntl(activated, F_SETFD, FD_CLOEXEC);
        listen_fd_ = activated;
        loop_.add_fd(listen_fd_, kRead, [this](int, unsigned) { on_listener_readable(); });
        return {};
    }

    if (path_.size() + 1 > sizeof(sockaddr_un::sun_path)) return errc(ENAMETOOLONG);

    const int fd = ::socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
    if (fd < 0) return errc(errno);

    // A socket file from a previous run blocks bind() with EADDRINUSE. Single-instance is
    // enforced separately (main.cpp lock socket), so an existing file here is stale.
    ::unlink(path_.c_str());

    sockaddr_un addr{};
    addr.sun_family = AF_UNIX;
    std::memcpy(addr.sun_path, path_.c_str(), path_.size());

    // umask can only tighten; set the mode explicitly after bind so it is exactly 0600.
    if (::bind(fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) != 0) {
        const int e = errno;
        ::close(fd);
        return errc(e);
    }
    bound_ = true;
    if (::chmod(path_.c_str(), 0600) != 0) {
        const int e = errno;
        ::close(fd);
        ::unlink(path_.c_str());
        bound_ = false;
        return errc(e);
    }
    if (::listen(fd, SOMAXCONN) != 0) {
        const int e = errno;
        ::close(fd);
        ::unlink(path_.c_str());
        bound_ = false;
        return errc(e);
    }

    listen_fd_ = fd;
    loop_.add_fd(listen_fd_, kRead, [this](int, unsigned) { on_listener_readable(); });
    return {};
}

void UdsServer::on_listener_readable() {
    for (;;) {
        const int cfd = ::accept4(listen_fd_, nullptr, nullptr, SOCK_NONBLOCK | SOCK_CLOEXEC);
        if (cfd < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) break;
            if (errno == EINTR || errno == ECONNABORTED) continue;
            break;  // EMFILE/ENFILE: stop accepting this pass; loop retries on next readable
        }

        ucred cred{};
        socklen_t len = sizeof(cred);
        if (::getsockopt(cfd, SOL_SOCKET, SO_PEERCRED, &cred, &len) != 0 ||
            cred.uid != ::geteuid()) {
            // Not the same user. The socket mode should already prevent this; refuse hard
            // regardless — this is the authorization on the Unix transport (docs/01 §2).
            ::close(cfd);
            continue;
        }

        auto conn = std::make_unique<Conn>();
        conn->fd = cfd;
        conns_.emplace(cfd, std::move(conn));
        loop_.add_fd(cfd, kRead, [this](int fd, unsigned ev) { on_conn_event(fd, ev); });
    }
}

void UdsServer::on_conn_event(int fd, unsigned events) {
    const auto it = conns_.find(fd);
    if (it == conns_.end()) return;
    Conn& c = *it->second;

    if (events & kWrite) {
        flush(c);
        if (conns_.find(fd) == conns_.end()) return;  // flush closed it
    }

    if (events & kRead) {
        char buf[64 * 1024];
        for (;;) {
            const ssize_t n = ::read(fd, buf, sizeof(buf));
            if (n > 0) {
                auto lines = c.reader.feed(std::string_view(buf, static_cast<std::size_t>(n)));
                const bool overflowed = c.reader.overflowed();
                for (auto& line : lines) {
                    handle_line(c, line);
                    // handle_line may have replied with a fatal error and closed the
                    // connection (e.g. a protocol-major mismatch). Once that happens `c`
                    // is dangling — stop touching it.
                    if (conns_.find(fd) == conns_.end()) return;
                }
                if (overflowed) {
                    close_conn(fd);
                    return;
                }
                continue;
            }
            if (n == 0) {  // peer closed
                close_conn(fd);
                return;
            }
            if (errno == EAGAIN || errno == EWOULDBLOCK) break;
            if (errno == EINTR) continue;
            close_conn(fd);
            return;
        }
    }
}

void UdsServer::handle_line(Conn& c, const std::string& line) {
    json req = json::parse(line, nullptr, /*allow_exceptions=*/false);
    if (req.is_discarded()) {
        queue_reply(c, rpc_error(nullptr, proto::ErrorCode::ParseError, "invalid JSON"));
        return;
    }

    const json id = req.is_object() && req.contains("id") ? req.at("id") : json(nullptr);
    const std::string method =
        req.is_object() && req.contains("method") && req.at("method").is_string()
            ? req.at("method").get<std::string>()
            : std::string{};

    if (!method.empty()) {
        json reply;
        if (handle_session_method(c, method, req, reply)) {
            if (!reply.is_null()) queue_reply(c, reply);
            return;
        }
    }

    if (method == "download.probe") {
        handle_download_probe(c, req);
        return;
    }
    if (method == "download.refreshUrl") {
        handle_download_refreshUrl(c, req);
        return;
    }

    // Everything else: the generated router. It returns a null json for a notification
    // that needs no reply.
    json reply = proto::dispatch(dispatcher_, proto::Transport::Uds, req);
    if (!reply.is_null()) queue_reply(c, reply);
}

void UdsServer::handle_download_probe(Conn& c, const json& request) {
    const json id = request.contains("id") ? request.at("id") : json(nullptr);
    const json params_json = request.contains("params") ? request.at("params") : json::object();

    auto parsed = proto::parse<proto::DownloadProbeParams>(params_json, "params");
    if (!parsed) {
        queue_reply(c, rpc_error(id, proto::ErrorCode::InvalidParams, parsed.error().message,
                                 json{{"path", parsed.error().path}}));
        return;
    }
    if (!actions_) {
        queue_reply(c, rpc_error(id, proto::ErrorCode::InternalError,
                                 "not implemented in this build: download.probe"));
        return;
    }

    const int fd = c.fd;
    actions_->probe_now(
        *parsed, [this, fd, id](proto::HandlerResult<proto::DownloadProbeResult> r) {
            auto it = conns_.find(fd);
            if (it == conns_.end()) return;  // client gone while the probe was outstanding
            if (r) {
                queue_reply(*it->second, proto::make_result(id, *r));
            } else {
                queue_reply(*it->second,
                           rpc_error(id, r.error().code, r.error().message, r.error().data));
            }
        });
}

void UdsServer::handle_download_refreshUrl(Conn& c, const json& request) {
    const json id = request.contains("id") ? request.at("id") : json(nullptr);
    const json params_json = request.contains("params") ? request.at("params") : json::object();

    auto parsed = proto::parse<proto::DownloadRefreshUrlParams>(params_json, "params");
    if (!parsed) {
        queue_reply(c, rpc_error(id, proto::ErrorCode::InvalidParams, parsed.error().message,
                                 json{{"path", parsed.error().path}}));
        return;
    }
    if (!actions_) {
        queue_reply(c, rpc_error(id, proto::ErrorCode::InternalError,
                                 "not implemented in this build: download.refreshUrl"));
        return;
    }

    const int fd = c.fd;
    actions_->refresh_url(
        parsed->taskId, parsed->url, parsed->headers, parsed->cookies,
        [this, fd, id](proto::HandlerResult<proto::DownloadRefreshUrlResult> r) {
            auto it = conns_.find(fd);
            if (it == conns_.end()) return;  // client gone while the probe was outstanding
            if (r) {
                queue_reply(*it->second, proto::make_result(id, *r));
            } else {
                queue_reply(*it->second,
                           rpc_error(id, r.error().code, r.error().message, r.error().data));
            }
        });
}

bool UdsServer::handle_session_method(Conn& c, const std::string& method, const json& request,
                                      json& reply) {
    const json id = request.contains("id") ? request.at("id") : json(nullptr);
    const json params = request.contains("params") ? request.at("params") : json::object();

    if (method == "session.hello") {
        auto p = proto::parse<proto::SessionHelloParams>(params, "params");
        if (!p) {
            reply = rpc_error(id, proto::ErrorCode::InvalidParams, p.error().message,
                              json{{"path", p.error().path}});
            return true;
        }
        const int want = major_of(std::string(proto::kProtocolVersion));
        const int got = major_of(p->protocolVersion);
        if (got != want) {
            reply = rpc_error(
                id, proto::ErrorCode::VersionMismatch,
                "protocol major version mismatch: daemon speaks " + std::to_string(want) +
                    ".x, client speaks " + std::to_string(got < 0 ? 0 : got) + ".x",
                json{{"expected", std::string(proto::kProtocolVersion)},
                     {"actual", p->protocolVersion}});
            c.close_after_flush = true;  // no method is served on a mismatched major
            return true;
        }

        c.hello_ok = true;
        if (c.session_id.empty()) c.session_id = uuid4();

        proto::SessionHelloResult r;
        r.daemonVersion = std::string(velox::daemon::kDaemonVersion);
        r.protocolVersion = std::string(proto::kProtocolVersion);
        r.sessionId = c.session_id;
        r.transport = proto::SessionHelloResultTransport::Uds;
        reply = proto::make_result(id, r);
        return true;
    }

    if (method == "session.subscribe") {
        auto p = proto::parse<proto::SessionSubscribeParams>(params, "params");
        if (!p) {
            reply = rpc_error(id, proto::ErrorCode::InvalidParams, p.error().message,
                              json{{"path", p.error().path}});
            return true;
        }
        if (!c.sub_id) {
            const int fd = c.fd;
            c.sub_id = hub_.subscribe([this, fd](const json& n) {
                if (const auto it = conns_.find(fd); it != conns_.end()) queue_reply(*it->second, n);
            });
        }
        std::vector<proto::Event> events;
        proto::SessionSubscribeResult r;
        r.ok = true;
        for (const auto& ev : p->events) {
            const auto name = proto::to_string(ev);
            r.events.emplace_back(name);
            if (auto e = proto::event_from_string(name)) events.push_back(*e);
        }
        hub_.set_filter(*c.sub_id, std::move(events), p->taskIds);
        reply = proto::make_result(id, r);
        return true;
    }

    return false;  // session.pair falls through to dispatch() -> -32003 on the Unix socket
}

void UdsServer::queue_reply(Conn& c, const json& reply) {
    c.outbuf += frame(reply.dump());
    if (c.outbuf.size() - c.out_off > kMaxOutBytes) {
        close_conn(c.fd);
        return;
    }
    flush(c);
}

void UdsServer::flush(Conn& c) {
    while (c.out_off < c.outbuf.size()) {
        const ssize_t n =
            ::write(c.fd, c.outbuf.data() + c.out_off, c.outbuf.size() - c.out_off);
        if (n > 0) {
            c.out_off += static_cast<std::size_t>(n);
            continue;
        }
        if (n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
            loop_.mod_fd(c.fd, kRead | kWrite);
            return;
        }
        if (n < 0 && errno == EINTR) continue;
        close_conn(c.fd);
        return;
    }

    c.outbuf.clear();
    c.out_off = 0;
    if (c.close_after_flush) {
        close_conn(c.fd);
        return;
    }
    loop_.mod_fd(c.fd, kRead);
}

void UdsServer::close_conn(int fd) {
    if (const auto it = conns_.find(fd); it != conns_.end()) {
        if (it->second->sub_id) hub_.unsubscribe(*it->second->sub_id);
        loop_.del_fd(fd);
        ::close(fd);
        conns_.erase(it);
    }
}

}  // namespace velox::daemon::rpc