#define _POSIX_C_SOURCE 200809L
#define _DEFAULT_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/stat.h>
#include <dirent.h>
#include <time.h>
#include "main.h"
#include "microtar/microtar.h"

// Suppress warnings from miniz header
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-function"
#include "miniz/miniz.h"
#pragma GCC diagnostic pop

////////////////////////////////////////////////////////////////////////////////
//      DIRECTORY ARCHIVING FUNCTIONS
////////////////////////////////////////////////////////////////////////////////

// Helper function to recursively add directory contents to TAR archive
static int add_directory_to_tar(mtar_t* tar, const char* base_path, const char* relative_path) {
    DIR* dir = opendir(base_path);
    if (!dir) {
        printf("Error: Cannot open directory '%s'\n", base_path);
        return 1;
    }

    struct dirent* entry;
    while ((entry = readdir(dir)) != NULL) {
        // Skip . and ..
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) {
            continue;
        }

        // Build full path
        char full_path[2048];
        snprintf(full_path, sizeof(full_path), "%s/%s", base_path, entry->d_name);

        // Build relative path for TAR
        char tar_path[2048];
        if (strlen(relative_path) > 0) {
            snprintf(tar_path, sizeof(tar_path), "%s/%s", relative_path, entry->d_name);
        } else {
            snprintf(tar_path, sizeof(tar_path), "%s", entry->d_name);
        }

        struct stat st;
        if (stat(full_path, &st) != 0) {
            printf("Warning: Cannot stat '%s', skipping\n", full_path);
            continue;
        }

        if (S_ISDIR(st.st_mode)) {
            // Recursively add subdirectory
            if (add_directory_to_tar(tar, full_path, tar_path) != 0) {
                closedir(dir);
                return 1;
            }
        } else if (S_ISREG(st.st_mode)) {
            // Add regular file
            FILE* fp = fopen(full_path, "rb");
            if (!fp) {
                printf("Warning: Cannot open '%s', skipping\n", full_path);
                continue;
            }

            // Get file size
            fseek(fp, 0, SEEK_END);
            size_t file_size = ftell(fp);
            fseek(fp, 0, SEEK_SET);

            // Read file data
            unsigned char* file_data = malloc(file_size);
            if (!file_data) {
                printf("Error: Memory allocation failed for '%s'\n", full_path);
                fclose(fp);
                closedir(dir);
                return 1;
            }

            size_t bytes_read = fread(file_data, 1, file_size, fp);
            fclose(fp);

            if (bytes_read != file_size) {
                printf("Warning: Could not read entire file '%s', skipping\n", full_path);
                free(file_data);
                continue;
            }

            // Write to TAR
            if (mtar_write_file_header(tar, tar_path, file_size) != MTAR_ESUCCESS) {
                printf("Error: Failed to write TAR header for '%s'\n", tar_path);
                free(file_data);
                closedir(dir);
                return 1;
            }

            if (mtar_write_data(tar, file_data, file_size) != MTAR_ESUCCESS) {
                printf("Error: Failed to write TAR data for '%s'\n", tar_path);
                free(file_data);
                closedir(dir);
                return 1;
            }

            free(file_data);
        }
    }

    closedir(dir);
    return 0;
}

// Create TAR archive from directory
int create_tar_archive(const char* dir_path, const char* tar_output_path) {
    mtar_t tar;
    
    if (mtar_open(&tar, tar_output_path, "w") != MTAR_ESUCCESS) {
        printf("Error: Cannot create TAR file '%s'\n", tar_output_path);
        return 1;
    }

    // Get directory name for relative paths
    char dir_name[512];
    const char* last_slash = strrchr(dir_path, '/');
    if (last_slash) {
        strncpy(dir_name, last_slash + 1, sizeof(dir_name) - 1);
    } else {
        strncpy(dir_name, dir_path, sizeof(dir_name) - 1);
    }
    dir_name[sizeof(dir_name) - 1] = '\0';

    // Add directory contents to TAR
    int result = add_directory_to_tar(&tar, dir_path, dir_name);

    // Finalize and close TAR
    mtar_finalize(&tar);
    mtar_close(&tar);

    return result;
}

