#include "server_pipe.h"
#include "Driver.h"
#include "check_vmware.h"

struct ClientInfo
{
    DWORD id = 0;
    DWORD pid = 0;
    LONG64 last_tick = 0;
    LONG active = 1;
    HANDLE pipe = nullptr;
    std::wstring label;
};

static HANDLE g_exit_event = nullptr;
static LONG g_shutdown_requested = 0;
static LONG g_next_client_id = 0;
static LONG g_seen_client = 0;
static LONG64 g_session_id = 0;
static std::mutex g_clients_mutex;
static std::vector<std::shared_ptr<ClientInfo>> g_clients;
static constexpr DWORD kFirstClientTimeoutMs = 3000;


static void ExitServerProcess(UINT exit_code)
{
    //printf("[server] ExitProcess code=%u\n", exit_code);
    fflush(stdout);
    ExitProcess(exit_code);
}

static bool ReadExact(HANDLE pipe, void* buffer, DWORD size)
{
    BYTE* ptr = static_cast<BYTE*>(buffer);
    DWORD total = 0;
    while (total < size)
    {
        DWORD got = 0;
        if (!ReadFile(pipe, ptr + total, size - total, &got, nullptr))
            return false;
        if (got == 0)
            return false;
        total += got;
    }
    return true;
}

static bool WriteExact(HANDLE pipe, const void* buffer, DWORD size)
{
    const BYTE* ptr = static_cast<const BYTE*>(buffer);
    DWORD total = 0;
    while (total < size)
    {
        DWORD wrote = 0;
        if (!WriteFile(pipe, ptr + total, size - total, &wrote, nullptr))
            return false;
        if (wrote == 0)
            return false;
        total += wrote;
    }
    return true;
}

static std::vector<std::shared_ptr<ClientInfo>> SnapshotClients()
{
    std::lock_guard<std::mutex> lock(g_clients_mutex);
    return g_clients;
}

static void RemoveClient(DWORD id)
{
    std::lock_guard<std::mutex> lock(g_clients_mutex);
    for (auto& client : g_clients)
    {
        if (client && client->id == id)
            InterlockedExchange(&client->active, 0);
    }

    g_clients.erase(
        std::remove_if(g_clients.begin(), g_clients.end(),
            [id](const std::shared_ptr<ClientInfo>& c) { return c && c->id == id; }),
        g_clients.end());

    if (InterlockedCompareExchange(&g_seen_client, 0, 0) != 0 && g_clients.empty())
    {
        //printf("[server] all clients exited, server shutdown\n");
        fflush(stdout);
        SetEvent(g_exit_event);
    }
}

static std::shared_ptr<ClientInfo> RegisterClient(HANDLE pipe, const Msg& hello)
{
    auto client = std::make_shared<ClientInfo>();
    client->id = static_cast<DWORD>(InterlockedIncrement(&g_next_client_id));
    client->pid = hello.pid;
    client->last_tick = static_cast<LONG64>(hello.tick);
    client->pipe = pipe;
    InterlockedExchange(&g_seen_client, 1);

    std::lock_guard<std::mutex> lock(g_clients_mutex);
    g_clients.push_back(client);
    return client;
}

static void BroadcastExit(const char* reason)
{
    auto clients = SnapshotClients();
    Msg msg{};
    msg.type = static_cast<uint32_t>(MsgType::Exit);
    msg.session_id = static_cast<uint64_t>(g_session_id);

    //printf("[server] broadcast exit: %s\n", reason ? reason : "unknown");
    fflush(stdout);

    for (auto& client : clients)
    {
        if (client && client->pipe)
            WriteExact(client->pipe, &msg, sizeof(msg));
    }
}

void RequestShutdown(const char* reason)
{
    if (InterlockedCompareExchange(&g_shutdown_requested, 1, 0) == 0)
    {
        SetEvent(g_exit_event);
        BroadcastExit(reason);
    }
}

static bool ReadMessage(HANDLE pipe, Msg& msg)
{
    return ReadExact(pipe, &msg, sizeof(msg));
}

static bool WriteMessage(HANDLE pipe, const Msg& msg)
{
    return WriteExact(pipe, &msg, sizeof(msg));
}

static bool SendCommandReply(
    HANDLE pipe,
    const Msg& req,
    uint32_t status,
    const void* output,
    uint32_t output_size)
{
    Msg reply{};
    reply.type = static_cast<uint32_t>(MsgType::CommandReply);
    reply.client_id = req.client_id;
    reply.pid = GetCurrentProcessId();
    reply.version = kProtocolVersion;
    reply.command = req.command;
    reply.status = status;
    reply.input_size = 0;
    reply.output_size = 0;
    reply.tick = static_cast<uint64_t>(NowMs());
    reply.session_id = static_cast<uint64_t>(g_session_id);
    reply.request_id = req.request_id;

    if (output && output_size)
    {
        if (output_size > kPayloadBytes)
            return false;

        memcpy(reply.payload, output, output_size);
        reply.output_size = output_size;
    }

    return WriteMessage(pipe, reply);
}

