#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")

#include "Net.h"
#include <algorithm>
#include <chrono>
#include <cstring>
#include <iostream>

#ifndef SIO_UDP_CONNRESET
#define SIO_UDP_CONNRESET _WSAIOW(IOC_VENDOR, 12)
#endif

namespace {
    constexpr uint32_t       NET_MAGIC = 0x45494946;  // "EIIF"
    constexpr uint16_t       NET_VERSION = 2;
    constexpr unsigned short NET_PORT = 27015;
    constexpr double         NET_TIMEOUT = 2.0;       // seconds of silence before a player is dropped

    // One datagram = Header followed by 'count' Entry records.
    // Client -> host: count 1 (the client's own state, id ignored).
    // Host -> client: one entry per other player, host included (id 0).
#pragma pack(push, 1)
    struct Header {
        uint32_t magic;
        uint16_t version;
        uint32_t seq;
        uint8_t  count;
    };
    struct Entry {
        uint8_t id;
        float   pos[3];
        float   rot[4];  // x, y, z, w
        uint8_t sliding;
    };
#pragma pack(pop)

    constexpr int MAX_ENTRIES = Lan::MaxPlayers;
    constexpr int MAX_DATAGRAM = (int)sizeof(Header) + MAX_ENTRIES * (int)sizeof(Entry);

    double Now() {
        using namespace std::chrono;
        return duration<double>(steady_clock::now().time_since_epoch()).count();
    }

    Entry MakeEntry(uint8_t id, const NetState& s) {
        Entry e;
        e.id = id;
        e.pos[0] = s.position.x; e.pos[1] = s.position.y; e.pos[2] = s.position.z;
        e.rot[0] = s.rotation.x; e.rot[1] = s.rotation.y; e.rot[2] = s.rotation.z; e.rot[3] = s.rotation.w;
        e.sliding = s.sliding ? 1 : 0;
        return e;
    }

    NetState ReadEntry(const Entry& e) {
        NetState s;
        s.position = glm::vec3(e.pos[0], e.pos[1], e.pos[2]);
        s.rotation = glm::quat(e.rot[3], e.rot[0], e.rot[1], e.rot[2]);  // glm ctor order is w, x, y, z
        s.sliding = e.sliding != 0;
        return s;
    }

    bool SameAddress(const sockaddr_in& a, const sockaddr_in& b) {
        return a.sin_addr.s_addr == b.sin_addr.s_addr && a.sin_port == b.sin_port;
    }
}

struct Lan::Impl {
    struct Peer {
        sockaddr_in addr{};
        uint8_t     id = 0;
        double      lastRecv = 0.0;
        uint32_t    lastSeq = 0;
        bool        started = false;  // has sent at least one datagram (so lastSeq is valid)
        NetState    state;
        bool        hasState = false;
    };

    SOCKET      sock = INVALID_SOCKET;
    bool        wsaStarted = false;
    uint32_t    sendSeq = 0;

    // Host
    std::vector<Peer> peers;

    // Client
    sockaddr_in hostAddr{};
    bool        anyRecv = false;
    uint32_t    lastSeq = 0;
    double      lastRecvTime = 0.0;

    void SendDatagram(const sockaddr_in& to, const std::vector<Entry>& entries) {
        char buf[MAX_DATAGRAM];
        Header h;
        h.magic = NET_MAGIC;
        h.version = NET_VERSION;
        h.seq = ++sendSeq;
        h.count = (uint8_t)entries.size();
        std::memcpy(buf, &h, sizeof(h));
        std::memcpy(buf + sizeof(h), entries.data(), entries.size() * sizeof(Entry));
        sendto(sock, buf, (int)(sizeof(h) + entries.size() * sizeof(Entry)), 0,
               (const sockaddr*)&to, (int)sizeof(to));
    }
};

Lan::Lan() : impl(new Impl()) {}

Lan::~Lan() {
    Stop();
    delete impl;
}

bool Lan::Start(Mode m, const std::string& hostIp) {
    if (m == Mode::Off) return false;

    WSADATA wsa;
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0) {
        std::cout << "[net] WSAStartup failed\n";
        return false;
    }
    impl->wsaStarted = true;

    impl->sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (impl->sock == INVALID_SOCKET) {
        std::cout << "[net] socket() failed: " << WSAGetLastError() << "\n";
        Stop();
        return false;
    }

    u_long nonBlocking = 1;
    ioctlsocket(impl->sock, FIONBIO, &nonBlocking);

    // Don't let an ICMP "port unreachable" (other side not up yet) kill recvfrom with WSAECONNRESET
    BOOL newBehavior = FALSE;
    DWORD bytes = 0;
    WSAIoctl(impl->sock, SIO_UDP_CONNRESET, &newBehavior, sizeof(newBehavior), NULL, 0, &bytes, NULL, NULL);

    sockaddr_in local{};
    local.sin_family = AF_INET;
    local.sin_addr.s_addr = htonl(INADDR_ANY);
    local.sin_port = htons(m == Mode::Host ? NET_PORT : 0);  // clients use any free port
    if (bind(impl->sock, (sockaddr*)&local, sizeof(local)) == SOCKET_ERROR) {
        std::cout << "[net] bind() failed: " << WSAGetLastError()
                  << " (is another copy already hosting on this PC?)\n";
        Stop();
        return false;
    }

    if (m == Mode::Client) {
        impl->hostAddr.sin_family = AF_INET;
        impl->hostAddr.sin_port = htons(NET_PORT);
        if (inet_pton(AF_INET, hostIp.c_str(), &impl->hostAddr.sin_addr) != 1) {
            std::cout << "[net] '" << hostIp << "' is not a valid IPv4 address\n";
            Stop();
            return false;
        }
    }

    mode = m;
    active = true;
    return true;
}