// Extract TAR archive to directory
int extract_tar_archive(const char* tar_path, const char* output_dir) {
    mtar_t tar;
    mtar_header_t header;

    if (mtar_open(&tar, tar_path, "r") != MTAR_ESUCCESS) {
        printf("Error: Cannot open TAR file '%s'\n", tar_path);
        return 1;
    }

    // Create output directory if it doesn't exist
    mkdir(output_dir, 0755);

    // Extract each file
    while (mtar_read_header(&tar, &header) == MTAR_ESUCCESS) {
        char output_path[2048];
        snprintf(output_path, sizeof(output_path), "%s/%s", output_dir, header.name);

        // Create parent directories
        char* last_slash = strrchr(output_path, '/');
        if (last_slash) {
            char parent_dir[2048];
            strncpy(parent_dir, output_path, last_slash - output_path);
            parent_dir[last_slash - output_path] = '\0';
            
            // Create directories recursively
            char* p = parent_dir;
            while (*p) {
                if (*p == '/') {
                    *p = '\0';
                    mkdir(parent_dir, 0755);
                    *p = '/';
                }
                p++;
            }
            mkdir(parent_dir, 0755);
        }

        // Extract file data
        unsigned char* data = malloc(header.size);
        if (!data) {
            printf("Error: Memory allocation failed\n");
            mtar_close(&tar);
            return 1;
        }

        if (mtar_read_data(&tar, data, header.size) != MTAR_ESUCCESS) {
            printf("Error: Failed to read data for '%s'\n", header.name);
            free(data);
            mtar_close(&tar);
            return 1;
        }

        // Write to file
        FILE* fp = fopen(output_path, "wb");
        if (!fp) {
            printf("Error: Cannot create file '%s'\n", output_path);
            free(data);
            mtar_close(&tar);
            return 1;
        }

        fwrite(data, 1, header.size, fp);
        fclose(fp);
        free(data);

        mtar_next(&tar);
    }

    mtar_close(&tar);
    return 0;
}

////////////////////////////////////////////////////////////////////////////////
//      COMPRESSION FUNCTIONS
////////////////////////////////////////////////////////////////////////////////

// Compress file with gzip (miniz)
int compress_file_gzip(const char* input_path, const char* output_path) {
    // Read input file
    FILE* in = fopen(input_path, "rb");
    if (!in) {
        printf("Error: Cannot open input file '%s'\n", input_path);
        return 1;
    }

    fseek(in, 0, SEEK_END);
    size_t input_size = ftell(in);
    fseek(in, 0, SEEK_SET);

    unsigned char* input_data = malloc(input_size);
    if (!input_data) {
        printf("Error: Memory allocation failed\n");
        fclose(in);
        return 1;
    }

    size_t bytes_read = fread(input_data, 1, input_size, in);
    fclose(in);

    if (bytes_read != input_size) {
        printf("Error: Failed to read input file\n");
        free(input_data);
        return 1;
    }

    // Compress with miniz
    mz_ulong compressed_size = compressBound(input_size);
    unsigned char* compressed_data = malloc(compressed_size);
    if (!compressed_data) {
        printf("Error: Memory allocation failed\n");
        free(input_data);
        return 1;
    }

    int result = compress2(compressed_data, &compressed_size,
                          input_data, input_size,
                          MZ_BEST_COMPRESSION);

    free(input_data);

    if (result != MZ_OK) {
        printf("Error: Compression failed (error code: %d)\n", result);
        free(compressed_data);
        return 1;
    }

    // Write compressed data
    FILE* out = fopen(output_path, "wb");
    if (!out) {
        printf("Error: Cannot create output file '%s'\n", output_path);
        free(compressed_data);
        return 1;
    }

    fwrite(compressed_data, 1, compressed_size, out);
    fclose(out);
    free(compressed_data);

    return 0;
}

// Decompress gzip file (miniz)
int decompress_file_gzip(const char* input_path, const char* output_path) {
    // Read compressed file
    FILE* in = fopen(input_path, "rb");
    if (!in) {
        printf("Error: Cannot open compressed file '%s'\n", input_path);
        return 1;
    }

    fseek(in, 0, SEEK_END);
    size_t compressed_size = ftell(in);
    fseek(in, 0, SEEK_SET);

    unsigned char* compressed_data = malloc(compressed_size);
    if (!compressed_data) {
        printf("Error: Memory allocation failed\n");
        fclose(in);
        return 1;
    }

    size_t bytes_read = fread(compressed_data, 1, compressed_size, in);
    fclose(in);

    if (bytes_read != compressed_size) {
        printf("Error: Failed to read compressed file\n");
        free(compressed_data);
        return 1;
    }

    // Estimate decompressed size (try multiple times if needed)
    mz_ulong output_size = compressed_size * 10;
    unsigned char* output_data = NULL;
    int result;

    for (int attempt = 0; attempt < 3; attempt++) {
        output_data = realloc(output_data, output_size);
        if (!output_data) {
            printf("Error: Memory allocation failed\n");
            free(compressed_data);
            return 1;
        }

        mz_ulong temp_size = output_size;
        result = uncompress(output_data, &temp_size, compressed_data, compressed_size);

        if (result == MZ_OK) {
            output_size = temp_size;
            break;
        } else if (result == MZ_BUF_ERROR) {
            // Buffer too small, try larger
            output_size *= 2;
        } else {
            printf("Error: Decompression failed (error code: %d)\n", result);
            free(compressed_data);
            free(output_data);
            return 1;
        }
    }

    free(compressed_data);

    if (result != MZ_OK) {
        printf("Error: Decompression failed after multiple attempts\n");
        free(output_data);
        return 1;
    }

    // Write decompressed data
    FILE* out = fopen(output_path, "wb");
    if (!out) {
        printf("Error: Cannot create output file '%s'\n", output_path);
        free(output_data);
        return 1;
    }

    fwrite(output_data, 1, output_size, out);
    fclose(out);
    free(output_data);

    return 0;
}