static bool SendCommandStatus(HANDLE pipe, const Msg& req, uint32_t status)
{
    return SendCommandReply(pipe, req, status, nullptr, 0);
}

static uint32_t HandlePrintf(const std::shared_ptr<ClientInfo>& client, const Msg& req)
{
    if (req.input_size < sizeof(PrintfRequest))
        return ERROR_INVALID_PARAMETER;

    const auto* in = reinterpret_cast<const PrintfRequest*>(req.payload);
    //printf("[server][printf][client=%lu pid=%lu] %s\n", client->id, client->pid, in->text);
    fflush(stdout);
    return ERROR_SUCCESS;
}
    
static uint32_t HandleLoadDriver(Msg& req)
{
    if (req.input_size < sizeof(LoadDriverRequest))
        return ERROR_INVALID_PARAMETER;

    auto* in = reinterpret_cast<LoadDriverRequest*>(req.payload);

    DWORD status = ERROR_SUCCESS;

    auto result =
        Driver::fn_get_instance()->driver_install(&status);

    in->status = status;

    return result ? ERROR_SUCCESS : ERROR_INVALID_PARAMETER;
}

static uint32_t HandleUnloadDriver(const Msg& req)
{
    if (req.input_size < sizeof(UnloadDriverRequest))
        return ERROR_INVALID_PARAMETER;
    
    const auto* in = reinterpret_cast<const UnloadDriverRequest*>(req.payload);
    if (!in->service_name[0])
        return ERROR_INVALID_PARAMETER;
    
    SC_HANDLE scm = OpenSCManagerW(nullptr, nullptr, SC_MANAGER_CONNECT);
    if (!scm)
        return GetLastError();
    
    SC_HANDLE service = OpenServiceW(scm, in->service_name, SERVICE_STOP | DELETE | SERVICE_QUERY_STATUS);
    if (!service)
    {
        DWORD err = GetLastError();
        CloseServiceHandle(scm);
        return err;
    }
    
    SERVICE_STATUS status{};
    ControlService(service, SERVICE_CONTROL_STOP, &status);
    
    DWORD err = ERROR_SUCCESS;
    if (!DeleteService(service))
    {
        err = GetLastError();
        if (err == ERROR_SERVICE_MARKED_FOR_DELETE)
            err = ERROR_SUCCESS;
    }
    
    CloseServiceHandle(service);
    CloseServiceHandle(scm);
    return err;
}

static uint32_t HandleBreakpointRemoved(const std::shared_ptr<ClientInfo>& client, const Msg& req)
{
    if (req.input_size < sizeof(BreakpointRemovedRequest))
        return ERROR_INVALID_PARAMETER;

    const auto* in = reinterpret_cast<const BreakpointRemovedRequest*>(req.payload);
    //printf("[server][breakpoint-removed][client=%lu pid=%lu] tid=%u dr0=0x%llx dr7=0x%llx\n",
    //    client->id,
    //    client->pid,
    //    in->thread_id,
    //    static_cast<unsigned long long>(in->dr0),
    //    static_cast<unsigned long long>(in->dr7));
    fflush(stdout);
    return ERROR_SUCCESS;
}

uint32_t HandleCheckVirtualSystem()
{

    BOOL bIsSuc[9]{ 0 };

    bIsSuc[0] = check_vmware::cpuid();
    bIsSuc[1] = check_vmware::cpuid2();
    bIsSuc[2] = check_vmware::diskname();
    bIsSuc[3] = check_vmware::dxgiGpuName();
    bIsSuc[4] = check_vmware::file();
    bIsSuc[5] = check_vmware::In();
    bIsSuc[6] = check_vmware::mousename();
    bIsSuc[7] = check_vmware::registry();
    bIsSuc[8] = check_vmware::process();

    for (size_t i = 0; i < 9; i++)
    {
        if (bIsSuc[i] == TRUE)
            return ERROR_SUCCESS;
    }

    return ERROR_INVALID_FUNCTION;
}
uint32_t HandleCheckVmxSetting()
{
    return check_vmware::virutal_check_vmcall() ? ERROR_SUCCESS : ERROR_INVALID_FUNCTION;
}

