#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <stdint.h>
#include <zlib.h>

#define BUFSIZE 8192
#define SPLIT 128

// gcc -Wno-all -ggdb -O0 main.c -lz -o main && ./main hide mountains.png text.txt

// Structure of each chunks
// http://www.libpng.org/pub/png/spec/1.2/PNG-Structure.html#Chunk-layout

// http://www.libpng.org/pub/png/spec/1.2/PNG-Structure.html#PNG-file-signature
struct signature{
    unsigned char bytes[8];
} __attribute__((packed));
// __attribute__((packed)) -> for removing padding
// GCC add padding in the struct, 3 bytes
//https://stackoverflow.com/questions/1841863/size-of-struct-in-c

struct chunk{
    unsigned char length[4];
    char type[4];
} __attribute__((packed));

// http://www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html#C.IHDR
struct iHDR{
    unsigned char width[4];
    unsigned char height[4];
    char depth;
    char type;
    char compression;
    char filter;
    char interlace;
    // Chunk crc
    unsigned char crc[4];
} __attribute__((packed));

/* PROTOTYPES */
static void list_chunks(int, const char *, int);
static int hide_message(int, char **);
static int unhide_message(int, char **);
static int getHeader(int, char *);
static size_t insertChunk(int, const char *, size_t, int);
static void convert_from_uint(unsigned char *, const unsigned int);
static void convert_from_bytes(const unsigned char *, unsigned int *);
static void usage(const char *);

int main(int argc, char *argv[], char **envp){
    int fd_r, dump = -1;

    // The two statements below check if the arguments passed are correct
    if (argc < 3){
        usage(argv[0]);
        return -1;
    }
    if (strcmp(argv[1], "hide") != 0 && strcmp(argv[1], "unhide") != 0 && strcmp(argv[1], "chunks") != 0){
        usage(argv[0]);
        return -1;
    }

    if (strncmp(argv[1], "hide", 4) == 0 && argc < 4){
        usage(argv[0]);
        return -1;
    }
    else if (strncmp(argv[1], "unhide", 6) == 0 && argc < 4){
        usage(argv[0]);
        return -1;
    }
    else if (strncmp(argv[1], "chunks", 6) == 0 && argc == 4){
        dump = atoi(argv[3]);    
    }

    // https://man7.org/linux/man-pages/man2/open.2.html
    // The next two statements open the picture src and destination
    if ((fd_r = open(argv[2], O_RDONLY)) == -1){
        printf("Failed to open the png file %s\n", argv[1]);
        return -1;
    }

    if (strcmp(argv[1], "hide") == 0)
        hide_message(fd_r, argv);
    if (strcmp(argv[1], "unhide") == 0)
        unhide_message(fd_r, argv);
    if (strcmp(argv[1], "chunks") == 0)
        list_chunks(fd_r, argv[2], dump);
    close(fd_r);
}
/*
 * This function print the usage of the program
 */
static void usage(const char *program){
    printf("Usage: %s <hide|unhide|chunks> <picture> <text file|id>\n", program);
}

/*
 * In this function, we get all chunks and dump a chunk if specified
 */
static void list_chunks(int fd, const char *filename, int dump){
    struct chunk s_chunk;
    int c = 1;
    int nbChunks = 0;

    // We bypass the signature
    lseek(fd, 8, SEEK_SET);

    while(c){
        read(fd, &s_chunk, sizeof(struct chunk));
        off_t off = lseek(fd, 0, SEEK_CUR);
        unsigned int s = 0;
        convert_from_bytes(s_chunk.length, &s);
        if (dump == -1)
            printf("chunk %d at 0x%0x type %c%c%c%c, length %x\n", 
                nbChunks, off, s_chunk.type[0], s_chunk.type[1], s_chunk.type[2], s_chunk.type[3],
                s); 
       
        if (nbChunks == dump){
            char *tmp = malloc(s);
            read(fd, tmp, s);
            printf("%s\n", tmp);
            //write(0, tmp, s); // Write to stdout
            lseek(fd, off, SEEK_SET); // Go back to the chunk header
            free(tmp);
        }

        if (s_chunk.type[0] == 'I' && s_chunk.type[1] == 'E' &&
            s_chunk.type[2] == 'N' && s_chunk.type[3] == 'D')
            c = 0;

        /* +4 for the CRC */
        lseek(fd, s + 4, SEEK_CUR);

        nbChunks += 1;
        memset(&s_chunk, 0, sizeof(struct chunk));
    }
}
/*
 * In this function, we will hide the text provide in argument of the program to 
 * the destination image
 */
