#include "browse_bonjour.hpp"

#include <deque>
#include <iostream>
#include <memory>
#ifdef WINDOWS
#include <WinSock2.h>
#include <Ws2tcpip.h>
#else
#include <netdb.h>
#include <sys/socket.h>
#endif

#include "common/aixlog.hpp"
#include "common/snap_exception.hpp"

using namespace std;

struct DNSServiceRefDeleter
{
    void operator()(DNSServiceRef* ref)
    {
        DNSServiceRefDeallocate(*ref);
        delete ref;
    }
};

typedef std::unique_ptr<DNSServiceRef, DNSServiceRefDeleter> DNSServiceHandle;

string BonjourGetError(DNSServiceErrorType error)
{
    switch (error)
    {
        case kDNSServiceErr_NoError:
            return "NoError";

        default:
        case kDNSServiceErr_Unknown:
            return "Unknown";

        case kDNSServiceErr_NoSuchName:
            return "NoSuchName";

        case kDNSServiceErr_NoMemory:
            return "NoMemory";

        case kDNSServiceErr_BadParam:
            return "BadParam";

        case kDNSServiceErr_BadReference:
            return "BadReference";

        case kDNSServiceErr_BadState:
            return "BadState";

        case kDNSServiceErr_BadFlags:
            return "BadFlags";

        case kDNSServiceErr_Unsupported:
            return "Unsupported";

        case kDNSServiceErr_NotInitialized:
            return "NotInitialized";

        case kDNSServiceErr_AlreadyRegistered:
            return "AlreadyRegistered";

        case kDNSServiceErr_NameConflict:
            return "NameConflict";

        case kDNSServiceErr_Invalid:
            return "Invalid";

        case kDNSServiceErr_Firewall:
            return "Firewall";

        case kDNSServiceErr_Incompatible:
            return "Incompatible";

        case kDNSServiceErr_BadInterfaceIndex:
            return "BadInterfaceIndex";

        case kDNSServiceErr_Refused:
            return "Refused";

        case kDNSServiceErr_NoSuchRecord:
            return "NoSuchRecord";

        case kDNSServiceErr_NoAuth:
            return "NoAuth";

        case kDNSServiceErr_NoSuchKey:
            return "NoSuchKey";

        case kDNSServiceErr_NATTraversal:
            return "NATTraversal";

        case kDNSServiceErr_DoubleNAT:
            return "DoubleNAT";

        case kDNSServiceErr_BadTime:
            return "BadTime";

        case kDNSServiceErr_BadSig:
            return "BadSig";

        case kDNSServiceErr_BadKey:
            return "BadKey";

        case kDNSServiceErr_Transient:
            return "Transient";

        case kDNSServiceErr_ServiceNotRunning:
            return "ServiceNotRunning";

        case kDNSServiceErr_NATPortMappingUnsupported:
            return "NATPortMappingUnsupported";

        case kDNSServiceErr_NATPortMappingDisabled:
            return "NATPortMappingDisabled";

        case kDNSServiceErr_NoRouter:
            return "NoRouter";

        case kDNSServiceErr_PollingMode:
            return "PollingMode";

        case kDNSServiceErr_Timeout:
            return "Timeout";
    }
}

struct mDNSReply
{
    string name, regtype, domain;
};

struct mDNSResolve
{
    string fullName;
    uint16_t port;
};

#define CHECKED(err)                                                                                                                                           \
    if ((err) != kDNSServiceErr_NoError)                                                                                                                       \
        throw SnapException(BonjourGetError(err) + ":" + to_string(__LINE__));

void runService(const DNSServiceHandle& service)
{
    if (!*service)
        return;

    auto socket = DNSServiceRefSockFD(*service);
    fd_set set;
    FD_ZERO(&set);
    FD_SET(socket, &set);

    timeval timeout;
    timeout.tv_sec = 0;
    timeout.tv_usec = 500000;

    while (select(FD_SETSIZE, &set, NULL, NULL, &timeout))
    {
        CHECKED(DNSServiceProcessResult(*service));
        timeout.tv_sec = 0;
        timeout.tv_usec = 500000;
    }
}