static bool DispatchCommand(HANDLE pipe, const std::shared_ptr<ClientInfo>& client, Msg& req)
{
    if (!client)
        return SendCommandStatus(pipe, req, ERROR_INVALID_PARAMETER);

    if (req.input_size > kPayloadBytes || req.output_size > kPayloadBytes)
        return SendCommandStatus(pipe, req, ERROR_INVALID_PARAMETER);

    uint32_t status = ERROR_INVALID_FUNCTION;

    switch (static_cast<CommandId>(req.command))
    {
    case CommandId::Printf:
        status = HandlePrintf(client, req);
        break;

    case CommandId::LoadDriver:

        status = HandleLoadDriver(req);
        return SendCommandReply(pipe, req, status, req.payload, sizeof(LoadDriverRequest));

    case CommandId::UnloadDriver:
        //status = HandleUnloadDriver(req);
        break;

    case CommandId::QueryStatus:
        status = ERROR_SUCCESS;
        break;
    case CommandId::BreakpointRemoved:
        status = HandleBreakpointRemoved(client, req);
        break;
    case CommandId::CheckVMware:
        status = HandleCheckVirtualSystem();
        break;
    case CommandId::CheckVmx:
        status = HandleCheckVmxSetting();
        break;
    default:
       /* printf("[server] unknown command=%u from client=%lu pid=%lu\n",
            req.command, client->id, client->pid);*/
        fflush(stdout);
        status = ERROR_INVALID_FUNCTION;
        break;
    }

    return SendCommandStatus(pipe, req, status);
}

static DWORD WINAPI FirstClientTimeoutProc(LPVOID)
{
    DWORD wait = WaitForSingleObject(g_exit_event, kFirstClientTimeoutMs);
    if (wait == WAIT_TIMEOUT &&
        InterlockedCompareExchange(&g_seen_client, 0, 0) == 0)
    {
        //printf("[server] no client connected in %lu ms, server shutdown\n", kFirstClientTimeoutMs);
        fflush(stdout);
        RequestShutdown("first client timeout");
    }

    return 0;
}

static DWORD WINAPI MonitorThreadProc(LPVOID)
{
    while (WaitForSingleObject(g_exit_event, kHeartbeatIntervalMs) == WAIT_TIMEOUT)
    {
        const LONG64 now = NowMs();
        auto clients = SnapshotClients();

        for (const auto& client : clients)
        {
            if (!client)
                continue;

            if (InterlockedCompareExchange(&client->active, 0, 0) == 0)
                continue;

            const LONG64 last = InterlockedCompareExchange64(&client->last_tick, 0, 0);
            if (last != 0 && now - last > kHeartbeatTimeoutMs)
            {
                //printf("[server] timeout pid=%lu id=%lu\n", client->pid, client->id);
                fflush(stdout);
                InterlockedExchange64(&client->last_tick, now);
            }
        }
    }
    return 0;
}

static DWORD WINAPI ClientSessionProc(LPVOID param)
{
    HANDLE pipe = reinterpret_cast<HANDLE>(param);

    Msg hello{};
    if (!ReadMessage(pipe, hello) || hello.type != static_cast<uint32_t>(MsgType::Hello))
    {
        CloseHandle(pipe);
        return 0;
    }

    auto client = RegisterClient(pipe, hello);
    /*printf("[server] client connected id=%lu pid=%lu session=%llu\n",
        client->id, client->pid, static_cast<unsigned long long>(hello.session_id));*/
    fflush(stdout);

    Msg ack{};
    ack.type = static_cast<uint32_t>(MsgType::HelloAck);
    ack.client_id = client->id;
    ack.pid = GetCurrentProcessId();
    ack.version = kProtocolVersion;
    ack.tick = static_cast<uint64_t>(NowMs());
    ack.session_id = static_cast<uint64_t>(g_session_id);
    if (!WriteMessage(pipe, ack))
    {
        RemoveClient(client->id);
        CloseHandle(pipe);
        return 0;
    }

    for (;;)
    {
        if (WaitForSingleObject(g_exit_event, 0) == WAIT_OBJECT_0)
            break;

        Msg msg{};
        if (!ReadMessage(pipe, msg))
        {
            //printf("[server] client disconnected id=%lu pid=%lu\n", client->id, client->pid);
            fflush(stdout);
            RemoveClient(client->id);
            CloseHandle(pipe);
            break;
        }

        if (msg.session_id != static_cast<uint64_t>(g_session_id))
        {
            /*printf("[server] client session mismatch id=%lu pid=%lu msg_session=%llu server_session=%llu\n",
                client->id,
                client->pid,
                static_cast<unsigned long long>(msg.session_id),
                static_cast<unsigned long long>(g_session_id));*/
            fflush(stdout);
            RemoveClient(client->id);
            CloseHandle(pipe);
            break;
        }

        if (msg.type == static_cast<uint32_t>(MsgType::Heartbeat))
        {
            InterlockedExchange64(&client->last_tick, static_cast<LONG64>(msg.tick));
            client->pid = msg.pid;
        }
        else if (msg.type == static_cast<uint32_t>(MsgType::Goodbye))
        {
            //printf("[server] client goodbye id=%lu pid=%lu\n", client->id, client->pid);
            fflush(stdout);
            RemoveClient(client->id);
            CloseHandle(pipe);
            break;
        }
        else if (msg.type == static_cast<uint32_t>(MsgType::Exit))
        {
            //printf("[server] client exit id=%lu pid=%lu, remove only\n", client->id, client->pid);
            fflush(stdout);
            RemoveClient(client->id);
            CloseHandle(pipe);
            break;
        }
        else if (msg.type == static_cast<uint32_t>(MsgType::CommandRequest))
        {
            if (!DispatchCommand(pipe, client, msg))
            {
               /* printf("[server] command reply failed id=%lu pid=%lu err=%lu\n",
                    client->id, client->pid, GetLastError());*/
                fflush(stdout);
                RemoveClient(client->id);
                CloseHandle(pipe);
                break;
            }
        }
    }

    return 0;
}

