#include <iostream>
#include <fstream>
#include <streambuf>

#include "transposition.hpp"
#include "polybius.hpp"
#include "parameter.hpp"

using std::cout;
using std::endl;

std::string groupsOfFive( std::string s );
std::string removeFormat( std::string s );

int main( int argc, char** argv )
{
    CommandlineParameter    cmdParam;
    PolybiusSquare          prng( PolybiusSquare::MODE_RNG, 10 );
    PolybiusSquare          poly( PolybiusSquare::MODE_ALPHANUM, 10 );
    Transposition           trn0;
    Transposition           trn1;

    cmdParam.setArgs( argc, argv );

        cout << endl <<
        "===================================================================\n"
        "ASAC --- A Strong(er) Adfgvx Cipher                                \n"
        "===================================================================\n"
        "                                                                   \n"
        "(C)2016 Stefan Fendt (for the MysteryTwister C3 crypto challenge)  \n"
        "                                                                   \n";

    if( argc == 1 )
    {
        cout <<
        "Usage:   ASAC <parameter> <parameter> ...                          \n"
        "                                                                   \n"
        "Available parameters:                                              \n"
        "---------------------                                              \n"
        "                                                                   \n"
        "  -encode | -decode      This should be rather obvious. There is   \n"
        "                         no default with this parameter and if it  \n"
        "                         does not exist, the program will exit.    \n"
        "                                                                   \n"
        "  -challenge=1|2|3|4     This sets the challenge. Challenge number \n"
        "                         one just is the 10x10 Polybius square     \n"
        "                         alone. Challenge number two is the        \n"
        "                         Polybius square followed by a carry-free  \n"
        "                         addition of the output of the PRNG.       \n"
        "                                                                   \n"
        "  -password1=\"ABCDE\"     This sets the password for the Polybius \n"
        "                         square. It should be twenty letters long. \n"
        "                         It should only contain letters (and no    \n"
        "                         numbers!). If it is shorter than twenty   \n"
        "                         letters it will be extended by repetition.\n"
        "                                                                   \n"
        "  -password2=\"EFGHI\"     This password is used for the PRNG. If  \n"
        "                         not specified and challenge 2 or above is \n"
        "                         wanted, a warning will be generated and   \n"
        "                         the same password as that associated with \n"
        "                         the initial Polybius square will be used. \n"
        "                         The same rules apply as those for the     \n"
        "                         password1-option.                         \n"
        "                                                                   \n"
        "  -password3=\"JKLMN\"     This password is used for the first     \n"
        "                         tabular transposition. It should be longer\n"
        "                         than SQRT(2.0*NR_OF_CHARACTERS) and its   \n"
        "                         length should be pseudo-prime to password4\n"
        "                         if double transposition is being used.    \n"
        "                                                                   \n"
        "  -password4=\"OPQRS\"     This password is used for the first     \n"
        "                         tabular transposition. It should be longer\n"
        "                         than SQRT(2.0*NR_OF_CHARACTERS) and its   \n"
        "                         length should be pseudo-prime to password \n"
        "                         number three.                             \n"
        "                                                                   \n"
        "  -input=\"some text\"     If your text contains more than one     \n"
        "                         word, be sure to use quotation marks...   \n"
        "                                                                   \n"
        "  -file=\"filename\"       This offers the options to enter (a lot \n"
        "                         of) text from within a file.              \n"
        "                                                                   \n"
        "  -prng-no-period        Without this flag the numbers generated by\n"
        "                         the PRNG will repeat after 100 numbers.   \n"
        "                         With this flag, the PRNG will be          \n"
        "                         reshuffled after 100 numbers have been    \n"
        "                         generated.                                \n"
        "                                                                   \n"
        "  -verbose               Be verbose about the internal states and  \n"
        "                         boxes used for the various passes.        \n"
        "Example:                                                           \n"
        "--------                                                           \n"
        "                                                                   \n"
        "ASHC -encode -challenge=1 -password1=\"YADA\" -INPUT=\"SECRETMSG\" \n"
        << endl;
        return( 0 );
    }

    std::string optionStr;

    int mode = 0;

    if( cmdParam.checkToken( "-verbose" ) )
    {
        poly.setVerbose( true );
        prng.setVerbose( true );
        trn0.setVerbose( true );
        trn1.setVerbose( true );
    }

    if( cmdParam.checkToken( "-challenge", optionStr ) )
    {
        if( optionStr == "1" )
        {
            mode = 1;

            if( cmdParam.checkToken( "-password1", optionStr ) )
            {
                poly.setKey( optionStr );
                poly.transpose();
            }
            else
            {
                cout << endl;
                cout << "--[!ERROR!]--------------------------------------------------------" << endl;
                cout << "Setting password #1 for Polybius square is required!               " << endl;
                cout << "-------------------------------------------------------------------" << endl;
                return( 0 );
            }
        }
        else
        if( optionStr == "2" )
        {
            mode = 2;

            if( cmdParam.checkToken( "-password1", optionStr ) )
            {
                poly.setKey( optionStr );
                poly.transpose();
                if( !cmdParam.checkToken( "-password2", optionStr ) )
                {
                    prng.setKey( optionStr );
                    prng.transpose();
                }
            }
            else
            {
                cout << endl;
                cout << "--[!ERROR!]--------------------------------------------------------" << endl;
                cout << "Setting password #1 for Polybius square is required!               " << endl;
                cout << "-------------------------------------------------------------------" << endl;
                return( 0 );
            }

            if( cmdParam.checkToken( "-password2", optionStr ) )
            {
                prng.setKey( optionStr );
                prng.transpose();
            }
            else
            {
                cout << endl;
                cout << "--[WARNING]--------------------------------------------------------" << endl;
                cout << "Are you sure, that you don't want to set a separate PRNG-password? " << endl;
                cout << "-------------------------------------------------------------------" << endl;
            }
        }
        else
        if( optionStr == "3" )
        {
            mode = 3;

            if( cmdParam.checkToken( "-password1", optionStr ) )
            {
                poly.setKey( optionStr );
                poly.transpose();
                if( !cmdParam.checkToken( "-password2", optionStr ) )
                {
                    prng.setKey( optionStr );
                    prng.transpose();
                }
            }
            else
            {
                cout << endl;
                cout << "--[!ERROR!]--------------------------------------------------------" << endl;
                cout << "Setting password #1 for Polybius square is required!               " << endl;
                cout << "-------------------------------------------------------------------" << endl;
                return( 0 );
            }

            if( cmdParam.checkToken( "-password2", optionStr ) )
            {
                prng.setKey( optionStr );
                prng.transpose();
            }
            else
            {
                cout << endl;
                cout << "--[WARNING]--------------------------------------------------------" << endl;
                cout << "Are you sure, that you don't want to set a separate PRNG-password? " << endl;
                cout << "-------------------------------------------------------------------" << endl;
            }

            if( cmdParam.checkToken( "-password3", optionStr ) )
            {
                trn0.setKey( optionStr );
            }
            else
            {
                cout << endl;
                cout << "--[!ERROR!]--------------------------------------------------------" << endl;
                cout << "Setting password #3 for 1st columnar transposition is required!    " << endl;
                cout << "-------------------------------------------------------------------" << endl;
                return( 0 );
            }
        }
        else
        {
            mode = 4;
            if( cmdParam.checkToken( "-password1", optionStr ) )
            {
                poly.setKey( optionStr );
                poly.transpose();
                if( !cmdParam.checkToken( "-password2", optionStr ) )
                {
                    prng.setKey( optionStr );
                    prng.transpose();
                }
            }
            else
            {
                cout << endl;
                cout << "--[!ERROR!]--------------------------------------------------------" << endl;
                cout << "Setting password #1 for Polybius square is required!               " << endl;
                cout << "-------------------------------------------------------------------" << endl;
                return( 0 );
            }

            if( cmdParam.checkToken( "-password2", optionStr ) )
            {
                prng.setKey( optionStr );
                prng.transpose();
            }
            else
            {
                cout << endl;
                cout << "--[WARNING]--------------------------------------------------------" << endl;
                cout << "Are you sure, that you don't want to set a separate PRNG-password? " << endl;
                cout << "-------------------------------------------------------------------" << endl;
            }

            if( cmdParam.checkToken( "-password3", optionStr ) )
            {
                trn0.setKey( optionStr );
            }
            else
            {
                cout << endl;
                cout << "--[!ERROR!]--------------------------------------------------------" << endl;
                cout << "Setting password #3 for 1st columnar transposition is required!    " << endl;
                cout << "-------------------------------------------------------------------" << endl;
                return( 0 );
            }

            if( cmdParam.checkToken( "-password4", optionStr ) )
            {
                trn1.setKey( optionStr );
            }
            else
            {
                cout << endl;
                cout << "--[!ERROR!]--------------------------------------------------------" << endl;
                cout << "Setting password #4 for 2nd columnar transposition is required!" << endl;
                cout << "-------------------------------------------------------------------" << endl;
                return( 0 );
            }
        }
    }

    if( cmdParam.checkToken( "-prng-no-period" ) )
    {
        // makes the prng a lot stronger (at least presumably)
        prng.setPRNGReshuffle();
    }

    std::string text;

    if( !cmdParam.checkToken( "-input", text ) )
    {
        if( cmdParam.checkToken( "-file", text ) )
        {
            std::fstream file( text.c_str() );

            text.clear();

            if( file.is_open() )
            {
                file.seekg( 0, std::ios::end );
                text.reserve( (size_t)file.tellg() );
                file.seekg( 0, std::ios::beg );
                text.assign( std::istreambuf_iterator<char>(file),
                             std::istreambuf_iterator<char>() );
                file.close();
            }
        }
        else
        {
            cout << "NO INPUT GIVEN (or unable to open file)!?" << endl;
            return( 0 );
        }
    }

    cout << endl << "Output:" << endl << endl;

    if( cmdParam.checkToken( "-encode" ) )
    {
        switch( mode )
        {
            case 1:
            {
                text = poly.encode( text );
                text = groupsOfFive( text );
                cout << text << endl << endl;
                break;
            }
            case 2:
            {
                text = poly.encode( text );
                text = prng.addPRNG( text );
                text = groupsOfFive( text );
                cout << text << endl << endl;
                break;
            }
            case 3:
            {
                text = poly.encode( text );
                text = prng.addPRNG( text );
                text = trn0.encode( text );
                text = groupsOfFive( text );
                cout << text << endl << endl;
                break;
            }
            case 4:
            {
                text = poly.encode( text );
                text = prng.addPRNG( text );
                text = trn0.encode( text );
                text = trn1.encode( text );
                text = groupsOfFive( text );
                cout << text << endl << endl;
                break;
            }
        }
    }
    else
    if( cmdParam.checkToken( "-decode" ) )
    {
        switch( mode )
        {
            case 1:
            {
                text = removeFormat( text );
                text = poly.decode( text );
                cout << text << endl << endl;
                break;
            }
            case 2:
            {
                text = removeFormat( text );
                text = prng.subPRNG( text );
                text = poly.decode( text );
                cout << text << endl << endl;
                break;
            }
            case 3:
            {
                text = removeFormat( text );
                text = trn0.decode( text );
                text = prng.subPRNG( text );
                text = poly.decode( text );
                cout << text << endl << endl;
                break;
            }
            case 4:
            {
                text = removeFormat( text );
                text = trn1.decode( text );
                text = trn0.decode( text );
                text = prng.subPRNG( text );
                text = poly.decode( text );
                cout << text << endl << endl;
                break;
            }
        }
    }
    else
    {
        cout << endl << "You *must* specify if you want to encode or decode something." << endl;
        return( 0 );
    }
}

std::string groupsOfFive( std::string s )
{
    int const groupSize = 5;

    std::string out;

    for( int n=0; n<(int)s.size(); ++n )
    {
        out += s[n];

        if( n%(groupSize*10) == (groupSize*10-1) )
        {
            out += '\n';
        }
        else
        if( n%(groupSize) == (groupSize-1) )
        {
            out += ' ';
        }
    }
    return( out );
}

std::string removeFormat( std::string s )
{
    std::string out;

    for( int n=0; n<(int)s.size(); ++n )
    {
        char c = s[n];
        if( c >= '0' && c <='9' )
        {
            out += c;
        }
    }
    return( out );
}


