from typing import TextIO
import re
from pathlib import Path
import argparse

ASCII_MODULO = 128

def encrypt_string(string: str, a: int, b: int):
    def encrypt_char(m): return (a * m + b) % ASCII_MODULO
    return ''.join(chr(encrypt_char(ord(char))) for char in string)


def encrypt(plaintext_file: TextIO, output_file: TextIO, a: int, b: int) -> None:
    valid_a = a > 0 and a < ASCII_MODULO and egcd(ASCII_MODULO, a)[0] == 1
    valid_b = b >= 0 and b < ASCII_MODULO

    if (not (valid_a and valid_b)):
        print(f"The key pair ({a}, {b}) is invalid, please select another key")
        return

    plaintext_characters: str = ''.join(plaintext_file.readlines())
    output_file.write(encrypt_string(plaintext_characters, a, b))


def decrypt_string(string: str, inverse_a: int, b: int) -> str:
    def decrypt_char(m): return (inverse_a * (m - b)) % ASCII_MODULO
    return ''.join(chr(decrypt_char(ord(char))) for char in string)


def decrypt(ciphertext_file: TextIO, output_file: TextIO, a: int, b: int) -> None:
    valid_a = a > 0 and a < ASCII_MODULO and egcd(ASCII_MODULO, a)[0] == 1
    valid_b = b >= 0 and b < ASCII_MODULO

    if not (valid_a and valid_b):
        print(f"The key pair ({a}, {b}) is invalid, please select another key")
        return

    inverse_a = modular_inverse(a, ASCII_MODULO)
    ciphered_text: str = ''.join(ciphertext_file.readlines())
    output_file.write(decrypt_string(ciphered_text, inverse_a, b))


def decipher(ciphertext_file: TextIO, output_file: TextIO, dictionary_file: TextIO) -> None:
    dictionary = set(word.strip().lower()
                     for word in dictionary_file.readlines())

    ciphered_text: str = ''.join(ciphertext_file.readlines())

    best_word_count = -1
    best_a = -1
    best_b = -1

    for a in range(1, ASCII_MODULO, 2):
        inverse_a = modular_inverse(a, ASCII_MODULO)

        for b in range(0, ASCII_MODULO):
            decrypted_string = decrypt_string(ciphered_text, inverse_a, b)
            word_count = count_words(decrypted_string, dictionary)

            if word_count > best_word_count:
                best_word_count = word_count
                best_a = a
                best_b = b

    best_inverse_a = modular_inverse(best_a, ASCII_MODULO)
    deciphered_text = decrypt_string(ciphered_text, best_inverse_a, best_b)
    output_file.write(f"{best_a} {best_b}\n")
    output_file.write("DECIPHERED MESSAGE:\n")
    output_file.write(f"{deciphered_text}")


def count_words(string: str, dictionary: set[str]) -> int:
    words: list[str] = re.findall(r'\b\w+\b', string)
    return sum(1 * len(word) for word in words if word.strip().lower() in dictionary)


def modular_inverse(a: int, mod: int) -> int:
    d, s, _ = egcd(a, mod)
    if d != 1:
        return -1  # No modular inverse exists
    return s % mod  # Ensure it's positive


def egcd(a: int, b: int) -> tuple[int, int, int]:
    s, t, u, v = 1, 0, 0, 1

    while b != 0:
        q = a // b
        a, b = b, a % b
        s, t, u, v = u, v, s - u * q, t - v * q

    d = a

    return d, s, t


def create_arg_parser():
    parser = argparse.ArgumentParser(
        description="CLI tool for encryption, decryption, and deciphering.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Encrypt command
    encrypt_parser = subparsers.add_parser(
        "encrypt", help="Encrypt a plaintext file.")
    encrypt_parser.add_argument(
        "plaintext_file", help="Path to the plaintext .txt file.")
    encrypt_parser.add_argument(
        "output_file", help="Path to the output encrypted .txt file.")
    encrypt_parser.add_argument(
        "a", type=int, help="Parameter a for encryption.")
    encrypt_parser.add_argument(
        "b", type=int, help="Parameter b for encryption.")

    # Decrypt command
    decrypt_parser = subparsers.add_parser(
        "decrypt", help="Decrypt a ciphertext file.")
    decrypt_parser.add_argument(
        "ciphertext_file", help="Path to the ciphertext .txt file.")
    decrypt_parser.add_argument(
        "output_file", help="Path to the output decrypted .txt file.")
    decrypt_parser.add_argument(
        "a", type=int, help="Parameter a for decryption.")
    decrypt_parser.add_argument(
        "b", type=int, help="Parameter b for decryption.")

    # Decipher command
    decipher_parser = subparsers.add_parser(
        "decipher", help="Decipher a ciphertext file using a dictionary.")
    decipher_parser.add_argument(
        "ciphertext_file", help="Path to the ciphertext .txt file.")
    decipher_parser.add_argument(
        "output_file", help="Path to the output deciphered .txt file.")
    decipher_parser.add_argument(
        "dictionary_file", help="Path to the dictionary .txt file.")
    return parser


def valid_file_path(path: Path) -> bool:
    return path.exists() and path.suffix == ".txt"


if __name__ == "__main__":
    parser = create_arg_parser()
    args = parser.parse_args()

    if args.command == "encrypt":
        plaintext_file_path = Path(args.plaintext_file).resolve()
        output_file_path = Path(args.output_file).resolve()

        if valid_file_path(plaintext_file_path) and valid_file_path(output_file_path):
            with open(plaintext_file_path, 'r') as plaintext_file, open(output_file_path, 'w') as output_file:
                encrypt(plaintext_file, output_file, args.a, args.b)
        else:
            print("Invalid file path(s), check paths point to .txt files")
            parser.print_help()

    elif args.command == "decrypt":
        ciphertext_file_path = Path(args.ciphertext_file).resolve()
        output_file_path = Path(args.output_file).resolve()

        if valid_file_path(ciphertext_file_path) and valid_file_path(output_file_path):
            with open(ciphertext_file_path, 'r') as ciphertext_file, open(output_file_path, 'w') as output_file:
                decrypt(ciphertext_file, output_file, args.a, args.b)
        else:
            print("Invalid file path(s), check paths point to .txt files")
            parser.print_help()

    elif args.command == "decipher":
        ciphertext_file_path = Path(args.ciphertext_file).resolve()
        output_file_path = Path(args.output_file).resolve()
        dictionary_file_path = Path(args.dictionary_file).resolve()

        if valid_file_path(ciphertext_file_path) and valid_file_path(output_file_path) and valid_file_path(dictionary_file_path):
            with open(ciphertext_file_path, 'r') as ciphertext_file, open(output_file_path, 'w') as output_file, open(dictionary_file_path, 'r') as dictionary_file:
                decipher(ciphertext_file, output_file, dictionary_file)
        else:
            print("Invalid file path(s), check paths point to .txt files")
            parser.print_help()
