// Name: Logan Gillespie
// Class: CST 126
// Date: 5/16/24
// Assignment: Homework 2

#include <iostream>
#include <fstream>
#include <sstream>

using namespace std;

string encodeFileToBase64(const string& input, const string& output);
string decodeBase64(const string& encoded);
bool decodeBase64ToFile(const string& input, const string& output);

constexpr short ARG_COUNT = 4;

//bool Worker();
bool Worker(char** argv);

int main(const int argc, char* argv[])
{
    if (argc == ARG_COUNT) {
        Worker(argv);
    }
    else {
        int choice = 0;
        do {
            cout << "Welcome to my program! Here are your options:" << endl;
            cout << "0. Quit" << endl;
            cout << "1. Encode File to Base64" << endl;
            cout << "2. Decode Base64 to File" << endl;
            cin >> choice;
            switch (choice) {
            case 1:
            {
                string input;
                string output;
                cout << "Enter input filename: ";
                std::cin.clear();
                std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
                getline(cin, input);
                cout << "Enter output filename: ";
                getline(cin, output);
                string base64Data = encodeFileToBase64(input, output);
                cout << "Base64 Encoded Data:\n" << base64Data << endl;
                break;
            }
            case 2:
            {
                string input;
                cout << "Enter input filename: ";
                std::cin.clear();
                std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
                getline(cin, input);

                string output;
                cout << "Enter output filename: ";
                getline(cin, output);

                if (decodeBase64ToFile(input, output)) {
                    cout << "File successfully decoded and written: " << output << endl;
                }
                else {
                    cerr << "Failed to decode and write file." << endl;
                }
                break;
            }
            }
        } while (choice != 0);
    }
    //return Worker();
}

/*bool Worker() {
    char option = 'e';
    
    switch (option) {
    case 'e':
        return true;
    case 'd':
        return true;
    default:
        return false;
    }
}*/
bool Worker(char** argv) {
    const char* arg1 = argv[1];
    const char* arg2 = argv[2];
    const char* arg3 = argv[3];

    const char option = arg1[1];

    string str2 = string("") + arg2;
    string str3 = string("") + arg3;

    switch (option) {
    case 'e':
        encodeFileToBase64(str2, str3);
        return true;
    case 'd':
        decodeBase64ToFile(str2, str3);
        return true;
    default:
        cerr << "Invalid command option\n" << "Valid Commands:\n" << "\t-e source_file.exe destination_file.txt\n" << "\t\tEncodes file in source into text in destination txt file.\n" << "\t-d source_file.txt destination_file.exe\n" << "\t\tDecodes text in source file into the destination file.\n\n";
        return false;
    }
}

string encodeFileToBase64(const string& input, const string& output) {
    ifstream fileInput(input, ios::binary | ios::ate);
    if (!fileInput.is_open()) {
        cerr << "Error: Unable to open file " << input << endl;
        return "";
    }

    // Get file size
    streamsize fileSize = fileInput.tellg();
    fileInput.seekg(0, ios::beg);

    // Read file into buffer
    char* buffer = new char[fileSize];
    fileInput.read(buffer, fileSize);

    // Base64 encoding table
    const string base64_chars =
             "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
             "abcdefghijklmnopqrstuvwxyz"
             "0123456789+/";

    // Encode buffer to Base64
    stringstream encoded;
    int val = 0, valb = -6;
    for (int i = 0; i < fileSize; i++) {
        unsigned char c = buffer[i];
        val = (val << 8) + c;
        valb += 8;
        while (valb >= 0) {
            encoded << base64_chars[(val >> valb) & 0x3F];
            valb -= 6;
        }
    }

    // Handle any remaining bits and add padding
    if (valb > -6) {
        encoded << base64_chars[((val << 8) >> (valb + 8)) & 0x3F];
    }
    while (encoded.str().size() % 4 != 0) {
        encoded << '=';
    }

    delete[] buffer;
    fileInput.close();

    ofstream fileOutput(output);

    if (!fileOutput.is_open()) {
        cerr << "Error: Unable to open file " << input << endl;
        return "";
    }

    fileOutput << encoded.str();
    fileOutput.close();

    return encoded.str();
}

string decodeBase64(const string& encoded) {
    // Base64 decoding table
    const string base64_chars =
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz"
        "0123456789+/";

    string decoded;
    int val = 0, valb = -8;
    for (char c : encoded) {
        if (c == '=') break;
        if (base64_chars.find(c) == string::npos) {
            cerr << "Error: Invalid Base64 character detected!" << endl;
            return "";
        }
        val = (val << 6) + base64_chars.find(c);
        valb += 6;
        if (valb >= 0) {
            decoded.push_back(char((val >> valb) & 0xFF));
            valb -= 8;
        }
    }
    return decoded;
}


bool decodeBase64ToFile(const string& input, const string& output) {
    string base64Data;
    ifstream inputfile(input);
    inputfile >> base64Data;
    inputfile.close();

    ofstream file(output, ios::binary);
    if (!file.is_open()) {
        cerr << "Error: Unable to create file " << output << endl;
        return false;
    }

    string decodedData = decodeBase64(base64Data);
    if (decodedData.empty()) {
        cerr << "Error: Failed to decode Base64 data" << endl;
        return false;
    }

    file.write(decodedData.c_str(), decodedData.size());
    file.close();
    return true;
}