#pragma once
#include "Utils/ThreadPool.hpp"
#include "Socket/WSAManager.hpp"
#include "Socket/TCPSocket.hpp"
#include "Socket/Log.hpp"
#include <functional>

namespace Chattr {

struct IOCPPASSINDATA {
    OVERLAPPED overlapped;
    TCPSocket socket;
    char buf[1501];
    int recvbytes;
    int sendbytes;
    WSABUF wsabuf;
};

class IOCP {
public:
    static void iocpWather(ThreadPool* threadPool, HANDLE completionPort_, std::function<void(ThreadPool*, IOCPPASSINDATA*)> callback) {
        DWORD tid = GetCurrentThreadId();
        spdlog::debug("Waiting IO to complete on TID: {}.", tid);
        IOCPPASSINDATA* data;
        SOCKET sock;
        DWORD cbTransfrred;
        int retVal = GetQueuedCompletionStatus(completionPort_, &cbTransfrred, (PULONG_PTR)&sock, (LPOVERLAPPED*)&data, INFINITE);
        if (retVal == 0 || cbTransfrred == 0) {
            spdlog::info("Client disconnected. [{}]", (std::string)(data->socket.remoteAddr));
            delete data;
            threadPool->enqueueJob(iocpWather, completionPort_, callback);
            return;
        }
        threadPool->enqueueJob(callback, data);
        threadPool->enqueueJob(iocpWather, completionPort_, callback);
    };

    template<typename _Callable>
    void init(ThreadPool* __IOCPThread, _Callable&& callback) {
        IOCPThread_ = __IOCPThread;
#ifdef _WIN32
        completionPort_ = ::CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
        if (completionPort_ == NULL)
            log::critical("CreateIoCompletionPort()");

        auto boundFunc = [callback = std::move(callback)](ThreadPool* __IOCPThread, IOCPPASSINDATA* data) mutable {
            callback(__IOCPThread, data);
        };

        int tCount = __IOCPThread->threadCount;

        spdlog::info("Resizing threadpool size to: {}", tCount * 2);

        __IOCPThread->respawnWorker(tCount * 2);

        spdlog::info("Set IOCP Worker count to: {}", tCount);
        for (int i = 0; i < tCount; i++) {
            std::function<void(ThreadPool*, IOCPPASSINDATA*)> task(boundFunc);
            __IOCPThread->enqueueJob(iocpWather, completionPort_, task);
        }
#elif __linux__

#endif
    }

    void registerSocket(SOCKET sock);

    int recv(void* __restrict __buf, size_t __n, int __flags);
    int send(const void* __buf, size_t __n, int __flags);

private:
    struct Chattr::WSAManager wsaManager;
    ThreadPool* IOCPThread_;

#ifdef _WIN32
    HANDLE completionPort_ = INVALID_HANDLE_VALUE;
#elif __linux__

#endif
};

}