#include <iostream>
#include <string>
#include <random>
#include <thread>
#include <atomic>
#include <cstdint>
#include "Header.h"
#include <vector>
#include <omp.h>
#include "ttmath-0.9.3\ttmath\ttmath.h"
using namespace ttmath;
using namespace std;
UInt<512> GenerateRandom ( size_t _sizeInBits, bool _testForPrimility )
{
    PRIMES;
    _sizeInBits = ( _sizeInBits % 32 ) ? _sizeInBits + ( _sizeInBits % 32 ) : _sizeInBits;
    UInt<512> random_number ( 0 );
    random_device rand_dev;
    uniform_int_distribution<uint32_t> engine ( UINT32_MAX/2 ,UINT32_MAX );
    while ( true )
    {
        random_number = 0;
        for ( auto i = 0; i < _sizeInBits; i+=32 )
        {
            UInt<512> temp = engine ( rand_dev );
            temp <<= i;
            random_number += temp;
        }
        if ( _testForPrimility )
        {
            bool bFlag = true;
            for ( auto & i : primes )
            {
                if ( random_number % i == 0 )
                {
                    bFlag = false;
                    break;
                }
            }
            if ( bFlag )
            {
                return random_number;
            }
        }
        else
        {
            return random_number;
        }
    }
}

bool MillerRabinTest ( UInt<512> m )
{
    random_device rand_dev;
    uniform_int_distribution<uint32_t> engine ( 2 ,UINT32_MAX );
    UInt<512> t = m - 1;
    UInt<512> s = 0;
    UInt<512> x = 0;
    UInt<512> r1 = 2;
    UInt<512> r2 = m - 2;
    UInt<512> a;
    while ( t != 0 && ( t % 2 == 0 ) )
    {
        s++;
        t /= 2;
    }
    UInt<512> r = 1;
    for ( UInt<512> i = 0; i < r; i++ )
    {
        a = r1 + UInt<512> ( engine ( rand_dev ) ) % ( r2 - r1 );
        a.PowMod ( t ,m );
        x = a;
        if( x == 1 || x == m - 1 )
        {
            continue;
        }
        for ( UInt<512> j = 0; j < s - 1; j++ )
        {
            x.PowMod ( 2 ,m );
            if ( x == 1 )
            {
                return false;
            }
            if ( x == m - 1 )
            {
                break;
            }
        }
        if ( x == m - 1 )
        {
            continue;
        }
        return false;
    }
    return true;
}