bool BrowseBonjour::browse(const string& serviceName, mDNSResult& result, int /*timeout*/)
{
    result.valid = false;
    // Discover
    deque<mDNSReply> replyCollection;
    {
        DNSServiceHandle service(new DNSServiceRef(NULL));
        CHECKED(DNSServiceBrowse(service.get(), 0, 0, serviceName.c_str(), "local.",
                                 [](DNSServiceRef /*service*/, DNSServiceFlags /*flags*/, uint32_t /*interfaceIndex*/, DNSServiceErrorType errorCode,
                                    const char* serviceName, const char* regtype, const char* replyDomain, void* context) {
                                     auto replyCollection = static_cast<deque<mDNSReply>*>(context);

                                     CHECKED(errorCode);
                                     replyCollection->push_back(mDNSReply{string(serviceName), string(regtype), string(replyDomain)});
                                 },
                                 &replyCollection));

        runService(service);
    }

    // Resolve
    deque<mDNSResolve> resolveCollection;
    {
        DNSServiceHandle service(new DNSServiceRef(NULL));
        for (auto& reply : replyCollection)
            CHECKED(DNSServiceResolve(service.get(), 0, 0, reply.name.c_str(), reply.regtype.c_str(), reply.domain.c_str(),
                                      [](DNSServiceRef /*service*/, DNSServiceFlags /*flags*/, uint32_t /*interfaceIndex*/, DNSServiceErrorType errorCode,
                                         const char* /*fullName*/, const char* hosttarget, uint16_t port, uint16_t /*txtLen*/,
                                         const unsigned char* /*txtRecord*/, void* context) {
                                          auto resultCollection = static_cast<deque<mDNSResolve>*>(context);

                                          CHECKED(errorCode);
                                          resultCollection->push_back(mDNSResolve{string(hosttarget), ntohs(port)});
                                      },
                                      &resolveCollection));

        runService(service);
    }

    // DNS/mDNS Resolve
    deque<mDNSResult> resultCollection(resolveCollection.size(), mDNSResult{IPVersion::IPv4, 0, "", "", 0, false});
    {
        DNSServiceHandle service(new DNSServiceRef(NULL));
        unsigned i = 0;
        for (auto& resolve : resolveCollection)
        {
            resultCollection[i].port = resolve.port;
            CHECKED(DNSServiceGetAddrInfo(service.get(), kDNSServiceFlagsLongLivedQuery, 0, kDNSServiceProtocol_IPv4, resolve.fullName.c_str(),
                                          [](DNSServiceRef /*service*/, DNSServiceFlags /*flags*/, uint32_t interfaceIndex, DNSServiceErrorType /*errorCode*/,
                                             const char* hostname, const sockaddr* address, uint32_t /*ttl*/, void* context) {
                                              auto result = static_cast<mDNSResult*>(context);

                                              result->host = string(hostname);
                                              result->ip_version = (address->sa_family == AF_INET) ? (IPVersion::IPv4) : (IPVersion::IPv6);
                                              result->iface_idx = static_cast<int>(interfaceIndex);

                                              char hostIP[NI_MAXHOST];
                                              char hostService[NI_MAXSERV];
                                              if (getnameinfo(address, sizeof(*address), hostIP, sizeof(hostIP), hostService, sizeof(hostService),
                                                              NI_NUMERICHOST | NI_NUMERICSERV) == 0)
                                                  result->ip = string(hostIP);
                                              else
                                                  return;
                                              result->valid = true;
                                          },
                                          &resultCollection[i++]));
        }
        runService(service);
    }

    if (resultCollection.size() == 0)
        return false;

    if (resultCollection.size() != 1)
        LOG(NOTICE) << "Multiple servers found.  Using first" << endl;

    result = resultCollection[0];

    return true;
}

#undef CHECKED