/***
    Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: AuNetAdapter.Linux.cpp
    Date: 2022-11-15
    Author: Reece
***/
#include "Networking.hpp"
#include "AuNetAdapter.hpp"
#include "AuNetEndpoint.hpp"

#include <asm/types.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <sys/socket.h>

namespace Aurora::IO::Net
{
    AuString NetAdapter::GetHostname()
    {
        char name[128];
        if (::gethostname(name, AuArraySize(name)) != 0)
        {
            return "";
        }       

        return name;
    }

    // Based on https://gist.github.com/Yawning/c70d804d4b8ae78cc698

    struct NetlinkDevice
    {
        NetlinkDevice();
        ~NetlinkDevice();

        bool Connect();

        AuSInt Send(const void *pBuffer, size_t uLength);
        AuSInt Recv(void *pBuffer, size_t uLength);

        AuList<NetAdapter> GetAdaptersForFamily(sa_family_t family);
        AuList<AuPair<int, NetEndpoint>> ReadRoutes(int index, int family, AuString &name);

        void UpdateAdapterInfo(int family, NetAdapter &adapter);

        int fd { -1 };
    };

    static NetlinkDevice gNetlinkDevice;

    NetlinkDevice::NetlinkDevice()
    {

    }

    NetlinkDevice::~NetlinkDevice()
    {
        if (this->fd != -1)
        {
            ::close(this->fd);
        }
    }

    bool NetlinkDevice::Connect()
    {
        sockaddr_nl nladdr {};

        this->fd = socket(AF_NETLINK, SOCK_DGRAM, NETLINK_ROUTE);
        if (this->fd < 0)
        {
            return false;
        }

        int flags = ::fcntl(this->fd, F_GETFL, 0);
        if (flags != -1)
        {
            ::fcntl(this->fd, F_SETFL, flags | FD_CLOEXEC);
        }

        nladdr.nl_family = AF_NETLINK;
        if (::bind(this->fd, (struct sockaddr *) &nladdr, sizeof(nladdr)) != 0)
        {
            return false;
        }

        return true;
    }

    AuSInt NetlinkDevice::Send(const void *pBuffer, size_t uLength)
    {
        sockaddr_nl nladdr {};
        msghdr msg {};
        iovec vec {};

        nladdr.nl_family = AF_NETLINK;

        msg.msg_name = &nladdr;
        msg.msg_namelen = sizeof(nladdr);
        vec.iov_base = (char *)pBuffer;
        vec.iov_len = uLength;
        msg.msg_iov = &vec;
        msg.msg_iovlen = 1;

        return ::sendmsg(fd, &msg, 0);
    }

    AuSInt NetlinkDevice::Recv(void *pBuffer, size_t uLength)
    {
        sockaddr_nl nladdr {};
        msghdr msg {};
        iovec vec;

        nladdr.nl_family = AF_NETLINK;

        msg.msg_name = &nladdr;
        msg.msg_namelen = sizeof(nladdr);
        vec.iov_base = pBuffer;
        vec.iov_len = uLength;
        msg.msg_iov = &vec;
        msg.msg_iovlen = 1;

        auto ret = ::recvmsg(fd, &msg, 0);
        if (msg.msg_flags & MSG_TRUNC)
        {
            return -1;
        }

        return ret;
    }

    struct addr_t
    {
        sa_family_t family;
        union
        {
            in_addr in_addr;
            in6_addr in6_addr;
        } u;
    };

