#define WIN32_LEAN_AND_MEAN
#include <winsock2.h>
#include <ws2tcpip.h>
#include <iphlpapi.h>
#include <windows.h>

#include "mdns_advertise.h"
#include "daemon_log.h"
#include "mdns.h"

#include <cstring>
#include <cstdio>
#include <thread>
#include <chrono>

// Two-service record block: one set for _oscjson._tcp, one for _osc._udp.
// Each service advertises PTR (service type -> instance), SRV (instance ->
// host+port), and TXT (metadata). Shared A record resolves host -> 127.0.0.1.
struct ServiceBlock {
    mdns_string_t    service;             // e.g. "_oscjson._tcp.local."
    mdns_string_t    service_instance;    // e.g. "Beyond Backglow._oscjson._tcp.local."
    mdns_string_t    hostname_qualified;  // e.g. "BeyondBackglow.local."
    mdns_record_t    record_ptr;
    mdns_record_t    record_srv;
    mdns_record_t    record_a;
    mdns_record_txt_t txt_record[2];
    int              port;
};

// Instance-scoped helper state held by MdnsAdvertiser. Stored in .cpp because
// the types are C-ABI and leaking them through the header complicates callers.
// IN-01 (iter-3) fix: MdnsImpl is no longer file-scope `g_impl`; each
// MdnsAdvertiser owns its own std::unique_ptr<MdnsImpl> via the forward-decl
// in mdns_advertise.h. Multiple advertiser instances can now coexist without
// silently stomping each other at Start()-time.
struct MdnsImpl {
    std::string   instance_lower;   // "beyond backglow" (used as TXT value)
    std::string   service_tcp_str;  // "_oscjson._tcp.local."
    std::string   service_udp_str;  // "_osc._udp.local."
    std::string   inst_tcp_str;     // "Beyond Backglow._oscjson._tcp.local."
    std::string   inst_udp_str;     // "Beyond Backglow._osc._udp.local."
    std::string   hostname_q_str;   // "BeyondBackglow.local."
    sockaddr_in   addr_ipv4;
    ServiceBlock  blk_tcp;
    ServiceBlock  blk_udp;
};

namespace {
mdns_string_t MakeStr(const std::string& s) {
    mdns_string_t out;
    out.str = s.c_str();
    out.length = s.length();
    return out;
}
} // namespace

// IN-01 (iter-3): ctor/dtor defined out-of-line in the .cpp so std::unique_ptr<MdnsImpl>
// sees a complete type. The header only forward-declares MdnsImpl.
MdnsAdvertiser::MdnsAdvertiser() = default;
MdnsAdvertiser::~MdnsAdvertiser() = default;

static void BuildServiceBlock(ServiceBlock& blk,
                               const std::string& service_str,
                               const std::string& instance_str,
                               const std::string& hostname_q,
                               const sockaddr_in& addr,
                               uint16_t port) {
    blk.service            = MakeStr(service_str);
    blk.service_instance   = MakeStr(instance_str);
    blk.hostname_qualified = MakeStr(hostname_q);
    blk.port               = port;

    // PTR: _oscjson._tcp.local. -> <instance>._oscjson._tcp.local.
    blk.record_ptr.name        = blk.service;
    blk.record_ptr.type        = MDNS_RECORDTYPE_PTR;
    blk.record_ptr.data.ptr.name = blk.service_instance;
    blk.record_ptr.rclass      = 0;
    blk.record_ptr.ttl         = 0;

    // SRV: <instance> -> hostname:port
    blk.record_srv.name              = blk.service_instance;
    blk.record_srv.type              = MDNS_RECORDTYPE_SRV;
    blk.record_srv.data.srv.name     = blk.hostname_qualified;
    blk.record_srv.data.srv.priority = 0;
    blk.record_srv.data.srv.weight   = 0;
    blk.record_srv.data.srv.port     = port;
    blk.record_srv.rclass            = 0;
    blk.record_srv.ttl               = 0;

    // A: hostname -> 127.0.0.1
    blk.record_a.name         = blk.hostname_qualified;
    blk.record_a.type         = MDNS_RECORDTYPE_A;
    blk.record_a.data.a.addr  = addr;
    blk.record_a.rclass       = 0;
    blk.record_a.ttl          = 0;

    // TXT: two attributes -- "txtvers=1" and "service=Beyond Backglow".
    // VRChat OSCQuery doesn't require specific TXT keys but all reference
    // implementations include txtvers. Keep it minimal.
    static const char kTxtVersKey[] = "txtvers";
    static const char kTxtVersVal[] = "1";
    blk.txt_record[0].key.str = kTxtVersKey;
    blk.txt_record[0].key.length = sizeof(kTxtVersKey) - 1;
    blk.txt_record[0].value.str = kTxtVersVal;
    blk.txt_record[0].value.length = sizeof(kTxtVersVal) - 1;
    blk.txt_record[1].key.str = nullptr;  // sentinel -- only one TXT entry used
    blk.txt_record[1].key.length = 0;
}

