#!/usr/bin/python
# -*- encoding: utf-8 -*-
import secretpy.alphabets as al


class Enigmini:
    """
    The Enigmini Cipher
    """

    class Rotor:
        def __init__(self, name, key, turnover, position=0, ringsetting=0):
            self.name = name
            self.position = position % 6
            self.ring = ringsetting % 6
            self.key = [int(c) - 1 for c in key]
            self.next_rotor = None
            self.turnover = turnover

        def encrypt_coordinate(self, c):
            i = (c - self.position) % 6
            e = (self.key[i] + self.position) % 6
            #print(f"    rotor {self.name} encrypt {c} -> wiring[{c} + {self.position}] = {e}      wiring: {self.key}")
            return e

        def inverse_encrypt(self, c):
            i = (c - self.position) % 6
            e = (self.key.index(i) + self.position) % 6
            #print(f"    rotor {self.name} inv encrypt {c} -> {self.key.index(i)} + {self.position} = {e}")
            return e

        def rotate(self):
            if str(self.position + 1) in self.turnover:
                if self.next_rotor:
                    self.next_rotor.rotate()
            self.position = (self.position + 1) % 6

        def set_next_rotor(self, rotor):
            self.next_rotor = rotor

        # end of class Rotor

    def encryptCoordinate(self, c, rotors, plugboard):
        # see Errata, plugboard is only used in one direction in the Enigmini
        # c = plugboard[c]

        # we go from right to left
        for r in reversed(rotors):
            c = r.encrypt_coordinate(c)

        it = iter(rotors)
        next(it)  # omit reflector

        # and then go back
        for r in it:
            c = r.inverse_encrypt(c)

        c = plugboard[c]

        # rotate after send a char
        rotors[-1].rotate()

        return c


    def encrypt(self, text, key=None):
        """
        Encryption method

        :param text: Text to encrypt
        :param key: Encryption key
        :type text: string
        :type key: dict
        :return: encrypted text
        :rtype: string
        """
        if not isinstance(key, dict):
            key = {
                'square': '3Z9TWIR7E0XJUNLBSD6VK8GO1A5P4HFYCQM2', # JKP
                'reflector': 'JKP',
                'rotor_offsets': (2, 0), # Grundstellung
                'ring_offsets': (0, 0), # Ringstellung
                'rotor_order': ('JKP_II', 'JKP_I'),
                'steckers': []
            }

        exist_rotors = {
            'JKP_I':  {'wiring': '152346', 'overturn': '6'},
            'JKP_II': {'wiring': '152643', 'overturn': '6'},
            'KP_I':  {'wiring': '254613', 'overturn': '6'},
            'KP_II': {'wiring': '315624', 'overturn': '6'}
        }

        exist_reflectors = {
            'JKP': '645231',
            'KP': '546213'
        }

        plugboard = [int(c) - 1 for c in '123456']
        for s1, s2 in key['steckers']:
            plugboard[s1 - 1] = s2 - 1
            plugboard[s2 - 1] = s1 - 1

        # set reflector, it behaves as a rotor without rotation
        rotors = [self.Rotor('reflector', exist_reflectors[key['reflector']], '', 0, 0)]

        # set other rotors
        for i, ri in enumerate(key['rotor_order']):
            r = exist_rotors[ri]
            rotor = self.Rotor(ri, r['wiring'], r['overturn'], key['rotor_offsets'][i], key['ring_offsets'][i])
            rotors.append(rotor)

        # connect rotors
        for i in range(len(rotors) - 1, 1, -1):
            rotors[i].set_next_rotor(rotors[i - 1])

        # puntuation
        punctuation = "!.\",:? ;_'"
        punct_replace = "1358620974"

        res = []
        for nextChar in text:


            if nextChar in punctuation:
                nextChar = punct_replace[punctuation.index(nextChar)]

            nextIndex = key['square'].index(nextChar)

            #print(f"encrypt {nextChar} ({nextIndex}): Y {nextIndex // 6}")
            encY = self.encryptCoordinate(nextIndex // 6, rotors, plugboard)
            #print(f"encrypt {nextChar} ({nextIndex}): X {nextIndex % 6}")
            encX = self.encryptCoordinate(nextIndex % 6, rotors, plugboard)

            c = key['square'][encY * 6 + encX]
            res.append(c)

        return "".join(res)

    def decrypt(self, text, key=None):
        """
        Decryption method

        :param text: Text to decrypt
        :param key: Decryption key
        :type text: string
        :type key: dict
        :return: decrypted text
        :rtype: string
        """
        return self.encrypt(text, key)