static int hide_message(int fd_r, char **argv){
    int fd_w;
    int fd_r_text;
    size_t length;
    size_t length_text_to_hide;
    char *buf;
    char text[BUFSIZE];
    char tmp[8];
    struct chunk *s_chunk = {0};
    char buffer[sizeof(struct signature) + sizeof(struct chunk) + sizeof(struct iHDR)];
    int end = 0;

    if ((fd_w = open("dest.png", O_WRONLY | O_CREAT, 00644)) == -1){
        printf("Failed to create the dest file\n");
        close(fd_r);
        return -1;
    }

    // The two next statements, we open the text to hide and buffer it
    if ((fd_r_text = open(argv[3], O_RDONLY)) == -1){
        printf("Failed to open the png file %s\n", argv[3]);
        close(fd_r);
        close(fd_w);
        return -1;
    }

    if ((length_text_to_hide = read(fd_r_text, text, BUFSIZE)) == -1){
        printf("Failed to read the file %s\n", argv[3]);
        close(fd_r);
        close(fd_w);
        return -1;
    }

    // Get the header and the IHDR and copy to the dest file
    length = getHeader(fd_r, buffer);

    if (length == 0){
        printf("Failed to get the png header\n");
        close(fd_w);
        return -1;
    }
    write(fd_w, buffer, length);

    // TODO
    // For the text, we split it into different chunk. We split the size defined by the global variable SPLIT
    // Each chunk are inserted between differents chunks
    // We insert the text
    /*if (length_text_to_hide > SPLIT){

    }*/
    length = insertChunk(fd_w, text, length_text_to_hide, 0);
    

    // We will parse each chunk until the IEND chunk
    do{
        unsigned int s = 0;
        // read the 8 bytes (for the chunk length and type)
        length = read(fd_r, tmp, 8);
        s_chunk = (struct chunk*)(tmp);

        // Write the chunk block (length and type) without the CRC
        write(fd_w, tmp, length);

        // Get the size of the data
        convert_from_bytes(s_chunk->length, &s);

        // We add 4 bytes for the CRC
        s += 4;

        // Identify the chunk
        if (s_chunk->type[0] == 'I' && s_chunk->type[1] == 'E' &&
            s_chunk->type[2] == 'N' && s_chunk->type[3] == 'D'){
            end = 1;
        }

        // The next statements allocate the buf with the chunk data size
        // and read the chunk data with the CRC
        buf = malloc(s);

        // read the file
        length = read(fd_r, buf, s);
        write(fd_w, buf, length);

        memset(tmp, 0, 8);
        free(buf);
    }while(end != 1);

    close(fd_w);
    close(fd_r_text);

    return 0;
}
static int unhide_message(int fd_r, char **argv){
    int fd_w;
    int length;
    char buffer[sizeof(struct signature) + sizeof(struct chunk) + sizeof(struct iHDR)];
    struct chunk *s_chunk;
    char tmp[8];
    char *buf;
    int end = 0;

    // The statement below check if we can create the file for the output
    if ((fd_w = open(argv[3], O_WRONLY | O_CREAT, 00644)) == -1){
        printf("Failed to write the output text file %s\n", argv[3]);
        close(fd_r);
        return -1;
    }

    // Get the header and the IHDR
    length = getHeader(fd_r, buffer);
    if (length == 0){
        printf("Failed to get the png header\n");
        close(fd_w);
        return -1;
    }

    // We will parse each chunk until the IEND chunk 
    do{
        unsigned int s = 0;
        unsigned char crc[4];
        // read the 8 bytes (for the chunk length and type)
        length = read(fd_r, tmp, 8);
        s_chunk = (struct chunk*)(tmp);

        // Identify the chunk
        if (s_chunk->type[0] == 'I' && s_chunk->type[1] == 'E' &&
            s_chunk->type[2] == 'N' && s_chunk->type[3] == 'D'){
            end = 1;
        }

        // Get the size of the data
        convert_from_bytes(s_chunk->length, &s);

        // The next statements allocate the buf with the chunk data size
        // and read the chunk data with the CRC
        buf = malloc(s);

        // read the file
        length = read(fd_r, buf, s);
        if (s_chunk->type[0] == 's' && s_chunk->type[1] == 'T' &&
            s_chunk->type[2] == 'E' && s_chunk->type[3] == 'G'){
            write(fd_w, buf, length);
        }

        // Read the crc
        read(fd_r, crc, 4);

        memset(tmp, 0, 8);
        free(buf);
    }while(end != 1);

    close(fd_w);
    return 0;
}
/*
 * In this function, we put our chunk to the file
 * this chunk contain our text message to hide
 */
