/* Copyright (C) CJ Affiliate
  *
  * You may use, distribute and modify this code under  the
  * terms of the  GNU General Public License  version 2  or
  * later.
  *
  * You should have received a copy of the license with this
  * file. If not, you will find a copy in the "LICENSE" file
  * at https://github.com/cjdev/dual-control.
  */
 

 #include <bitset>

 #include <unordered_map>

 #include <memory>

 #include <iostream>

 #include "base32.h"

 
 static const std::vector<char> alphabet = {
     'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
     'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
     'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
     'Y', 'Z', '2', '3', '4', '5', '6', '7'
 };
 

 namespace
 {

 class invalid_input_size : public std::exception
 {
 private:
     using string_size = std::string::size_type;
 public:
     string_size received_size;

     invalid_input_size (string_size received_size) :

         received_size (received_size)
     {}
 };
 

 class base32_impl : public base32_ifc
 {
 private:

     uint64_t calculate_padding_bytes (uint64_t extra_bytes) const

     {
         auto padding_chars = 0;
 
         switch (extra_bytes) {
         case 1:
             padding_chars = 6;
             break;
 
         case 2:
             padding_chars = 4;
             break;
 
         case 3:
             padding_chars = 3;
             break;
 
         case 4:
             padding_chars = 1;
             break;
         }

         return padding_chars;
     }

     std::string &pad_string (uint64_t extra_bytes, std::string &input) const

     {
         auto padding_chars = calculate_padding_bytes (extra_bytes);

         if (padding_chars > 0) {
             auto replace_end   = input.end();
             auto replace_start = replace_end - padding_chars;
             auto replace_count = replace_end - replace_start;

             input.replace (replace_start, replace_end, replace_count, '=');

         }
 

         return input;

     }
 

 public:

     std::string encode (const std::vector<uint8_t> data) const override

     {
         std::string result;

         uint8_t tmp = 0;
 
         auto data_size = data.size();
         auto extra_bytes = data_size % 5;
 
         auto begin = data.begin();
         auto end = data.end();
         auto leftover = (5 - extra_bytes) % 5;
 
         for (auto cur = begin; cur+5 < (end + leftover + 1); cur+=5) {
             std::vector<uint8_t> batch;
 
             if (cur+5 < end) {
                 batch = std::vector<uint8_t> (cur, cur+5);
             } else {
                 batch = std::vector<uint8_t> (cur, end);
 
                 for (int x = 0; x < leftover; x++) {
                     batch.push_back (0);
                 }
             }
 
             std::bitset<40> everything;

             for (int x = 0, y = 32; x < 5; x+=1, y=32-x*8) {
                 uint64_t item = batch[x];
                 everything |= item << y;
             }

             std::bitset<40> mask = 31;
             int offset = 35;
             mask <<= offset;
 
             for (/* see above */; offset >= 0; mask >>= 5, offset -= 5) {
                 auto idx = ((everything & mask) >> offset).to_ullong();
                 result.push_back (alphabet[idx]);
             }

         }
 

         result = pad_string (extra_bytes, result);
         return result;

     }
 

 private:

     using reverse_map = std::unordered_map <char, unsigned char>;
     using character_type = std::string::value_type;

     reverse_map construct_lookup_table() const

     {
         std::unordered_map<char, unsigned char> lookup_table;

         for (int i = 0; i < alphabet.size(); i += 1) {
             lookup_table[alphabet[i]] = i;
         }
 
         return lookup_table;

     }
 

     std::pair<uint8_t, uint8_t> split_value (uint8_t item, uint8_t pos,

             uint8_t width) const

     {
         auto shift = width-pos;
         uint8_t part2_mask = (1 << (shift)) - 1;
         uint8_t part2 = item & part2_mask;
         item >>= shift;
         return { item, part2 };
     }

     void set_vector_at_bit (std::vector<uint8_t> &data, uint8_t item,
                             std::vector<uint8_t>::size_type byte_offset, uint8_t bit_offset_from_start,

                             uint8_t width) const

     {
         bool need_to_split = bit_offset_from_start + width > 8;
         uint8_t bit_offset_from_end = 7 - bit_offset_from_start;
 
         if (need_to_split) {
             uint8_t split_pos = bit_offset_from_end + 1;
             auto pieces = split_value (item, split_pos, width);
             data[byte_offset] |= pieces.first;
 
             uint8_t leftover_bits = width - split_pos;
             uint8_t second_part_shift = 8 - leftover_bits;
             data[byte_offset+1] |= pieces.second << second_part_shift;
         } else {
             uint8_t last_bit = (bit_offset_from_end - width) + 1;
             data[byte_offset] |= item << last_bit;
         }

     }

     std::string::size_type calculate_decoded_size (std::string input) const

     {
         std::string::size_type input_size = input.size();
         std::string::size_type first_equals = input.find_first_of ('=');

         if (input_size == 0 || (input_size % 8 != 0)) {

             throw invalid_input_size (input_size);

         }
 

         if (first_equals != std::string::npos) {
             input_size = first_equals;
         }

         return (input_size * 5) / 8;

     }

     uint8_t get_next_input (reverse_map &lookup_table, std::string &input,
                             std::string::size_type idx) const
     {

         character_type character = input[idx];
 

         if (character != '='
             && lookup_table.find (character) == lookup_table.end()) {

             throw invalid_data_exception ();
         }
 
         return lookup_table[input[idx]];
     }
 

 public:

     std::vector<uint8_t> decode (std::string input) const override

     {
         auto lookup_table = construct_lookup_table();
         auto input_size = calculate_decoded_size (input);

         std::vector<uint8_t> result (input_size, 0);

         unsigned long long bits_written = 0;

         for (std::string::size_type idx = 0; idx < input.size(); idx++) {
             uint8_t start_bit = bits_written % 8;
             std::vector<unsigned char>::size_type current_byte = bits_written/8;

             uint8_t val = get_next_input (lookup_table, input, idx);

             set_vector_at_bit (result, val, current_byte, start_bit, 5);

             bits_written += 5;
         }

         return result;

     }

 };

 }

 
 template class std::vector<unsigned char>;

 base32::base32 ():

     delegate_ (std::make_shared<base32_impl> ())
 {}