// tunnelmanager_win.cpp — Windows WireGuard tunnel implementation.
//
// Uses Wintun (kernel TUN driver) for the virtual NIC and boringtun (via FFI)
// for the WireGuard crypto.  Requires elevation: the process must run as
// Administrator (see DragonMoonlight.manifest) so that WintunCreateAdapter
// can install the kernel driver.
//
// Packet format on Windows: raw IP (no 4-byte AF header unlike macOS utun).
//
// I/O model
// ---------
//   TUN -> UDP  thread: WintunReceivePacket -> wireguard_write -> sendto
//   UDP -> TUN  thread: recvfrom -> wireguard_read -> WintunAllocSendPacket/Send
//   Ticker      thread: wireguard_tick every 100 ms (keepalives / handshake)
//
// Shutdown: SetEvent(stopEvent) unblocks WaitForMultipleObjects in all threads.
//
// Tunnel-up detection: udpToTunThread emits tunnelUp() the first time a
// decrypted packet arrives.  A 1500 ms fallback timer fires the same signal
// if no packet has arrived yet, so the UI moves out of "Connecting" state.
// This corresponds to the roadmap item "connected() signal from first
// received packet rather than a fixed timer".

#include "tunnelmanager.h"

#ifdef Q_OS_WIN

#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <iphlpapi.h>
#include <netioapi.h>

#include <atomic>
#include <cstring>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

#include <QDebug>
#include <QPointer>
#include <QTimer>

#include "boringtun_ffi.h"
#include "wireguardconfig.h"
#include "wintun.h"

// ----- helpers --------------------------------------------------------------

static std::string lastWinError(const char *context)
{
    char buf[256] = {};
    FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                   nullptr, GetLastError(), 0, buf, sizeof(buf), nullptr);
    return std::string(context) + ": " + buf;
}

// Parse "x.x.x.x/prefix" into binary address + prefix length.  Default prefix
// is 32 if none is supplied.
static bool parseCIDR(const std::string &cidr, IN_ADDR *addr, UINT8 *prefix)
{
    auto slash = cidr.find('/');
    std::string ip = (slash != std::string::npos) ? cidr.substr(0, slash) : cidr;
    *prefix        = (slash != std::string::npos)
                       ? static_cast<UINT8>(std::stoi(cidr.substr(slash + 1)))
                       : 32;
    return InetPtonA(AF_INET, ip.c_str(), addr) == 1;
}

// ----- private state --------------------------------------------------------

struct WinTunnelPriv {
    WintunLib              wintun;
    WINTUN_ADAPTER_HANDLE  adapter      = nullptr;
    WINTUN_SESSION_HANDLE  session      = nullptr;
    wireguard_tunnel      *wg           = nullptr;
    SOCKET                 udpSock      = INVALID_SOCKET;
    HANDLE                 stopEvent    = nullptr;
    NET_LUID               adapterLuid  {};

    std::atomic<bool>      running      {false};
    std::atomic<bool>      gotFirstPkt  {false};
    std::thread            thdTunToUdp;
    std::thread            thdUdpToTun;
    std::thread            thdTicker;

    WireGuardConfig        cfg;

    // Owning TunnelManager - so I/O threads can post status to the Qt main
    // thread once the first decrypted packet arrives.
    QPointer<TunnelManager> owner;
};

// ----- TUN -> UDP thread ----------------------------------------------------

static void tunToUdpThread(WinTunnelPriv *p)
{
    sockaddr_in peer {};
    peer.sin_family = AF_INET;
    InetPtonA(AF_INET, p->cfg.endpointHost().toUtf8().constData(), &peer.sin_addr);
    peer.sin_port = htons(static_cast<u_short>(p->cfg.endpointPort()));

    HANDLE readEvent = p->wintun.GetReadWaitEvent(p->session);
    HANDLE events[2] = { readEvent, p->stopEvent };

    static constexpr size_t BUF = 65536;
    std::vector<uint8_t> dst(BUF);

    while (p->running.load()) {
        DWORD wait = WaitForMultipleObjects(2, events, FALSE, INFINITE);
        if (wait != WAIT_OBJECT_0)
            break;

        DWORD pktSize = 0;
        BYTE *pkt = p->wintun.ReceivePacket(p->session, &pktSize);
        if (!pkt)
            continue;

        size_t dstLen = dst.size();
        wireguard_result r = wireguard_write(p->wg, pkt, pktSize,
                                             dst.data(), &dstLen);
        p->wintun.ReleaseReceivePacket(p->session, pkt);

        if (r.op == WRITE_TO_NETWORK && r.size > 0) {
            sendto(p->udpSock,
                   reinterpret_cast<const char *>(dst.data()),
                   static_cast<int>(r.size),
                   0, reinterpret_cast<sockaddr *>(&peer), sizeof(peer));
        }
    }
}

