import logging
from argparse import ArgumentParser
from pathlib import Path

_VERBOSE_SBOX = False
_VERBOSE_PRGA = False
_STEPWISE_PRGA = False


def sbox(key: list[int]) -> list[int]:
    """
    Creates and permutates the substitution box given a key.
    """
    # create sbox
    S = [i for i in range(256)]

    # permutate sbox
    j = 0
    for i in range(256):
        j = (j + S[i] + key[i % len(key)]) % 256
        if _VERBOSE_SBOX:
            print(f"Swapping S[{i=:3d}] = {S[i]:3d} and S[{j=:3d}] = {S[j]:3d}.")
        S[i], S[j] = S[j], S[i]  # swap

    return S


def prga(text: list[int], S: list[int]) -> list[int]:
    """
    Encrypts/Decrypts text given a substitution box.
    """
    text = list(text)
    i = 0
    j = 0
    for x in range(len(text)):
        i = (i + 1) % 256
        j = (j + S[i]) % 256

        if _VERBOSE_PRGA:
            print(f"Swapping S[{i=:3d}] = {S[i]:3d} and S[{j=:3d}] = {S[j]:3d}.")

        if _STEPWISE_PRGA:
            try:
                input("Press ENTER to continue ...")  # Step through
            except KeyboardInterrupt:
                print("KeyboardInterrupt.")
                exit(0)

        S[i], S[j] = S[j], S[i]  # swap
        K = S[(S[i] + S[j]) % 256]
        text[x] ^= K
    return text


def convert(text: bytes | list[int]) -> bytes | list[int]:
    """
    Converts text from bytes -> list[int] and back.
    """
    if isinstance(text, bytes):
        return [int(c) for c in text]
    elif isinstance(text, list):
        return bytes(text)
    else:
        raise ValueError(f"Unsupported input type '{type(text)}'")


def rc4(in_file: str | Path, out_file: str | Path, key_file: str | Path):
    """
    Execute the RC4 encryption/decryption algorithm on the input file,
    creating the output file using the contents from the key file.
    """

    with open(key_file, "rb") as key_file:
        S = sbox(convert(key_file.read()))

    with open("sbox_solution", "wb") as fd:
        fd.write(convert(S))

    # read input
    with open(in_file, "rb") as data_in_file:
        data_in = data_in_file.read()

    data_in = convert(data_in)  # bytes -> list[int]
    data_out = prga(data_in, S)
    data_out = convert(data_out)  # list[int] -> bytes

    # write output
    with open(out_file, "wb") as data_out_file:
        data_out_file.write(data_out)


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Encrypts plaintext using the RC4 encryption algorithm, or decrypts RC4 ciphertext (based on input)"
    )
    parser.add_argument("input", type=str, help="Input file")
    parser.add_argument("key", type=str, help="Key file")
    parser.add_argument("output", type=str, help="Output file")
    parser.add_argument("--verbose-sbox", const=True, default=False, action="store_const", help="Show S-Box swaps")
    parser.add_argument("--verbose-prga", const=True, default=False, action="store_const", help="Show PRGA swaps")
    parser.add_argument("--step-prga", const=True, default=False, action="store_const", help="Step-by-step PRGA swaps")
    args = parser.parse_args()

    _VERBOSE_SBOX = args.verbose_sbox
    _VERBOSE_PRGA = args.verbose_prga
    _STEPWISE_PRGA = args.step_prga

    rc4(args.input, args.output, args.key)