// Announce one service block via mdns_announce_multicast. Returns 0 on success.
static int AnnounceBlock(int sock, const ServiceBlock& blk) {
    char buffer[2048];
    // additional[0] = SRV, additional[1] = A, additional[2] = TXT
    mdns_record_t additional[3];
    additional[0] = blk.record_srv;
    additional[1] = blk.record_a;

    additional[2].name      = blk.service_instance;
    additional[2].type      = MDNS_RECORDTYPE_TXT;
    additional[2].data.txt  = blk.txt_record[0];
    additional[2].rclass    = 0;
    additional[2].ttl       = 0;

    return mdns_announce_multicast(sock, buffer, sizeof(buffer),
                                   blk.record_ptr,
                                   nullptr, 0,
                                   additional, 3);
}

// Send goodbye packets so other clients invalidate our cache.
static int GoodbyeBlock(int sock, const ServiceBlock& blk) {
    char buffer[2048];
    mdns_record_t additional[3];
    additional[0] = blk.record_srv;
    additional[1] = blk.record_a;
    additional[2].name      = blk.service_instance;
    additional[2].type      = MDNS_RECORDTYPE_TXT;
    additional[2].data.txt  = blk.txt_record[0];
    additional[2].rclass    = 0;
    additional[2].ttl       = 0;

    return mdns_goodbye_multicast(sock, buffer, sizeof(buffer),
                                  blk.record_ptr,
                                  nullptr, 0,
                                  additional, 3);
}

bool MdnsAdvertiser::Start(uint16_t oscUdpPort, uint16_t httpTcpPort, const char* serviceName) {
    // WR-04 fix: guard against double-start. If Start() is ever called a
    // second time without an intervening Stop() — a future retry path, a
    // test harness, etc. — the previous MdnsImpl heap allocation and mDNS
    // socket would leak silently. Tear down any prior state first.
    // IN-01 (iter-3): the check now uses the per-instance m_impl instead of a
    // file-scope global, so multiple MdnsAdvertiser instances no longer stomp
    // each other's state.
    if (m_impl || m_socket >= 0 || m_thread.joinable()) {
        DaemonLog("mdns: Start() called with prior state live — calling Stop() first");
        Stop();
    }

    // IN-03 (re-review): previously cached oscUdpPort / httpTcpPort /
    // serviceName as members, but they were never read. Use the args directly
    // below and let the ServiceBlock / MdnsImpl structs hold the durable
    // state they need.

    // D-20: bind mDNS socket on 127.0.0.1 only.
    sockaddr_in saddr{};
    saddr.sin_family = AF_INET;
    saddr.sin_port = htons(MDNS_PORT);
    if (inet_pton(AF_INET, "127.0.0.1", &saddr.sin_addr) != 1) {
        DaemonLog("mdns: inet_pton 127.0.0.1 failed");
        return false;
    }

    m_socket = mdns_socket_open_ipv4(&saddr);
    if (m_socket < 0) {
        DaemonLog("mdns: socket open failed");
        return false;
    }

    // Build both service blocks up-front so ListenLoop / AnnounceBlock only
    // dereference stable strings. All strings must outlive the thread.
    // IN-01 (iter-3): state is per-instance, owned via unique_ptr.
    m_impl = std::make_unique<MdnsImpl>();
    m_impl->service_tcp_str = "_oscjson._tcp.local.";
    m_impl->service_udp_str = "_osc._udp.local.";
    m_impl->inst_tcp_str    = std::string(serviceName) + "._oscjson._tcp.local.";
    m_impl->inst_udp_str    = std::string(serviceName) + "._osc._udp.local.";
    // Sanitize hostname: strip spaces per RFC 6762.
    std::string hostname = serviceName;
    for (auto& c : hostname) if (c == ' ') c = '-';
    m_impl->hostname_q_str = hostname + ".local.";

    m_impl->addr_ipv4.sin_family = AF_INET;
    m_impl->addr_ipv4.sin_port   = 0;
    inet_pton(AF_INET, "127.0.0.1", &m_impl->addr_ipv4.sin_addr);

    BuildServiceBlock(m_impl->blk_tcp, m_impl->service_tcp_str, m_impl->inst_tcp_str,
                      m_impl->hostname_q_str, m_impl->addr_ipv4, httpTcpPort);
    BuildServiceBlock(m_impl->blk_udp, m_impl->service_udp_str, m_impl->inst_udp_str,
                      m_impl->hostname_q_str, m_impl->addr_ipv4, oscUdpPort);

    // D-18 LOCKED: announce BOTH service types. Failure to announce either is
    // a hard error -- fall back to D-19 legacy path (caller handles by observing
    // Start() == false).
    if (AnnounceBlock(m_socket, m_impl->blk_tcp) != 0) {
        DaemonLog("mdns: announce _oscjson._tcp failed");
        mdns_socket_close(m_socket);
        m_socket = -1;
        m_impl.reset();
        return false;
    }
    if (AnnounceBlock(m_socket, m_impl->blk_udp) != 0) {
        DaemonLog("mdns: announce _osc._udp failed");
        GoodbyeBlock(m_socket, m_impl->blk_tcp);  // best-effort cleanup
        mdns_socket_close(m_socket);
        m_socket = -1;
        m_impl.reset();
        return false;
    }

    DaemonLog("mdns: advertising '%s' on _oscjson._tcp:%u + _osc._udp:%u (127.0.0.1)",
              serviceName, (unsigned)httpTcpPort, (unsigned)oscUdpPort);

    // WR-04 fix: std::thread construction can throw std::system_error on
    // thread-creation failure. Without this guard, the socket, goodbye-needed
    // announce state, and heap-allocated impl all leak through the unwind
    // because Stop() is never reached.
    m_stop.store(false);
    try {
        m_thread = std::thread(&MdnsAdvertiser::ListenLoop, this);
    } catch (...) {
        DaemonLog("mdns: thread ctor threw — rolling back socket + impl");
        GoodbyeBlock(m_socket, m_impl->blk_tcp);
        GoodbyeBlock(m_socket, m_impl->blk_udp);
        mdns_socket_close(m_socket);
        m_socket = -1;
        m_impl.reset();
        throw;
    }
    return true;
}

