'''
Implementation of various substitution cyphers, some of which are 
already in the standard library but lack features/options I require.

Terminology:
1. "Simple": simple substitution cyphers operates on single letters
2. "Polygraphic": polygraphic substitution cyphers operate on groups
                  of letters larger than length 1.
3. "Monoalphabetic cypher": a monoalphabetic cypher applies the same fixed
                            substitutions to the entire message (ie ROT13)
4. "Polyalphabetic cypher": a polyalphabetic cypher applies different 
                            substitutions to the entire message (ie vigenere)
'''

from celeste.constants import ALPHA_LOWER, ALPHA_UPPER

'''
Constant Declarations 
'''
# Error Message Format Strings
ERRMSG_CHARSET_DUPL = 'Bad charsets for {0}: found duplicate character \'{1}\''
ERRMSG_NOT_IN_CHARSET = 'No charset for {0} has character \'{1}\''
# Charsets
CHARSET_SUB_DEFAULT = [ALPHA_LOWER, ALPHA_UPPER]

'''
==========================================
Simple Monoalphabetic Substitution Cyphers
==========================================
'''
def ROT(plaintext: str | list, 
        n: int, 
        charsets: list[str] | list[list] = CHARSET_SUB_DEFAULT, 
        ignore_noncharset: bool = True,
        as_string: bool = True) -> str:
    cyphertext = []
    for c in plaintext:
        i = None
        active_charset = None
        for charset in charsets:
            try:
                i = charset.index(c)
                # if we reached here then no ValueError occured
                if active_charset is not None:
                    raise ValueError(ERRMSG_CHARSET_DUPL.format('ROT', c))
                active_charset = charset
            except ValueError: pass
            
        if i is not None:
            cyphertext.append(active_charset[(i + n) % len(active_charset)])
        elif not ignore_noncharset:
            raise ValueError(ERRMSG_NOT_IN_CHARSET.format('ROT', c))
        else:
            # if ignore_noncharset then just append c
            cyphertext.append(c)
    if as_string:
        return ''.join(cyphertext)
    return cyphertext

# Common ROT aliases
def ROT13(plaintext: str, charsets: list[str] = CHARSET_SUB_DEFAULT) -> str: 
    return ROT(plaintext, 13, charsets=charsets)
def caesar(plaintext: str, charsets: list[str] = CHARSET_SUB_DEFAULT) -> str: 
    return ROT(plaintext, 1, charsets=charsets)

# The Atbash Cypher maps a character to its reverse (ie a -> z, b -> y, etc)
def atbash(plaintext: str | list, 
           charsets: list[str] | list[list] = CHARSET_SUB_DEFAULT,
           ignore_noncharset: bool = True,
           as_string: bool = True) -> str:
    cyphertext = []
    for c in plaintext:
        i = None
        active_charset = None
        for charset in charsets:
            try:
                i = charset.index(c)
                # if we reached here then no ValueError occured
                if active_charset is not None:
                    raise ValueError(ERRMSG_CHARSET_DUPL.format('Atbash', c))
                active_charset = charset
            except ValueError: pass
            
        if i is not None:
            L = len(active_charset)
            cyphertext.append(active_charset[(L - i - 1) % L])
        elif not ignore_noncharset:
            raise ValueError(ERRMSG_NOT_IN_CHARSET.format('Atbash', c))
        else:
            # if ignore_noncharset then just append c
            cyphertext.append(c)
    if as_string:
        return ''.join(cyphertext)
    return cyphertext

'''
=========================================
Simple Polyalphabetic Substituion Cyphers
========================================= 
'''
def vigenere(plaintext: str | list, 
             key: str | list, 
             charsets: list[str] | list[list] = CHARSET_SUB_DEFAULT,
             key_charset: str | list = ALPHA_UPPER,
             decrypt: bool = False,
             ignore_noncharset: bool = True,
             as_string: bool = True) -> str:
    cyphertext = []
    # map the key characters to their rotations
    rots = [key_charset.index(k) for k in key]
    pos = 0
    for c in plaintext:
        i = None
        active_charset = None
        for charset in charsets:
            try:
                i = charset.index(c)
                # if we reached here then no ValueError occured
                if active_charset is not None:
                    raise ValueError(ERRMSG_CHARSET_DUPL.format('Vigenere', c))
                active_charset = charset
            except ValueError: pass
            
        if i is not None:
            rot = rots[pos % len(rots)]
            if decrypt: rot = -rot
            cyphertext.append(active_charset[(i + rot) % len(active_charset)])
            pos += 1
        elif not ignore_noncharset:
            raise ValueError(ERRMSG_NOT_IN_CHARSET.format('Vigenere', c))
        else:
            # if ignore_noncharset then just append c
            cyphertext.append(c)
    if as_string:
        return ''.join(cyphertext)
    return cyphertext