#include <amsstring4/amsstring4.hpp>

namespace ams
{

//PGP / GPG text armor, binary encoding scheme: 
//  

//HTML embedded image file binary encoding scheme:
//example: 
// <img alt="Embedded Image" width="158" height="158" 
//   src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAJ4A..." />
// ref: https://stackoverflow.com/questions/11474346/how-to-encode-images-within-html
// Embed other stuff!
// Data URIs can potentially store any type of data, not just images! Try these examples on for size: (X)HTML CSS Embedding Example
// <link rel="stylesheet" type="text/css"
//   href="data:text/css;base64,LyogKioqKiogVGVtcGxhdGUgKioq..." />
// (X)HTML Javascript Embedding Example
// <script type="text/javascript"
//   src="data:text/javascript;base64,dmFyIHNjT2JqMSA9IG5ldyBzY3Jv..."></script>

//I think these are both base-64 encodings of a binary blob
//ref: https://security.stackexchange.com/questions/142043/how-are-pgp-messages-constructed

//also widely used for email attachments

//Base64: 3 bytes (24-bits) 11111111  22222222 33333333
//  converted to 4 base-64 digits
//                          111111 112222 222233 333333
// subtleties on termination of a string of bytes that doesn't divide by 3

//real reference:
//  https://datatracker.ietf.org/doc/html/rfc4648#section-4


    //              Table 1: The Base 64 Alphabet