    AuList<NetAdapter> NetlinkDevice::GetAdaptersForFamily(sa_family_t family)
    {
        static const auto kRouteBufferSize = 12 * 1024;

        struct nlmsghdr *hdr;
        struct rtmsg *rt;
        ssize_t len;
        size_t reqlen;

        reqlen    = NLMSG_SPACE(sizeof(*rt));
        auto pReq = AuMakeSharedArray<AuUInt8>(kRouteBufferSize);
        if (!pReq)
        {
            SysPushErrorMemory();
            return {};
        }

        hdr = (struct nlmsghdr *)pReq.get();
        hdr->nlmsg_len = NLMSG_LENGTH(sizeof(*rt));
        hdr->nlmsg_flags = NLM_F_REQUEST | NLM_F_DUMP;
        hdr->nlmsg_type = RTM_GETROUTE;

        rt = (struct rtmsg *)NLMSG_DATA(hdr);
        rt->rtm_family = family;
        rt->rtm_table = RT_TABLE_MAIN | RT_TABLE_LOCAL;

        if (this->Send(pReq.get(), reqlen) < 0)
        {
            SysPushErrorIO("netlink error");
            return {};
        }

        auto pResp = AuMakeSharedArray<AuUInt8>(kRouteBufferSize);
        if (!pResp)
        {
            SysPushErrorMemory();
            return {};
        }

        AuList<NetAdapter> ret;
        while (1)
        {
            len = this->Recv(pResp.get(), kRouteBufferSize);
            if (len < 0)
            {
                SysPushErrorIO("netlink error");
                return {};
            }

            for (hdr = (struct nlmsghdr *)pResp.get();
                 NLMSG_OK(hdr, len);
                 hdr = NLMSG_NEXT(hdr, len))
            {
                NetAdapter adapter;

                if (hdr->nlmsg_type == NLMSG_DONE)
                {
                    return ret;
                }

                if (hdr->nlmsg_type == NLMSG_ERROR)
                {
                    return {};
                }

                rt = (struct rtmsg *)NLMSG_DATA(hdr);

                auto attr    = RTM_RTA(rt);
                auto attrlen = RTM_PAYLOAD(hdr);
                for (; RTA_OK(attr, attrlen); attr = RTA_NEXT(attr, attrlen))
                {
                    const size_t expectedAddressLength = (family == AF_INET) ? 4 : 16;

                    size_t dstlen = RTA_PAYLOAD(attr);

                    switch (attr->rta_type)
                    {
                    case RTA_GATEWAY:
                    case RTA_DST:
                    case RTA_SRC:
                    {
                        if (dstlen != expectedAddressLength)
                        {
                            continue;
                        }

                        NetEndpoint ep;
                        addr_t *address2 = (addr_t*)ep.hint;
                        address2->family = family;

                        AuMemcpy(&address2->u.in_addr, RTA_DATA(attr), dstlen);
                        DeoptimizeEndpoint(ep);

                        if (attr->rta_type == RTA_GATEWAY)
                        {
                            adapter.gateway = ep.ip;
                        }

                        if (attr->rta_type == RTA_DST)
                        {
                            adapter.address = ep.ip;
                        }

                        break;
                    }
                    case RTA_OIF:
                    {
                        adapter.index = *(int*)RTA_DATA(attr);
                        break;
                    }
                    default:
                    {
                        break;
                    }
                    }
                }

                ret.push_back(adapter);
            }
        }

        return ret;
    }

    AuList<AuPair<int, NetEndpoint>> NetlinkDevice::ReadRoutes(int index, int family, AuString &name)
    {
        static const auto kAddrBufferSize = 12 * 1024;

        AuList<AuPair<int, NetEndpoint>> ret;
        nlmsghdr *hdr {};
        ifaddrmsg *ifa {};
        ssize_t len;
        size_t reqlen;

        /* Allocate space for the request. */
        reqlen    = NLMSG_SPACE(sizeof(*ifa));
        auto pReq = AuMakeSharedArray<AuUInt8>(kAddrBufferSize);
        if (!pReq)
        {
            SysPushErrorMemory();
            return {};
        }

        hdr = (struct nlmsghdr *)pReq.get();
        hdr->nlmsg_len = NLMSG_LENGTH(sizeof(*ifa));
        hdr->nlmsg_flags = NLM_F_REQUEST | NLM_F_ROOT;
        hdr->nlmsg_type = RTM_GETADDR;

        ifa = (struct ifaddrmsg *)NLMSG_DATA(hdr);
        ifa->ifa_family = family;
        ifa->ifa_index = index;
   
        if (this->Send(pReq.get(), reqlen) < 0)
        {
            SysPushErrorIO("netlink error");
            return {};
        }

        auto pResp = AuMakeSharedArray<AuUInt8>(kAddrBufferSize);
        if (!pResp)
        {
            SysPushErrorMemory();
            return {};
        }

        while (1)
        {
            len = this->Recv(pResp.get(), kAddrBufferSize);
            if (len < 0)
            {
                SysPushErrorIO("netlink error");
                return {};
            }

            for (hdr = (struct nlmsghdr *)pResp.get(); NLMSG_OK(hdr, len); hdr = NLMSG_NEXT(hdr, len))
            {
                if (hdr->nlmsg_type == NLMSG_DONE)
                {
                    return ret;
                }

                if (hdr->nlmsg_type == NLMSG_ERROR)
                {
                    return {};
                }

                ifa = (struct ifaddrmsg *)NLMSG_DATA(hdr);
                if ((ifa->ifa_family != family) ||
                    (ifa->ifa_index != index))
                {
                    continue;
                }

                if (ifa->ifa_scope != RT_SCOPE_UNIVERSE)
                {
                    continue;
                }

                auto attr    = IFA_RTA(ifa);
                auto attrlen = RTM_PAYLOAD(hdr);
                for (; RTA_OK(attr, attrlen); attr = RTA_NEXT(attr, attrlen))
                {
                    size_t dstlen = RTA_PAYLOAD(attr);
                    const size_t addrlen = (family == AF_INET) ? sizeof(struct in_addr) : sizeof(struct in6_addr);

                    if (attr->rta_type == IFA_LABEL)
                    {
                        name = (const char *)RTA_DATA(attr);
                    }
                    else if (addrlen == dstlen)
                    {
                        NetEndpoint ep;
                        addr_t *address2 = (addr_t*)ep.hint;
                        address2->family = family;

                        AuMemcpy(&address2->u.in_addr, RTA_DATA(attr), addrlen);
                        DeoptimizeEndpoint(ep);
                        ret.push_back(AuMakePair(attr->rta_type, ep));
                    }
                }
            }
        }

        return ret;
    }