////////////////////////////////////////////////////////////////////////////////
//      HIGH-LEVEL DIRECTORY ENCRYPTION/DECRYPTION
////////////////////////////////////////////////////////////////////////////////

// Encrypt directory: TAR → GZIP → Encrypt
int encrypt_directory(const char* dir_path, const char* pad_identifier, const char* output_file) {
    char temp_tar[512];
    char temp_gz[512];
    int result = 0;

    // Generate temporary file paths
    snprintf(temp_tar, sizeof(temp_tar), "/tmp/otp_tar_%d.tar", getpid());
    snprintf(temp_gz, sizeof(temp_gz), "/tmp/otp_gz_%d.tar.gz", getpid());

    printf("Creating TAR archive...\n");
    if (create_tar_archive(dir_path, temp_tar) != 0) {
        printf("Error: Failed to create TAR archive\n");
        return 1;
    }

    printf("Compressing archive...\n");
    if (compress_file_gzip(temp_tar, temp_gz) != 0) {
        printf("Error: Failed to compress archive\n");
        unlink(temp_tar);
        return 1;
    }

    printf("Encrypting compressed archive...\n");
    result = encrypt_file(pad_identifier, temp_gz, output_file, 0);

    // Cleanup temporary files
    unlink(temp_tar);
    unlink(temp_gz);

    if (result == 0) {
        printf("Directory encrypted successfully: %s\n", output_file);
    }

    return result;
}

// Detect if file is a compressed TAR archive
int is_compressed_tar_archive(const char* file_path) {
    FILE* fp = fopen(file_path, "rb");
    if (!fp) {
        return 0;
    }

    unsigned char magic[512];
    size_t bytes_read = fread(magic, 1, sizeof(magic), fp);
    fclose(fp);

    if (bytes_read < 2) {
        return 0;
    }

    // Check for GZIP magic bytes (0x1f 0x8b)
    if (magic[0] == 0x1f && magic[1] == 0x8b) {
        return 1;
    }

    // Check for TAR magic ("ustar" at offset 257)
    if (bytes_read >= 262 && memcmp(magic + 257, "ustar", 5) == 0) {
        return 1;
    }

    return 0;
}

// Decrypt and extract directory: Decrypt → GUNZIP → Extract TAR
int decrypt_and_extract_directory(const char* encrypted_file, const char* output_dir) {
    char temp_decrypted[512];
    char temp_tar[512];
    int result = 0;

    // Generate temporary file paths
    snprintf(temp_decrypted, sizeof(temp_decrypted), "/tmp/otp_decrypt_%d", getpid());
    snprintf(temp_tar, sizeof(temp_tar), "/tmp/otp_tar_%d.tar", getpid());

    printf("Decrypting file...\n");
    if (decrypt_file(encrypted_file, temp_decrypted) != 0) {
        printf("Error: Failed to decrypt file\n");
        return 1;
    }

    // Check if it's compressed
    FILE* fp = fopen(temp_decrypted, "rb");
    if (!fp) {
        printf("Error: Cannot open decrypted file\n");
        unlink(temp_decrypted);
        return 1;
    }

    unsigned char magic[2];
    fread(magic, 1, 2, fp);
    fclose(fp);

    if (magic[0] == 0x1f && magic[1] == 0x8b) {
        // GZIP compressed
        printf("Decompressing archive...\n");
        if (decompress_file_gzip(temp_decrypted, temp_tar) != 0) {
            printf("Error: Failed to decompress archive\n");
            unlink(temp_decrypted);
            return 1;
        }
        unlink(temp_decrypted);
    } else {
        // Not compressed, assume it's already TAR
        rename(temp_decrypted, temp_tar);
    }

    printf("Extracting archive...\n");
    result = extract_tar_archive(temp_tar, output_dir);

    // Cleanup
    unlink(temp_tar);

    if (result == 0) {
        printf("Directory extracted successfully to: %s\n", output_dir);
    }

    return result;
}