#include <iostream>
#include <fstream>
#include <vector>
#include <cstdint>

const uint32_t PNG_SIGNATURE = 0x89504E47;
const uint32_t PNG_END_SIGNATURE = 0xAE426082;

bool readFileToBuffer(const std::string& filename, std::vector<uint8_t>& buffer) {
    std::ifstream file(filename, std::ios::binary);
    if (!file) {
        std::cerr << "Unable to open file: " << filename << std::endl;
        return false;
    }

    // Read the entire file into a buffer
    buffer.assign((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    return true;
}

void extractPNGs(const std::string& outdir, const std::vector<uint8_t>& buffer) {
    size_t offset = 0;
    size_t fileCount = 0;

    while (offset < buffer.size() - 8) { // Ensure there's enough data for the end signature
        // Look for the PNG signature
        if (buffer[offset] == 0x89 && buffer[offset + 1] == 0x50 && 
            buffer[offset + 2] == 0x4E && buffer[offset + 3] == 0x47) {
            
            // Check for the end signature starting from the PNG signature position
            size_t endOffset = offset;
            while (endOffset <= buffer.size() - 4) {
                if (buffer[endOffset] == 0xAE && buffer[endOffset + 1] == 0x42 &&
                    buffer[endOffset + 2] == 0x60 && buffer[endOffset + 3] == 0x82) {
                    
                    // We found the end signature; extract the PNG
                    std::string outputFilename = std::to_string(fileCount++) + ".png";
                    std::ofstream outFile(outdir + outputFilename, std::ios::binary);
                    if (outFile) {
                        outFile.write(reinterpret_cast<const char*>(&buffer[offset]), endOffset - offset + 4); // Include end signature
                        std::cout << "Extracted: " << outputFilename << std::endl;
                    } else {
                        std::cerr << "Unable to create file: " << outputFilename << std::endl;
                    }

                    // Move offset to end of current PNG
                    offset = endOffset + 4; // Move past the end signature
                    break; // Break to search for the next PNG
                }
                endOffset++;
            }
        }
        offset++;
    }
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        std::cerr << "Usage: " << argv[0] << " <input_file> <output_dir>" << std::endl;
        return 1;
    }

    const std::string filename = argv[1]; // Take the input file from command-line argument
    const std::string outdir = argv[2];
    std::vector<uint8_t> buffer;

    if (readFileToBuffer(filename, buffer)) {
        extractPNGs(outdir, buffer);
    }

    return 0;
}