#include <cstddef>
#include <filesystem>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <stdexcept>
#include <algorithm>
#include "velvet_ring.h"
#include <unistd.h>


bool is_ascii(const std::string& str) {
    return std::all_of(str.begin(), str.end(), [](unsigned char c) {
        return std::isprint(c) || c == '\n' || c == '\r' || c == '\t';
    });
}

using ByteVec = std::vector<unsigned char>;

ByteVec read_file(const std::string path) {
    std::ifstream f(path, std::ios::binary);
    if (!f) throw std::runtime_error("failed to open: " + path);
    return ByteVec((std::istreambuf_iterator<char>(f)),
                    std::istreambuf_iterator<char>());
}

void write_file(const std::string path, const ByteVec data) {
    std::ofstream f(path, std::ios::binary);
    if (!f) throw std::runtime_error("failed to write: " + path);
    f.write(reinterpret_cast<const char*>(data.data()), data.size());
}

bool ends_with(const std::string s, const std::string suf) {
    return s.size() >= suf.size() &&
           s.compare(s.size() - suf.size(), suf.size(), suf) == 0;
}

std::string enc_out(const std::string in) {
    return in + ".vvt";
}

std::string dec_out(const std::string in) {
    if (ends_with(in, ".vvt"))
        return in.substr(0, in.size() - 4);
    return in + ".dec";
}

std::string ring_name(const std::string& enc) {
    return enc + ".ring";
}

std::string get_opt(int argc, char* argv[], const std::string opt) {
    for (int i = 3; i < argc - 1; ++i)
        if (argv[i] == opt)
            return argv[i + 1];
    return "";
}


std::string help = R"(love is a gentle thing(lx) by maki
v1.1

usage:
    ec-pass  <input file> -p <password> [-o <out>]   # generates a password-protected velvet file
    dec-pass <input file> -p <password> [-o <out>]   # decrypts password-protected velvet files
    ec-key   <input file> [-o <out>] [-p <key path>] # generates a 32-byte ring and a velvet file
    dec-key  <input file> [-o <out>] [-p <key path>] # decrypts a velvet file using a ring file. assumes that the ring file shares the same name as the input file, unless a key is supplied with -p

options:
    -p       provide a password or ring file
    -o       set the destination of processed files
    -r       print content if non-ascii file)";

void usage() {
    std::cout << help << "\n";
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        int retcode = 1;
        if (argc == 1)
            retcode = 0;
        usage();
        return retcode;
    }

    std::string mode = argv[1];
    std::string in = argv[2];
    std::string out = "";
    std::string pass = "";
    bool print_raw = false;
    
    int opt;
    while ((opt = getopt(argc, argv, "ro:p:")) != -1) {
        switch (opt) {
            case 'r':
                print_raw = true;
                break;

            case 'o':
                out = optarg;
                break;

            case 'p':
                pass = optarg;
                break;

            default:
                usage();
                return 1;
        }
    }

    try {
        ByteVec input = read_file(in);

        // using key (ring)
        if (mode == "ec-key") {
            if (out.empty()) out = enc_out(in);
            std::string ring = ring_name(out);
            if (!pass.empty())
                ring = pass;

            ByteVec key = rb(32);
            ByteVec enc = ec_wky(
                std::string(input.begin(), input.end()),
                key
            );

            write_file(out, enc);
            write_file(ring, key);

            std::cout << "encrypted file: " << out << "\n";
            std::cout << "key ring: " << ring << "\n";
        }

        else if (mode == "dec-key") {
            if (out.empty()) out = dec_out(in);
            std::string ring = ring_name(in);
            if (!std::filesystem::exists(ring)) {
                if (!pass.empty())
                    ring = pass;
                else
                    throw std::runtime_error("provide a ring file!");
            }

            ByteVec key = read_file(ring);
            if (key.size() != 32)
                throw std::runtime_error("invalid ring file!");

            std::string dec = dec_wky(input, key);
            write_file(out, ByteVec(dec.begin(), dec.end()));

            if (is_ascii(dec) || print_raw)
                std::cout << dec;

            //std::cout << "Decrypted " << out << "\n";
        }

        // using pass
        else if (mode == "ec-pass") {
            if (pass.empty())
                throw std::runtime_error("missing password! provide one using -p <password>.");

            if (out.empty()) out = enc_out(in);

            ByteVec enc = ec_p(
                std::string(input.begin(), input.end()),
                pass
            );
            write_file(out, enc);

            std::cout << "encrypted file: " << out << "\n";
        }

        else if (mode == "dec-pass") {
            if (pass.empty())
                throw std::runtime_error("missing password! provide one using -p <password>.");

            if (out.empty()) out = dec_out(in);

            std::string dec = dec_p(input, pass);
            write_file(out, ByteVec(dec.begin(), dec.end()));

            if (is_ascii(dec) || print_raw) {
                std::cout << dec;
            }
            
            //std::cout << "Decrypted " << out << "\n"; // i think this is redundnant now cuz duh it better be there
        }

        else {
            usage();
            return 1;
        }
    }
    catch (const std::exception& e) {
        std::cerr << "error: " << e.what() << "\n";
        return 1;
    }

    return 0;
}