    void NetlinkDevice::UpdateAdapterInfo(int family, NetAdapter &adapter)
    {
        static const auto kAddrBufferSize = 12 * 1024;

        nlmsghdr *hdr {};
        ifinfomsg *ifa {};

        auto reqlen = NLMSG_SPACE(sizeof(*ifa));
        auto pReq   = AuMakeSharedArray<AuUInt8>(kAddrBufferSize);
        if (!pReq)
        {
            SysPushErrorMemory();
            return;
        }

        hdr = (struct nlmsghdr *)pReq.get();
        hdr->nlmsg_len = NLMSG_LENGTH(sizeof(*ifa));
        hdr->nlmsg_flags = NLM_F_REQUEST | NLM_F_DUMP;
        hdr->nlmsg_type = RTM_GETLINK;

        ifa = (struct ifinfomsg *)NLMSG_DATA(hdr);
        ifa->ifi_family = family;
        ifa->ifi_index = adapter.index;

        if (this->Send(pReq.get(), reqlen) < 0)
        {
            SysPushErrorIO("netlink error");
            return;
        }

        auto pResp = AuMakeSharedArray<AuUInt8>(kAddrBufferSize);
        if (!pResp)
        {
            SysPushErrorMemory();
            return ;
        }

        while (1)
        {
            auto len = this->Recv(pResp.get(), kAddrBufferSize);
            if (len < 0)
            {
                SysPushErrorIO("netlink error");
                return;
            }

            for (hdr = (struct nlmsghdr *)pResp.get(); NLMSG_OK(hdr, len); hdr = NLMSG_NEXT(hdr, len))
            {
                if (hdr->nlmsg_type == NLMSG_DONE)
                {
                    return;
                }

                if (hdr->nlmsg_type == NLMSG_ERROR)
                {
                    return;
                }

                ifa = (struct ifinfomsg *)NLMSG_DATA(hdr);
                if ((ifa->ifi_index != adapter.index))
                {
                    continue;
                }

                auto attr    = IFLA_RTA(ifa);
                auto attrlen = hdr->nlmsg_len - NLMSG_LENGTH(sizeof(*ifa));

                for (; RTA_OK(attr, attrlen); attr = RTA_NEXT(attr, attrlen)) {

                    size_t dstlen = RTA_PAYLOAD(attr);

                    const size_t addrlen = (family == AF_INET) ? sizeof(struct in_addr) : sizeof(struct in6_addr);

                    if (attr->rta_type == IFLA_IFNAME)
                    {
                        adapter.device = (const char *)RTA_DATA(attr);
                    }

                    // TODO: we can pull MTU and other stats from here.
                }
            }
        }
    }

    static AuList<AuSPtr<INetAdapter>> GetForFamily(int family)
    {
        AuList<AuSPtr<INetAdapter>> ret;
        if (gNetlinkDevice.fd == -1)
        {
            if (!gNetlinkDevice.Connect())
            {
                return {};
            }
        }

        AuList<IPAddress> dups;
        auto adapters = gNetlinkDevice.GetAdaptersForFamily(family);
        for (auto &adapter : adapters)
        {
            gNetlinkDevice.UpdateAdapterInfo(family, adapter);

            auto pAdapter = AuMakeShared<NetAdapter>(adapter);
            if (!pAdapter)
            {
                SysPushErrorMemory();
                return {};
            }

            bool bBreak = false;

            auto ips = gNetlinkDevice.ReadRoutes(adapter.index, family, pAdapter->name);
            for (const auto &[type, endpoint] : ips)
            {
                if (type == IFA_ANYCAST)
                {
                    pAdapter->anycast = endpoint.ip;
                }

                if (type == IFA_ADDRESS)
                {
                    bBreak |= AuExists(dups, endpoint.ip);
                    dups.push_back(endpoint.ip);
                    pAdapter->address = endpoint.ip;
                }

                if (type == IFA_BROADCAST)
                {
                    pAdapter->broadcast = endpoint.ip;
                }
            }

            if (bBreak)
            {
                continue;
            }

            ret.push_back(pAdapter);
        }

        return ret;
    }

    AuList<AuSPtr<INetAdapter>> NetAdapter::GetIPv4s()
    {
        return GetForFamily(AF_INET);
    }

    AuList<AuSPtr<INetAdapter>> NetAdapter::GetIPv6s()
    {
        return GetForFamily(AF_INET6);
    }
}