UInt<512> PrimitiveRoot ( UInt<512> _prime )
{
    vector<UInt<512>> fact;
    UInt<512> phi = _prime - 1 ,n = phi;
    for ( UInt<512> i = 2; i*i <= n; ++i )
    {
        if ( n % i == 0 )
        {
            fact.push_back ( i );
            while ( n % i == 0 )
            {
                n /= i;
            }
        }
        i += ( i > 10 ) ? 2 : 0;
    }
    if ( n > 1 )
    {
        fact.push_back ( n );
    }

    cout << "Factorization done\n";
    for (UInt<512> res = 2; res <= _prime; ++res )
    {
        bool ok = true;
        for ( size_t i = 0; i < fact.size ( ) && ok; ++i )
        {
            UInt<512> temp = res;
            temp.PowMod ( phi / fact[ i ] ,_prime );
            ok = ( temp != 1 ) ? true : false;
        }
        if ( ok )
        {
            return res;
        }
    }
    return -1;
}
UInt<512> GeneratePrime ( size_t _sizeInBits )
{
    UInt<512> Prime[ 4 ];
    cout << "Mode - " << _sizeInBits << "bit - generation prime";
    atomic<int> id_done ( -1 );
    vector<thread> threads;
    auto PrimalityThread = [ &id_done ,&Prime ,&_sizeInBits ] ( size_t thread_id )-> void
    {
        while ( true )
        {
            if ( id_done != -1 )
            {
                break;
            }
            Prime[ thread_id ] = GenerateRandom ( _sizeInBits, true );
            cout << ".";
            if ( ( id_done == -1 ) && MillerRabinTest ( Prime[ thread_id ] ) )
            {
                id_done = thread_id;
                cout << "done ( thread "<< id_done << " )\n";
                break;
            }
        }
    };
    for ( int i = 0; i < ((_sizeInBits <= 128 ) ? 1 : 4); i++ )
    {
        threads.push_back ( thread(PrimalityThread, i ) );
    }
    for ( auto &i : threads )
    {
        i.join ( );
    }
    return Prime[ id_done ];
}
void AliceBob ( int key )
{
    UInt<512> PublicPrime = GeneratePrime ( key ); // p
    UInt<512> AliceSecretKey = GenerateRandom ( key ,false ); // a
    UInt<512> BobSecretKey = GenerateRandom ( key ,false );  // b
    cout << "Public Random Prime ( p ): " << PublicPrime << endl;
    UInt<512> PrimeRoot = PrimitiveRoot ( PublicPrime );
    cout << "Primitive root of public prime ( g ): " << PrimeRoot << endl;
    cout << "Alice secret key (a) : " << AliceSecretKey << endl;
    cout << "Bob secret key (b) : " << BobSecretKey << endl;
    UInt<512> temp = PrimeRoot;
    temp.PowMod ( AliceSecretKey ,PublicPrime );
    UInt<512> AlicePublicKey = temp;
    temp = PrimeRoot;
    temp.PowMod ( BobSecretKey ,PublicPrime );
    UInt<512> BobPublicKey = temp;
    cout << "Alice public key (A) : " << AlicePublicKey << endl;
    cout << "Bob public key (B) : " << BobPublicKey << endl;
    UInt<512> AliceValide = BobPublicKey;
    UInt<512> BobValide = AlicePublicKey;
    AliceValide.PowMod ( AliceSecretKey ,PublicPrime );
    BobValide.PowMod ( BobSecretKey ,PublicPrime );
    if ( AliceValide == BobValide )
    {
        cout << "\n\nALICE <-> BOB \nValidation successful!\n" << AliceValide << " = " << BobValide << endl;
    }
    else
    {
        cout << "Validation failed!\n" << AliceValide << " != " << BobValide << endl;
    }
}
void AliceMaloryBob ( int key )
{
    cout << "\n\nATTACK MAN IN THE MIDDLE " << endl;
    UInt<512> PublicPrime = GeneratePrime ( key ); // p
    UInt<512> AliceSecretKey = GenerateRandom ( key ,false ); // a
    UInt<512> BobSecretKey = GenerateRandom ( key ,false );  // b
    UInt<512> MalorySecretKey = GenerateRandom ( key ,false );
    cout << "Public Random Prime ( p ): " << PublicPrime << endl;
    UInt<512> PrimeRoot = PrimitiveRoot ( PublicPrime );
    cout << "Primitive root of public prime ( g ): " << PrimeRoot << endl;
    cout << "Alice secret key (a) : " << AliceSecretKey << endl;
    cout << "Bob secret key (b) : " << BobSecretKey << endl;
    cout << "Malory secret key (c) : " << MalorySecretKey << endl;
    UInt<512> temp = PrimeRoot;
    temp.PowMod ( AliceSecretKey ,PublicPrime );
    UInt<512> AlicePublicKey = temp;
    temp = PrimeRoot;
    temp.PowMod ( BobSecretKey ,PublicPrime );
    UInt<512> BobPublicKey = temp;
    temp = PrimeRoot;
    temp.PowMod ( MalorySecretKey ,PublicPrime );
    UInt<512> MaloryPublicKey = temp;
    cout << "Alice public key (A) : " << AlicePublicKey << endl;
    cout << "Bob public key (B) : " << BobPublicKey << endl;
    cout << "Malory public key (C) : " << MaloryPublicKey << endl;
    UInt<512> AliceValide = MaloryPublicKey;
    UInt<512> BobValide = MaloryPublicKey;
    UInt<512> MaloryBobValide = BobPublicKey;
    UInt<512> MaloryAliceValide = AlicePublicKey;
    AliceValide.PowMod ( AliceSecretKey ,PublicPrime );
    BobValide.PowMod ( BobSecretKey ,PublicPrime );
    MaloryAliceValide.PowMod ( MalorySecretKey ,PublicPrime );
    MaloryBobValide.PowMod ( MalorySecretKey ,PublicPrime );
    cout << "Alice - Bob " << AliceValide << "\t" << BobValide << endl;
    cout << "Alice - Malory " << AliceValide << "\t" << MaloryAliceValide << endl;
    cout << "Malory - Bob " << MaloryBobValide <<"\t" << BobValide << endl;
    if ( AliceValide == BobValide )
    {
        cout << "\n\nALICE <-> BOB \nValidation successful!\n" << AliceValide << " = " << BobValide << endl;
    }
    else
    {
        cout << "\n\nALICE <-> BOB \nValidation failed!\n" << AliceValide << " = " << BobValide << endl;
    }
}
int main ( )
{
    int key;
    cout << "Enter size of key : ";
    cin >> key;
    AliceBob ( key );
    AliceMaloryBob ( key );
    system ( "pause" );
}