/**
 * @file rc4.h
 * @brief RC4 stream cipher.
 *
 * Declares netplus's implementation of the RC4 stream cipher: a
 * byte-oriented, key-scheduled pseudo-random keystream generator whose
 * output is XORed with data to encrypt or decrypt it (the operation is its
 * own inverse). Retained for interoperability with legacy protocols/formats
 * that still require it.
 *
 * @note RC4 has well-known statistical biases in its keystream and is
 *       considered cryptographically broken for general use; prefer
 *       AES-GCM (see aes.h) for anything requiring real security margin.
 *       This class exists purely for compatibility with legacy consumers.
 */
#include <vector>
#include <string>
#include <cstdint>
#include <stdexcept>

// Using #pragma once is common but often combined with include guards for max portability.
// For simplicity here, we'll keep the structure you provided.
#pragma once

// The size of the state vector S is always 256 bytes
#define SBOX_SIZE 256

namespace netplus {
    /**
     * @brief Implements the RC4 stream cipher (key-scheduled XOR keystream generator).
     *
     * Holds the 256-byte internal permutation state (the "S-box") produced
     * by the Key-Scheduling Algorithm (KSA, see ksa()) and the two PRGA
     * indices used to step the keystream generator (see PRGA_byte()).
     * Because RC4 is a stateful stream cipher, an rc4 instance is
     * single-use per key/stream position: calling crypt() advances the
     * internal state, so it is not safe to call concurrently from multiple
     * threads on the same instance, and the same instance must not be
     * reused to encrypt two different messages under the same key/state
     * (see reset_prga_state() for restarting the keystream from the
     * post-KSA state).
     */
    class rc4 {
    private:
        // The internal state vector S (S-Box)
        // Holds a permutation of the numbers 0 to 255.
        /** @brief The 256-byte internal state vector (a permutation of 0..255), initialized by ksa() and mutated by every PRGA_byte() call. */
        uint8_t S[SBOX_SIZE];

        // Indices for the PRGA part (Pseudo-Random Generation Algorithm)
        /** @brief PRGA index `i`, advanced by one (mod 256) on every keystream byte. */
        int i_prga = 0;
        /** @brief PRGA index `j`, advanced pseudo-randomly (mod 256) based on the current state, on every keystream byte. */
        int j_prga = 0;

        // Helper function to swap two elements
        /** @brief Swaps two state-vector bytes in place. */
        void swap(uint8_t& a, uint8_t& b);

        /**
        * @brief Initializes the state vector S (Key-Scheduling Algorithm - KSA).
        *
        * Initializes S to the identity permutation, then scrambles it using
        * the key bytes (repeated cyclically if shorter than 256 bytes), per
        * the standard RC4 KSA. Also resets the PRGA indices to 0 via
        * reset_prga_state() so the instance is immediately ready to
        * generate keystream.
        * @param key The secret key.
        * @param key_len The length of the key in bytes; must be between 1 and 256 inclusive.
        * @throws std::invalid_argument if @p key_len is 0 or greater than 256.
        */
        void ksa(const uint8_t* key, size_t key_len); // ksa name is used here

        /**
        * @brief Generates a single keystream byte (Pseudo-Random Generation Algorithm - PRGA).
        *
        * Advances i_prga and j_prga, swaps the corresponding S-box entries,
        * and returns S[(S[i]+S[j]) mod 256] as the next keystream byte.
        * Mutates the internal state, so each call produces the next byte
        * in the keystream sequence.
        * @return The next keystream byte.
        */
        uint8_t PRGA_byte();

    public:
        /**
        * @brief Constructor, initializes the RC4 instance with a key.
        *
        * Runs the Key-Scheduling Algorithm (ksa()) once to derive the
        * initial S-box permutation from @p key.
        * @param key The secret key as a std::vector of bytes; length must be between 1 and 256 bytes.
        * @throws std::invalid_argument if @p key is empty or longer than 256 bytes.
        */
        rc4(const std::vector<uint8_t>& key);

        /**
        * @brief Encrypts or decrypts data.
        * Since RC4 is a stream cipher and uses XOR, the function is symmetric.
        *
        * Each call consumes the next |data| bytes of keystream (advancing
        * the internal PRGA state), so calling crypt() repeatedly on the
        * same instance continues the keystream from where the previous
        * call left off rather than restarting it.
        * @param data The data to be encrypted/decrypted.
        * @return The encrypted/decrypted data, the same length as @p data.
        */
        std::vector<uint8_t> crypt(const std::vector<uint8_t>& data);

        /**
        * @brief Resets the PRGA state to start a new encryption/decryption
        * with the same key and the same initialized S-box state.
        * The S-box itself remains preserved.
        *
        * Resets i_prga and j_prga to 0 without re-running the KSA, so the
        * S-box permutation as it stood at the time of the reset (whether
        * freshly keyed or already advanced by prior crypt() calls) becomes
        * the starting point for the next keystream.
        */
        void reset_prga_state();
    };
}