void Lan::Stop() {
    if (impl->sock != INVALID_SOCKET) {
        closesocket(impl->sock);
        impl->sock = INVALID_SOCKET;
    }
    if (impl->wsaStarted) {
        WSACleanup();
        impl->wsaStarted = false;
    }
    impl->peers.clear();
    impl->anyRecv = false;
    active = false;
    mode = Mode::Off;
}

void Lan::Send(const NetState& s) {
    if (!active) return;

    if (mode == Mode::Client) {
        impl->SendDatagram(impl->hostAddr, { MakeEntry(0, s) });
        return;
    }

    // Host: to each player, send the host's state plus every other player's latest state
    for (const Impl::Peer& to : impl->peers) {
        std::vector<Entry> entries;
        entries.push_back(MakeEntry(0, s));
        for (const Impl::Peer& other : impl->peers) {
            if (other.id != to.id && other.hasState && (int)entries.size() < MAX_ENTRIES)
                entries.push_back(MakeEntry(other.id, other.state));
        }
        impl->SendDatagram(to.addr, entries);
    }
}

void Lan::Receive(std::vector<NetUpdate>& out) {
    if (!active) return;

    char buf[MAX_DATAGRAM + 64];
    for (int i = 0; i < 512; i++) {  // bounded so a flood can't stall the frame
        sockaddr_in from{};
        int fromLen = (int)sizeof(from);
        int r = recvfrom(impl->sock, buf, (int)sizeof(buf), 0, (sockaddr*)&from, &fromLen);
        if (r == SOCKET_ERROR) {
            if (WSAGetLastError() == WSAEWOULDBLOCK) break;
            continue;
        }
        if (r < (int)sizeof(Header)) continue;

        Header h;
        std::memcpy(&h, buf, sizeof(h));
        if (h.magic != NET_MAGIC || h.version != NET_VERSION) continue;
        if (h.count == 0 || h.count > MAX_ENTRIES) continue;
        if (r != (int)(sizeof(Header) + h.count * sizeof(Entry))) continue;

        const double now = Now();

        if (mode == Mode::Host) {
            // Find this sender, or admit them as a new player
            Impl::Peer* peer = nullptr;
            for (Impl::Peer& p : impl->peers)
                if (SameAddress(p.addr, from)) { peer = &p; break; }

            if (!peer) {
                if ((int)impl->peers.size() >= MaxPlayers - 1) continue;  // pit is full
                uint8_t id = 1;
                while (std::any_of(impl->peers.begin(), impl->peers.end(),
                                   [&](const Impl::Peer& p) { return p.id == id; }))
                    id++;
                Impl::Peer np;
                np.addr = from;
                np.id = id;
                impl->peers.push_back(np);
                peer = &impl->peers.back();
                std::cout << "[net] player " << (int)id << " joined\n";
            }

            if (peer->started && (int32_t)(h.seq - peer->lastSeq) <= 0) continue;  // stale / duplicate
            peer->started = true;
            peer->lastSeq = h.seq;
            peer->lastRecv = now;

            Entry e;
            std::memcpy(&e, buf + sizeof(Header), sizeof(Entry));
            peer->state = ReadEntry(e);
            peer->hasState = true;

            NetUpdate u;
            u.id = peer->id;
            u.state = peer->state;
            out.push_back(u);
        }
        else {
            if (!SameAddress(from, impl->hostAddr)) continue;  // only listen to the host

            // After a long silence the host may have restarted (sequence back to 1): accept anything
            if (impl->anyRecv && now - impl->lastRecvTime >= NET_TIMEOUT) impl->anyRecv = false;
            if (impl->anyRecv && (int32_t)(h.seq - impl->lastSeq) <= 0) continue;
            impl->anyRecv = true;
            impl->lastSeq = h.seq;
            impl->lastRecvTime = now;

            for (int k = 0; k < h.count; k++) {
                Entry e;
                std::memcpy(&e, buf + sizeof(Header) + k * sizeof(Entry), sizeof(Entry));
                NetUpdate u;
                u.id = e.id;
                u.state = ReadEntry(e);
                out.push_back(u);
            }
        }
    }

    // Host: forget players who stopped sending
    if (mode == Mode::Host) {
        const double now = Now();
        for (size_t i = 0; i < impl->peers.size();) {
            if (now - impl->peers[i].lastRecv > NET_TIMEOUT) {
                std::cout << "[net] player " << (int)impl->peers[i].id << " left\n";
                impl->peers.erase(impl->peers.begin() + i);
            }
            else i++;
        }
    }
}

bool Lan::IsConnected() const {
    if (!active) return false;
    if (mode == Mode::Host) return !impl->peers.empty();
    return impl->anyRecv && (Now() - impl->lastRecvTime) < NET_TIMEOUT;
}

int Lan::PeerCount() const {
    return (int)impl->peers.size();
}

std::string Lan::LocalAddresses() {
    WSADATA wsa;
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0) return "";

    std::string result;
    char name[256];
    if (gethostname(name, sizeof(name)) == 0) {
        addrinfo hints{};
        hints.ai_family = AF_INET;
        addrinfo* list = nullptr;
        if (getaddrinfo(name, nullptr, &hints, &list) == 0) {
            for (addrinfo* a = list; a; a = a->ai_next) {
                char ip[INET_ADDRSTRLEN];
                const sockaddr_in* sa = (const sockaddr_in*)a->ai_addr;
                if (inet_ntop(AF_INET, &sa->sin_addr, ip, sizeof(ip)))
                    result += std::string(result.empty() ? "" : ", ") + ip;
            }
            freeaddrinfo(list);
        }
    }
    WSACleanup();
    return result;
}