// ----- UDP -> TUN thread ----------------------------------------------------

static void udpToTunThread(WinTunnelPriv *p)
{
    static constexpr size_t BUF = 65536;
    std::vector<uint8_t> enc(BUF), plain(BUF);

    sockaddr_in peer {};
    peer.sin_family = AF_INET;
    InetPtonA(AF_INET, p->cfg.endpointHost().toUtf8().constData(), &peer.sin_addr);
    peer.sin_port = htons(static_cast<u_short>(p->cfg.endpointPort()));

    WSAEVENT sockEvent = WSACreateEvent();
    WSAEventSelect(p->udpSock, sockEvent, FD_READ);
    HANDLE events[2] = { sockEvent, p->stopEvent };

    while (p->running.load()) {
        DWORD wait = WaitForMultipleObjects(2, events, FALSE, INFINITE);
        if (wait != WAIT_OBJECT_0)
            break;
        WSAResetEvent(sockEvent);

        sockaddr_in from {};
        int fromLen = sizeof(from);
        int n = recvfrom(p->udpSock,
                         reinterpret_cast<char *>(enc.data()),
                         static_cast<int>(enc.size()),
                         0, reinterpret_cast<sockaddr *>(&from), &fromLen);
        if (n <= 0)
            continue;

        size_t plainLen = plain.size();
        wireguard_result r = wireguard_read(p->wg, enc.data(),
                                            static_cast<size_t>(n),
                                            plain.data(), &plainLen);

        if ((r.op == WRITE_TO_TUNNEL_IPV4 || r.op == WRITE_TO_TUNNEL_IPV6)
            && r.size > 0) {
            BYTE *outPkt = p->wintun.AllocateSendPacket(p->session,
                                                       static_cast<DWORD>(r.size));
            if (outPkt) {
                memcpy(outPkt, plain.data(), r.size);
                p->wintun.SendPacket(p->session, outPkt);
            }
            // First decrypted packet => handshake completed => tunnel is up.
            if (!p->gotFirstPkt.exchange(true) && p->owner) {
                const QString localIP = p->cfg.localIP();
                QPointer<TunnelManager> owner = p->owner;
                QMetaObject::invokeMethod(owner.data(), [owner, localIP]() {
                    if (owner) emit owner->tunnelUp(localIP);
                }, Qt::QueuedConnection);
            }
        } else if (r.op == WRITE_TO_NETWORK && r.size > 0) {
            sendto(p->udpSock,
                   reinterpret_cast<const char *>(plain.data()),
                   static_cast<int>(r.size),
                   0, reinterpret_cast<sockaddr *>(&peer), sizeof(peer));
        }
    }
    WSACloseEvent(sockEvent);
}

// ----- Ticker thread (keepalive / re-handshake) -----------------------------

static void tickerThread(WinTunnelPriv *p)
{
    sockaddr_in peer {};
    peer.sin_family = AF_INET;
    InetPtonA(AF_INET, p->cfg.endpointHost().toUtf8().constData(), &peer.sin_addr);
    peer.sin_port = htons(static_cast<u_short>(p->cfg.endpointPort()));

    static constexpr size_t BUF = 512;
    std::vector<uint8_t> dst(BUF);

    while (p->running.load()) {
        if (WaitForSingleObject(p->stopEvent, 100) != WAIT_TIMEOUT)
            break;

        size_t dstLen = dst.size();
        wireguard_result r = wireguard_tick(p->wg, dst.data(), &dstLen);
        if (r.op == WRITE_TO_NETWORK && r.size > 0) {
            sendto(p->udpSock,
                   reinterpret_cast<const char *>(dst.data()),
                   static_cast<int>(r.size),
                   0, reinterpret_cast<sockaddr *>(&peer), sizeof(peer));
        }
    }
}

// ----- Route setup helpers --------------------------------------------------

