#include <stdio.h>
#include <stdlib.h>
#include <cs50.h>

int main(int argc, char* argv[])
{

    // check single argument passed, returning error if not
    if (argc != 2)
    {
        fprintf(stderr, "Usage : ./recover filename\n");
        return 1;
    }

    char* filename = argv[1];

    // open file for reading, returning error if failure
    FILE *inptr = fopen(filename, "r");

    if (inptr == NULL)
    {
        fprintf(stderr, "Could not open %s.\n", filename);
        return 2;
    }

    // iterate through bytes in file
    // count how many bytes found and print result
    int header = 0x00;
    int checkHeader = 0x00;
    int Jheader = 0x00;
    int JcheckHeader = 0x00;
    int counter = 0;
    int filesize = 0;
    int junk = 0x00;

    // calculate size of file to be recovered
    while (!feof(inptr))
    {
        fread(&junk, 1, 1, inptr);
        filesize++;
    }

    // reset file pointer to start of file
    rewind(inptr);

    // iterate through file looking for JPEG headers
    for (int j = 0; j < filesize; j++)
    {
        fread(&header, 3, 1, inptr);
        if (header == 0xffd8ff)
        {
            fread(&checkHeader, 1, 1, inptr);
            if (checkHeader >= 0xe0 && checkHeader <= 0xef)
            {
                // copy file to new jpg in blocks of 512 bytes
                char outfile[8];
                sprintf(outfile, "%03i.jpg", counter);
                counter++;
                FILE *outptr = fopen(outfile, "w");
                if (outptr == NULL)
                {
                    fprintf(stderr, "Could not create %s.\n", outfile);
                }

                // loop through blocks of 512 bytes until JPEG header found
                bool isEOJ = false;
                fseek(inptr, -4, SEEK_CUR);

                do
                {
                    // read 512 bytes into buffer
                    int buffer[512];
                    j += 512;
                    for (int i = 0; i < 512; i++)
                    {
                        fread(&buffer[i], 1, 1, inptr);
                    }

                    // write buffer to file
                    for (int i = 0; i < 512; i++)
                    {
                        fwrite(&buffer[i], 1, 1, outptr);
                    }

                    // check if start of next block is JPEG header

                    fread(&Jheader, 3, 1, inptr);
                    if (Jheader == 0xffd8ff)
                    {
                        fread(&JcheckHeader, 1, 1, inptr);
                        if (JcheckHeader >= 0xe0 && JcheckHeader <= 0xef)
                        {
                            fseek(inptr, -4, SEEK_CUR);
                            isEOJ = true;
                        } else
                        {
                            fseek(inptr, -4, SEEK_CUR);
                        }
                    } else
                    {
                        fseek (inptr, -3, SEEK_CUR);
                    }

                    if (j >= filesize)
                    {
                        isEOJ = true;
                    }

                    } while (!isEOJ);

                fclose(outptr);

            } else
            {
                fseek(inptr, -1, SEEK_CUR);
            }
        }
        fseek(inptr, -2, SEEK_CUR);
    }

    return 0;
}