static DWORD WINAPI AcceptThreadProc(LPVOID)
{
    for (;;)
    {
        if (WaitForSingleObject(g_exit_event, 0) == WAIT_OBJECT_0)
            return 0;

        HANDLE pipe = CreateNamedPipeW(
            kPipeName,
            PIPE_ACCESS_DUPLEX,
            PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT,
            PIPE_UNLIMITED_INSTANCES,
            sizeof(Msg),
            sizeof(Msg),
            0,
            nullptr);

        if (pipe == INVALID_HANDLE_VALUE)
        {
            //printf("[server] CreateNamedPipeW failed: %lu\n", GetLastError());
            Sleep(1000);
            continue;
        }

        BOOL connected = ConnectNamedPipe(pipe, nullptr);
        if (!connected)
        {
            DWORD err = GetLastError();
            if (err != ERROR_PIPE_CONNECTED)
            {
                CloseHandle(pipe);
                if (WaitForSingleObject(g_exit_event, 0) == WAIT_OBJECT_0)
                    return 0;
                continue;
            }
        }

        HANDLE session_thread = CreateThread(nullptr, 0, ClientSessionProc, pipe, 0, nullptr);
        if (!session_thread)
        {
            //printf("[server] ClientSession thread failed: %lu\n", GetLastError());
            CloseHandle(pipe);
            continue;
        }

        CloseHandle(session_thread);
    }
}

static BOOL WINAPI CtrlHandler(DWORD type)
{
    switch (type)
    {
    case CTRL_C_EVENT:
    case CTRL_BREAK_EVENT:
    case CTRL_CLOSE_EVENT:
    case CTRL_LOGOFF_EVENT:
    case CTRL_SHUTDOWN_EVENT:
        RequestShutdown("console ctrl");
        return TRUE;
    default:
        return FALSE;
    }
}

static bool InitializeServer()
{
    g_exit_event = CreateEventW(nullptr, TRUE, FALSE, nullptr);
    if (!g_exit_event)
    {
        //printf("[server] CreateEventW failed: %lu\n", GetLastError());
        return false;
    }

    g_session_id = (static_cast<LONG64>(GetCurrentProcessId()) << 32) ^ NowMs();

   /* printf("[server] started protocol=%u session=%llu pipe=%ws\n",
        kProtocolVersion,
        static_cast<unsigned long long>(g_session_id),
        kPipeName);*/
    fflush(stdout);

    return true;
}

auto RunServerCore() -> BOOL
{
    HANDLE first_client_timeout_thread = CreateThread(nullptr, 0, FirstClientTimeoutProc, nullptr, 0, nullptr);
    if (!first_client_timeout_thread)
    {
        //printf("[server] first client timeout thread failed: %lu\n", GetLastError());
        return FALSE;
    }

    HANDLE monitor_thread = CreateThread(nullptr, 0, MonitorThreadProc, nullptr, 0, nullptr);
    if (!monitor_thread)
    {
        //printf("[server] monitor thread failed: %lu\n", GetLastError());
        return 1;
    }

    HANDLE accept_thread = CreateThread(nullptr, 0, AcceptThreadProc, nullptr, 0, nullptr);
    if (!accept_thread)
    {
        //printf("[server] accept thread failed: %lu\n", GetLastError());
        return 1;
    }

    WaitForSingleObject(g_exit_event, INFINITE);

    BroadcastExit("server shutdown");
    Sleep(400);


    CloseHandle(accept_thread);
    CloseHandle(monitor_thread);
    CloseHandle(first_client_timeout_thread);
    CloseHandle(g_exit_event);
    ExitServerProcess(0);

    return TRUE;
}

auto StartServer() -> BOOL
{
    if (!InitializeServer())
        return FALSE;

    return RunServerCore();
}