static void setAdapterAddress(const NET_LUID &luid, const QString &cidr)
{
    IN_ADDR addr {};
    UINT8   prefix = 32;
    if (!parseCIDR(cidr.toStdString(), &addr, &prefix))
        return;

    MIB_UNICASTIPADDRESS_ROW row {};
    InitializeUnicastIpAddressEntry(&row);
    row.InterfaceLuid              = luid;
    row.Address.si_family          = AF_INET;
    row.Address.Ipv4.sin_family    = AF_INET;
    row.Address.Ipv4.sin_addr      = addr;
    row.OnLinkPrefixLength         = prefix;
    row.DadState                   = IpDadStatePreferred;

    CreateUnicastIpAddressEntry(&row); // ignore error: may already exist
}

static void addRoute(const NET_LUID &luid, const QString &cidr)
{
    IN_ADDR net {};
    UINT8   prefix = 0;
    if (!parseCIDR(cidr.toStdString(), &net, &prefix))
        return;

    MIB_IPFORWARD_ROW2 row {};
    InitializeIpForwardEntry(&row);
    row.InterfaceLuid                              = luid;
    row.DestinationPrefix.Prefix.si_family         = AF_INET;
    row.DestinationPrefix.Prefix.Ipv4.sin_family   = AF_INET;
    row.DestinationPrefix.Prefix.Ipv4.sin_addr     = net;
    row.DestinationPrefix.PrefixLength             = prefix;
    row.NextHop.si_family                          = AF_INET;
    row.Metric                                     = 1;
    row.Protocol                                   = MIB_IPPROTO_NETMGMT;

    CreateIpForwardEntry2(&row); // ignore duplicate error
}

// ----- TunnelManager implementation -----------------------------------------

TunnelManager::TunnelManager(QObject *parent) : QObject(parent) {}

TunnelManager::~TunnelManager() { stop(); }

bool TunnelManager::start(const WireGuardConfig &cfg)
{
    if (!cfg.isValid()) {
        m_error = QStringLiteral("Invalid WireGuard configuration");
        emit tunnelError(m_error);
        return false;
    }

    if (m_running) stop();

    return platformStart(cfg);
}

void TunnelManager::stop()
{
    if (!m_running) return;
    m_running = false;
    platformStop();
    m_localAddr.clear();
    emit tunnelDown();
}