void MdnsAdvertiser::Stop() {
    m_stop.store(true);
    if (m_thread.joinable()) m_thread.join();
    if (m_socket >= 0 && m_impl) {
        GoodbyeBlock(m_socket, m_impl->blk_tcp);
        GoodbyeBlock(m_socket, m_impl->blk_udp);
        mdns_socket_close(m_socket);
        m_socket = -1;
    }
    m_impl.reset();
}

void MdnsAdvertiser::ListenLoop() {
    // IN-06 (re-review): add a periodic re-announce in the idle path. The
    // library still drops queries (nullptr callback), but unsolicited
    // announcements every ~90 s keep passive-cache peers (VRChat today, any
    // conformant discovery agent tomorrow) from aging our records out of
    // their caches. Value chosen below the typical 2 min TTL so at least one
    // refresh lands per cache lifetime. The previous implementation announced
    // exactly once in Start() and then silently decayed for long sessions.
    char recvBuf[2048];
    using clock = std::chrono::steady_clock;
    const auto kReannouncePeriod = std::chrono::seconds(90);
    auto lastAnnounce = clock::now();
    while (!m_stop.load()) {
        size_t n = mdns_socket_listen(m_socket, recvBuf, sizeof(recvBuf),
                                      nullptr, nullptr);
        // WR-03 (iter-4) fix: ALWAYS yield 100 ms after each listen call, not
        // only when no datagram is returned. On chatty multicast segments
        // (corporate LANs, other OSCQuery-capable apps) `n` is regularly
        // non-zero and the library's non-blocking socket gave us no inherent
        // backpressure — we would receive, drop (nullptr callback), re-check
        // the deadline, and immediately re-enter mdns_socket_listen, pinning
        // a full CPU core for the whole session. Unconditional 100 ms is
        // short enough that the 90 s re-announce deadline below stays accurate
        // (0.1 % granularity, well inside the TTL margin) and long enough to
        // prevent busy-spin on heavy multicast chatter.
        (void)n;  // packets dropped at the library level (nullptr callback)
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        // Library fires our callback inline when a query arrives. With a
        // nullptr callback the library just drops queries -- our announce
        // packets are unsolicited multicast, which is sufficient for VRChat's
        // passive discovery loop. A proper responder would take a callback
        // and answer queries via mdns_query_answer_multicast; sketched here
        // but not required because announce covers the primary path.

        // IN-06: periodic re-announce. Guarded by m_impl / m_socket so we
        // never dereference post-Stop state, and ignores announce errors
        // (a failed re-announce is non-fatal — next period will retry).
        // IN-01 (iter-3): was g_impl; state is now per-instance via unique_ptr.
        // IN-03 (iter-3) fix: previously assigned `lastAnnounce = clock::now()`
        // AFTER the two AnnounceBlock calls, which meant each cycle's period
        // drifted forward by whatever time the announces took. Now bump the
        // deadline by exactly kReannouncePeriod so the effective cadence is a
        // rigid 90 s regardless of announce latency. The comparison is shifted
        // to `now >= lastAnnounce + period` (deadline form).
        if (!m_stop.load() && m_impl && m_socket >= 0 &&
            clock::now() >= lastAnnounce + kReannouncePeriod) {
            AnnounceBlock(m_socket, m_impl->blk_tcp);
            AnnounceBlock(m_socket, m_impl->blk_udp);
            lastAnnounce += kReannouncePeriod;  // deadline-based, no drift
        }
    }
}