#include <iostream>
#include <string>

#include "cryptlib.h"
#include "filters.h"
#include "files.h"
#include "sha.h"
#include "hex.h"

using namespace CryptoPP;

class CombinedSource : public Source
{
public:
    // Pumps s1 first, and then s2
    CombinedSource(Source& s1, Source& s2, bool pumpAll,
        BufferedTransformation* attachment)
        : Source(attachment), m_s1(s1), m_s2(s2) {
            if (pumpAll)
                PumpAll();
        }

    virtual size_t Pump2(lword &byteCount, bool blocking=true)
    {    
        std::cout << "Pump2: " << byteCount << std::endl;

        if (byteCount)
            m_s1.TransferTo2(*AttachedTransformation(), byteCount, DEFAULT_CHANNEL, blocking);

        if (byteCount)
            m_s2.TransferTo2(*AttachedTransformation(), byteCount, DEFAULT_CHANNEL, blocking);

        return static_cast<size_t>(byteCount);
    }

    virtual size_t PumpMessages2(unsigned int &messageCount, bool blocking=true)
    {
        std::cout << "PumpMessages2: " << messageCount << std::endl;

        if (messageCount)
            m_s1.TransferMessagesTo2(*AttachedTransformation(), messageCount, DEFAULT_CHANNEL, blocking);

        if (messageCount)
            m_s2.TransferMessagesTo2(*AttachedTransformation(), messageCount, DEFAULT_CHANNEL, blocking);

        return messageCount;
    }

    virtual bool SourceExhausted() const {
        return m_s1.SourceExhausted() && m_s2.SourceExhausted();
    }

private:
    Source &m_s1, &m_s2;
};

int main(int argc, char* argv[])
{
    // Create a file of all 0's with:
    // dd if=/dev/zero of=./zero.dat bs=4096 count=1

    std::string digest;
    SHA256 sha256;

    // Create the hash on the file
    FileSource("zero.dat", true, new HashFilter(sha256, new StringSink(digest)));

    std::cout << "Digest: ";
    StringSource(digest, true, new HexEncoder(new FileSink(std::cout)));
    std::cout << std::endl;

    StringSource ss(digest, true);
    FileSource fs("zero.dat", true);

    byte result = 0;
    std::string digest2;
    CombinedSource cs(ss, fs, true,
        new HashVerificationFilter(sha256,
            // new ArraySink(&result, sizeof(result))));
            new StringSink(digest2), 4));

    std::cout << "Digest: ";
    StringSource(digest2, true, new HexEncoder(new FileSink(std::cout)));
    std::cout << std::endl;

    if (result)
        std::cout << "Verified hash on file" << std::endl;
    else
        std::cout << "Failed to verify hash on file" << std::endl;

    return 0;
}