bool TunnelManager::platformStart(const WireGuardConfig &cfg)
{
    auto *p = new WinTunnelPriv();
    p->cfg   = cfg;
    p->owner = this;

    // 1. Load wintun.dll
    if (!p->wintun.load()) {
        delete p;
        m_error = QStringLiteral(
            "Failed to load wintun.dll - is it next to the executable?");
        emit tunnelError(m_error);
        return false;
    }

    // 2. Create WireGuard tunnel (boringtun)
    {
        const QByteArray priv = cfg.privateKey.toUtf8();
        const QByteArray pubk = cfg.peerPublicKey.toUtf8();
        const QByteArray psk  = cfg.presharedKey.toUtf8();
        p->wg = new_tunnel(
            priv.constData(),
            pubk.constData(),
            psk.isEmpty() ? nullptr : psk.constData(),
            cfg.persistentKeepalive,
            /*log_level=*/1);
    }
    if (!p->wg) {
        p->wintun.unload();
        delete p;
        m_error = QStringLiteral("boringtun: failed to create WireGuard tunnel");
        emit tunnelError(m_error);
        return false;
    }

    // 3. Create Wintun adapter
    GUID guid;
    CoCreateGuid(&guid);
    p->adapter = p->wintun.CreateAdapter(L"DragonMoonlight", L"WireGuard", &guid);
    if (!p->adapter) {
        const QString err = QString::fromStdString(lastWinError("WintunCreateAdapter"));
        tunnel_free(p->wg);
        p->wintun.unload();
        delete p;
        m_error = err;
        emit tunnelError(m_error);
        return false;
    }
    p->wintun.GetAdapterLUID(p->adapter, &p->adapterLuid);

    // 4. Configure IP address + routes
    setAdapterAddress(p->adapterLuid, cfg.address);
    for (const QString &cidr : cfg.allowedIPs)
        addRoute(p->adapterLuid, cidr);

    // 5. Start Wintun session
    p->session = p->wintun.StartSession(p->adapter, 0x400000); // 4 MiB ring
    if (!p->session) {
        const QString err = QString::fromStdString(lastWinError("WintunStartSession"));
        p->wintun.CloseAdapter(p->adapter);
        tunnel_free(p->wg);
        p->wintun.unload();
        delete p;
        m_error = err;
        emit tunnelError(m_error);
        return false;
    }

    // 6. Create UDP socket
    WSADATA wsd {};
    WSAStartup(MAKEWORD(2, 2), &wsd);
    p->udpSock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (p->udpSock == INVALID_SOCKET) {
        p->wintun.EndSession(p->session);
        p->wintun.CloseAdapter(p->adapter);
        tunnel_free(p->wg);
        WSACleanup();
        p->wintun.unload();
        delete p;
        m_error = QStringLiteral("Failed to create UDP socket");
        emit tunnelError(m_error);
        return false;
    }
    sockaddr_in local {};
    local.sin_family      = AF_INET;
    local.sin_addr.s_addr = INADDR_ANY;
    local.sin_port        = 0;
    bind(p->udpSock, reinterpret_cast<sockaddr *>(&local), sizeof(local));

    // 7. Stop event for I/O threads
    p->stopEvent = CreateEvent(nullptr, TRUE, FALSE, nullptr);
    if (!p->stopEvent) {
        const QString err = QString::fromStdString(lastWinError("CreateEvent"));
        closesocket(p->udpSock);
        WSACleanup();
        p->wintun.EndSession(p->session);
        p->wintun.CloseAdapter(p->adapter);
        tunnel_free(p->wg);
        p->wintun.unload();
        delete p;
        m_error = err;
        emit tunnelError(m_error);
        return false;
    }

    // 8. Force initial handshake
    {
        static constexpr size_t BUF = 512;
        std::vector<uint8_t> hs(BUF);
        size_t hsLen = hs.size();
        wireguard_force_handshake(p->wg, hs.data(), &hsLen);
        if (hsLen > 0) {
            sockaddr_in peer {};
            peer.sin_family = AF_INET;
            InetPtonA(AF_INET, cfg.endpointHost().toUtf8().constData(),
                      &peer.sin_addr);
            peer.sin_port = htons(static_cast<u_short>(cfg.endpointPort()));
            sendto(p->udpSock,
                   reinterpret_cast<const char *>(hs.data()),
                   static_cast<int>(hsLen),
                   0, reinterpret_cast<sockaddr *>(&peer), sizeof(peer));
        }
    }

    // 9. Start I/O threads
    p->running     = true;
    m_running      = true;
    m_localAddr    = cfg.localIP();
    p->thdTunToUdp = std::thread(tunToUdpThread, p);
    p->thdUdpToTun = std::thread(udpToTunThread, p);
    p->thdTicker   = std::thread(tickerThread,   p);

    m_priv = p;

    qDebug() << "[DragonVPN] Wintun tunnel started, local IP" << m_localAddr;

    // Belt-and-braces: if no packet arrives, fire tunnelUp() after 1500 ms so
    // the UI doesn't hang on a green-but-disconnected state.
    QPointer<TunnelManager> owner(this);
    QTimer::singleShot(1500, this, [owner]() {
        if (!owner || !owner->m_running) return;
        auto *priv = static_cast<WinTunnelPriv *>(owner->m_priv);
        if (priv && !priv->gotFirstPkt.load())
            emit owner->tunnelUp(owner->localAddress());
    });

    return true;
}

void TunnelManager::platformStop()
{
    auto *p = static_cast<WinTunnelPriv *>(m_priv);
    if (!p) return;

    p->running = false;
    if (p->stopEvent) SetEvent(p->stopEvent);

    if (p->thdTunToUdp.joinable()) p->thdTunToUdp.join();
    if (p->thdUdpToTun.joinable()) p->thdUdpToTun.join();
    if (p->thdTicker.joinable())   p->thdTicker.join();

    if (p->session) p->wintun.EndSession(p->session);
    if (p->adapter) p->wintun.CloseAdapter(p->adapter);

    if (p->wg) tunnel_free(p->wg);

    // socket close BEFORE WSACleanup
    if (p->udpSock != INVALID_SOCKET) closesocket(p->udpSock);
    WSACleanup();
    if (p->stopEvent) CloseHandle(p->stopEvent);

    p->wintun.unload();

    delete p;
    m_priv = nullptr;
    qDebug() << "[DragonVPN] Wintun tunnel stopped";
}

#endif // Q_OS_WIN