    //  Value Encoding  Value Encoding  Value Encoding  Value Encoding
    //      0 A            17 R            34 i            51 z
    //      1 B            18 S            35 j            52 0
    //      2 C            19 T            36 k            53 1
    //      3 D            20 U            37 l            54 2
    //      4 E            21 V            38 m            55 3
    //      5 F            22 W            39 n            56 4
    //      6 G            23 X            40 o            57 5
    //      7 H            24 Y            41 p            58 6
    //      8 I            25 Z            42 q            59 7
    //      9 J            26 a            43 r            60 8
    //     10 K            27 b            44 s            61 9
    //     11 L            28 c            45 t            62 +
    //     12 M            29 d            46 u            63 /
    //     13 N            30 e            47 v
    //     14 O            31 f            48 w         (pad) =
    //     15 P            32 g            49 x
    //     16 Q            33 h            50 y


//ord is 0-63, 64 for '='
ams_chartype base64_char(int8_t ord)
{
    ams_chartype ret;
    if(ord==64)
    {
        ret = '=';
    }
    else if(ord>=0&&ord<26)
    {
        ret = (ams_chartype)(ord+65);
    }
    else if(ord>=26&&ord<52)
    {
        ret = (ams_chartype)((ord-26)+97);
    }
    else if(ord>=52&&ord<62)
    {
        ret = (ams_chartype)((ord-52)+48);
    }
    else if(ord==62)
        ret = '+';
    else if(ord==63)
        ret = '/';
    else
        ret = '\0'; //NULL character is for an out of range ord
    return ret;
}

int8_t base64_ord(ams_chartype ch)
{
    uint8_t ret;
    uint8_t ucval = (unsigned char) ch;
    if(ucval>=65 && ucval<91)
    {
        ret = ucval-65;
    }
    else if(ucval>=97 && ucval<123)
    {
        ret = ucval-97+26;
    }
    else if(ucval>=48 && ucval<58)
    {
        ret = ucval-48+52;
    }
    else if(ch=='+')
    {
        ret = 62;
    }
    else if(ch=='/')
    {
        ret = 63;
    }
    else if(ch=='=')
    {
        ret = 64;
    }
    else
    {
        ret = -1;
    }

    return ret;
}

static void b64_encode_3byteblock(uint8_t *bytes, ams_chartype *chars)
{
    uint8_t o1,o2,o3,o4;

    o1 = ((bytes[0] & 0b11111100)>>2);
    o2 = (((bytes[0] & 0b00000011)<<4) + ((bytes[1] & 0b11110000)>>4));
    o3 = (((bytes[1] & 0b00001111)<<2) + ((bytes[2] & 0b11000000)>>6));
    o4 = ((bytes[2] & 0b00111111));

    chars[0] = base64_char(o1);
    chars[1] = base64_char(o2);
    chars[2] = base64_char(o3);
    chars[3] = base64_char(o4);

    return;
}

static void b64_encode_2byteblock(uint8_t *bytes, ams_chartype *chars)
{
    uint8_t o1,o2,o3;

    o1 = ((bytes[0] & 0b11111100)>>2);
    o2 = (((bytes[0] & 0b00000011)<<4) + ((bytes[1] & 0b11110000)>>4));
    o3 = (((bytes[1] & 0b00001111)<<2));

    chars[0] = base64_char(o1);
    chars[1] = base64_char(o2);
    chars[2] = base64_char(o3);
    chars[3] = '=';

    return;
}

static void b64_encode_1byteblock(uint8_t *bytes, ams_chartype *chars)
{
    uint8_t o1,o2;

    o1 = ((bytes[0] & 0b11111100)>>2);
    o2 = (((bytes[0] & 0b00000011)<<4));

    chars[0] = base64_char(o1);
    chars[1] = base64_char(o2);
    chars[2] = '=';
    chars[3] = '=';

    return;
}

static int decode_4charblock(ams_chartype *chars, uint8_t *bytes)
{
    int ret = 1;

    uint8_t o1,o2,o3,o4,b1,b2,b3;
    int len;

    o1 = base64_ord(chars[0]);
    o2 = base64_ord(chars[1]);
    o3 = base64_ord(chars[2]);
    o4 = base64_ord(chars[3]);
    
    len = 3;
    if(chars[3]=='=') {len--; o4 = 0;}
    if(chars[2]=='=') {len--; o3 = 0;}
    if(chars[1]=='=') {len--; o2 = 0;}
    
    if(len>0)
    {
        //111111 222222 333333 444444
        //11111122 22223333 33444444

        b1 = ((o1<<2)+((o2 & 0b00110000)>>4));
        b2 = (((o2 & 0b00001111)<<4) + ((o3 & 0b00111100)>>2));
        b3 = (((o3 & 0b00000011)<<6) + (o4 & 0b00111111));
        if(len>=1) bytes[0] = b1;
        if(len>=2) bytes[1] = b2;
        if(len>=3) bytes[2] = b3;
    }

    return ret;
}

void base64encode(ams::amsarray<uint8_t> *bytes, amsstring *str)
{
    long I,J;
    
    //str->resize((bytes->length*4)/3+2);
    if( bytes->length % 3 == 0)
    {
        str->resize((bytes->length/3)*4);
    }
    else
    {
        str->resize((bytes->length/3+1)*4);
    }

    I = 0; J = 0;
    while(I<bytes->length)
    {
        if(bytes->length-I>=3)
        {
            //printf("debug3: %c%c%c\n",bytes->data[I],bytes->data[I+1],bytes->data[I+2]);
            b64_encode_3byteblock(&(bytes->data[I]),&(str->cstring[J]));
            I = I + 3;
            J = J + 4;
        }
        else if(bytes->length-I==2)
        {
            //printf("debug2: %c%c\n",bytes->data[I],bytes->data[I+1]);
            b64_encode_2byteblock(&(bytes->data[I]),&(str->cstring[J]));
            I = I + 2;
            J = J + 4;
        }
        else if(bytes->length-I==1)
        {
            //printf("debug1: %c\n",bytes->data[I]);
            b64_encode_1byteblock(&(bytes->data[I]),&(str->cstring[J]));
            I = I + 1;
            J = J + 4;
        }
        else
        {
            break;
        }
    }

    //str->shrinktofit();
    str->cstring[str->length] = '\0';

    return;
}

static long __intl_localstrlen(amsstring *str)
{
    long ret = 0;
    long I;
    for(I=0;I<str->length;I++)
    {
        if(str->cstring[I]=='\0')
        {
            ret = I;
            break;
        }
    }
    return ret;
}


//Only processes strings of length divisible by 4, with
//expected 0,1,2 padding chars at the end of the string,
//and no non-coding characters.
int base64decode_strict(amsstring *str, ams::amsarray<uint8_t> *bytes)
{
    int ret = 1;
    int fail = 0;
    long len;
    long I,J;
    uint8_t v;
    int bp;

    if(str->length == 0)
    {
        bytes->resize(0);
        ret = 1;
        return ret;
    }

    if(str->length % 4 !=0)
    {
        ret = -1; //invalid length
        bytes->resize(0);
        return ret;
    }

    //resize bytes to indicated size

    
    len = (str->length/4)*3;

    if(str->cstring[str->length-1]=='=') len--;
    if(str->cstring[str->length-2]=='=') len--;
    if(str->cstring[str->length-3]=='=')
    {
        ret = -2; //unexpected number of padding chars
        bytes->resize(0);
        return ret;
    }

    bytes->resize(len);

    fail = 0;
    I = 0; J = 0;
    //for(I=0;I<str->length && fail==0;I++)
    while(I<str->length && fail==0)
    {
        for(bp=0;bp<4;bp++)
        {
            v = base64_ord(str->cstring[I+bp]);
            //printf("debug: %d %d %c %d\n",I,I+bp,str->cstring[I+bp],v);
            if(v<0 || (v>=64 && (I+bp)<str->length-2) || v>=65)
            {
                //printf("dbg: fail!\n");
                fail = 1;
                break;
            }
        }
        decode_4charblock(&(str->cstring[I]),&(bytes->data[J]));

        // for(bp=0;bp<3;bp++)
        // {
        //     printf("debug: %d %d %c\n",J+bp,bytes->data[J+bp],bytes->data[J+bp]);
        // }

        I = I + 4;
        J = J + 3;
    }
    if(fail==1)
    {
        ret = -3; //invalid char encountered
        bytes->resize(0);
    }

    //fail on any non-coding chars while parsing

    return ret;
}

//decodes, ignoring (as in MIME spec) all characters that are not
//valid b64 alphabet chars, and all padding until the end of the string
int base64decode_liberal(amsstring *str, ams::amsarray<uint8_t> *bytes)
{
    int ret = 1;
    long I,J;
    amsstring str2;
    ams_chartype ch;

    int64_t v;

    str2.resize(str->length+4);
    
    //filter chars in str to remove all non-coding chars
    J = 0;
    for(I=0;I<str->length;I++)
    {
        ch = str->cstring[I];
        if(ch=='\0')
            break;
        
        v = base64_ord(ch);
        if(v>=0 && v<64)
        {
            str2.cstring[J] = ch;
            J = J + 1;
        }
    }

    if(J%4==0)
    {
        str2.resize(J);
        str2.cstring[str2.length] = '\0';
    }
    if(J%4==1)
    {
        ret = 0; //this isn't really valid
        str2.cstring[J] = '='; J++;
        str2.cstring[J] = '='; J++;
        str2.cstring[J] = '='; J++;
        str2.resize(J);
        str2.cstring[str2.length] = '\0';
    }
    if(J%4==2)
    {
        str2.cstring[J] = '='; J++;
        str2.cstring[J] = '='; J++;
        str2.resize(J);
        str2.cstring[str2.length] = '\0';
    }
    if(J%4==3)
    {
        str2.cstring[J] = '='; J++;
        str2.resize(J);
        str2.cstring[str2.length] = '\0';
    }

    str2.shrinktofit();

    ret = base64decode_strict(&str2,bytes);

    return ret;
}

static void _intl_convsb(amsstring *s, amsarray<uint8_t> *b)
{
    long I;
    b->resize(s->length);
    for(I=0;I<s->length;I++)
    {
        b->at(I) = (unsigned char) s->cstring[I];
    }
    return;
}

static void _intl_convbs(amsarray<uint8_t> *b,amsstring *s)
{
    long I;
    s->resize(b->length);
    for(I=0;I<b->length;I++)
    {
        s->cstring[I] = b->data[I];
    }
    s->cstring[s->length] = '\0';
    return;
}

void test_base64encode1()
{
    int I;
    int8_t o1,o2;
    ams_chartype c1,c2;

    for(I=-5;I<70;I++)
    {
        o1 = I;
        c1 = base64_char(o1);
        o2 = base64_ord(c1);
        c2 = base64_char(o2);

        printf("%d %d %c %d %c\n",(int)I,(int)o1,c1,(int)o2,c2);
    }
    return;
}

//segmentation faults in decode_liberal

void test_base64encode()
{
    printf("Tests of base64 encoding/deconding.\n");

    amsstring s1,s1e,s2;
    amsarray<uint8_t> b1,b2;
    int ret = 0;
    bool bstrict = 1;

    s1 = "light work.";
    _intl_convsb(&s1,&b1);
    base64encode(&b1,&s1e);
    ret = base64decode(&s1e,&b2,bstrict);
    _intl_convbs(&b2,&s2);

    // for(I=0;I<b1.length;I++)
    //     printf("%d:",b1.at(I));
    // printf("\n");

    printf("Original: '%s'\n",s1.cstring);
    printf("Encoded:  '%s'\n",s1e.cstring);
    printf("Decoded:  '%s', ret=%d\n",s2.cstring,ret);


    s1 = "light work";
    _intl_convsb(&s1,&b1);
    base64encode(&b1,&s1e);
    ret = base64decode(&s1e,&b2,bstrict);
    _intl_convbs(&b2,&s2);

    // for(I=0;I<b1.length;I++)
    //     printf("%d:",b1.at(I));
    // printf("\n");

    printf("Original: '%s'\n",s1.cstring);
    printf("Encoded:  '%s'\n",s1e.cstring);
    printf("Decoded:  '%s', ret=%d\n",s2.cstring,ret);

    s1 = "light wor";
    _intl_convsb(&s1,&b1);
    base64encode(&b1,&s1e);
    s1e.insert(3,"\n");
    s1e.insert(5,"\t");
    s1e.insert(7,"}");
    //s1e.insert(1,"}");
    ret = base64decode(&s1e,&b2,bstrict);
    //ret = base64decode_strict(&s1e,&b2);
    _intl_convbs(&b2,&s2);

    // for(I=0;I<b1.length;I++)
    //     printf("%d:",b1.at(I));
    // printf("\n");

    printf("Original: '%s'\n",s1.cstring);
    printf("Encoded:  '%s'\n",s1e.cstring);
    printf("Decoded:  '%s', ret=%d\n",s2.cstring,ret);

    return;
}

void test_base64encode_fuzztest()
{
    long I,J;
    ams::amsarray<uint8_t> bytes,bytes2;
    amsstring str,str2;

    int passed = 0;
    int failed = 0;

    for(I=0;I<100;I++)
    {
        printf("Test %d\n",(int)I);
        bytes.resize(I);
        for(J=0;J<bytes.size();J++)
        {
            bytes.data[J] = ams::rand::rand()*255;
        }
        base64encode(&bytes,&str);
        base64decode(&str,&bytes2,1);
        base64encode(&bytes2,&str2);
        if(bytes==bytes2 && str==str2)
        {
            passed++;
        }
        else
        {
            failed++;
        }
        
    }
    printf("passed: %d, failed %d\n",passed,failed);
    return;
}

int base64decode(amsstring *str, ams::amsarray<uint8_t> *bytes, bool bstrict=0)
{
    int ret = 0;
    if(bstrict==1)
    {
        ret = base64decode_strict(str,bytes);
    }
    else
    {
        ret = base64decode_liberal(str,bytes);
    }
    return ret;
}


};