static size_t insertChunk(int fd_w, const char *buf, size_t text_length, int start){
    int length;
    unsigned char crc[4];
    unsigned int crc_z;
    struct chunk s_chunk;
    memset(crc, 0, 4);

    // We create our chunk layout and write to the file
    convert_from_uint(s_chunk.length, text_length);

    s_chunk.type[0] = 's';
    s_chunk.type[1] = 'T';
    s_chunk.type[2] = 'E';
    s_chunk.type[3] = 'G';

    write(fd_w, s_chunk.length, 4);
    write(fd_w, s_chunk.type, 4);

    length = write(fd_w, buf, text_length);

    // Generate the CRC and write it to the file
    crc_z = crc32(0L, Z_NULL, 0);
    crc_z = crc32(crc_z, s_chunk.type, 4);
    crc_z = crc32(crc_z, buf, text_length);

    convert_from_uint(crc, crc_z);
    length = write(fd_w, crc, 4);

    return length;
}
/*
 * This function convert the unsigned int 32 bits to bytes and store to the buf variable
 */
static void convert_from_uint(unsigned char *buf, const unsigned int uint){
    buf[0] = (unsigned char)((uint >> 24) & 0xFFU);
    buf[1] = (unsigned char)((uint >> 16) & 0xFFU);
    buf[2] = (unsigned char)((uint >> 8) & 0xFFU);
    buf[3] = (unsigned char)(uint & 0xFFU);
}
/*
 * This function convert the bytes buf to unsigned int 32 bits
 */
static void convert_from_bytes(const unsigned char *buf, unsigned int *uint){
    *uint = buf[3] + (buf[2] << 8) + (buf[1] << 16) + (buf[0] << 24);
}
/*
 * This function read thefile for getting the signature, the 8 first bytes
 * and get the IHDR which contains some information regarding the picture
 */
static int getHeader(int fd_r, char *buf){
    size_t size_hdr = sizeof(struct signature) + sizeof(struct chunk) + sizeof(struct iHDR);
    int length;
    int i;
    unsigned int width;
    unsigned int height;
    struct signature *s_sign = {0};
    struct iHDR *s_ihdr = {0};

    // Get the PNG signature and the header
    length = read(fd_r, buf, size_hdr);

    if (length == 0)
        return 0;

    // Display signature data
    s_sign = (struct signature *)buf;
    printf("Signature\n");
    for (i = 0; i < sizeof(struct signature); i++)
        printf("%x ", s_sign->bytes[i]);
    printf("\n\n");

    // Display header data
    s_ihdr = (struct iHDR*)(buf + sizeof(struct signature) + sizeof(struct chunk));

    printf("Size of the image:\n");
    convert_from_bytes(s_ihdr->width, &width);
    convert_from_bytes(s_ihdr->height, &height);
    printf("Width: %upx\nHeight: %upx\n", width, height);

    return length;
}
