import secrets
from enum import Enum
from pydantic import BaseModel


class DarcKeyType(str, Enum):
    outer_shuffle = "outer_shuffle"
    inner_shuffle = "inner_shuffle"
    substitution = "substitution"
    other = "other"


class DarcKey(BaseModel):
    matrix: list[list[int]]
    key_type: DarcKeyType

    @classmethod
    def _random_shuffle(cls, array: list[int]) -> list[int]:
        for i in range(len(array) - 1, 0, -1):
            j = cls._random_number_in_range(0, i)
            array[i], array[j] = array[j], array[i]
        return array

    @classmethod
    def _random_mask(cls, array_len: int, max_val: int) -> list[int]:
        return [cls._random_number_in_range(0, max_val) for _ in range(array_len)]

    @staticmethod
    def _random_number_in_range(min_val: int, max_val: int) -> int:
        """ originally secure_rand """
        return min_val + secrets.randbelow(max_val - min_val + 1)

    def __lshift__(self, other):
        assert len(self.matrix) == len(other.matrix)
        assert len(self.matrix[0]) == len(other.matrix[0])
        assert other.key_type == DarcKeyType.inner_shuffle or other.key_type == DarcKeyType.outer_shuffle
        output = []
        for i, el in enumerate(self.matrix):
            row = []
            for j, el2 in enumerate(el):
                other_el = other.matrix[i][j]
                row.append(self.matrix[i][other_el])
            output.append(row)
        return DarcKey(
            matrix=output,
            key_type=self.key_type
        )

    def __lt__(self, other):
        assert len(other.matrix) == 1
        assert len(self.matrix) == len(other.matrix[0])
        assert other.key_type == DarcKeyType.outer_shuffle
        output = []
        for i, _ in enumerate(self.matrix):
            output.append(self.matrix[other.matrix[0][i]])
        return DarcKey(
            matrix=output,
            key_type=self.key_type
        )

    def __xor__(self, other):
        assert len(self.matrix) == len(other.matrix)
        assert len(self.matrix[0]) == len(other.matrix[0])
        assert self.key_type == DarcKeyType.substitution
        assert other.key_type == DarcKeyType.substitution
        output = []
        for i, el in enumerate(self.matrix):
            row = []
            for j, el2 in enumerate(el):
                row.append(el2 ^ other.matrix[i][j])
            output.append(row)
        return DarcKey(
            matrix=output,
            key_type=self.key_type
        )

    def __invert__(self):
        assert self.key_type == DarcKeyType.inner_shuffle or self.key_type == DarcKeyType.outer_shuffle

        output = []
        for i, el in enumerate(self.matrix):
            row = [(j, el2) for j, el2 in enumerate(el)]
            row.sort(key=lambda x: x[1])
            output.append([el2[0] for el2 in row])

        return DarcKey(
            matrix=output,
            key_type=self.key_type
        )

    def __eq__(self, other) -> bool:
        if self.key_type != other.key_type:
            return False

        if len(self.matrix) != len(other.matrix):
            return False

        if len(self.matrix[0]) != len(other.matrix[0]):
            return False

        for i in range(len(self.matrix)):
            for j in range(len(self.matrix[0])):
                if self.matrix[i][j] != other.matrix[i][j]:
                    return False
        return True


class OuterKey(DarcKey):
    key_type: DarcKeyType = DarcKeyType.outer_shuffle

    @classmethod
    def init_matrix(cls, height: int):
        """ originally buildRandom """
        matrix = [cls._random_shuffle(list(range(height)))]
        return OuterKey(
            matrix=matrix,
        )

    @classmethod
    def init_identity_matrix(cls, height: int):
        """ originally buildRandom """
        matrix = [list(range(height))]
        return OuterKey(
            matrix=matrix,
        )


class InnerKey(DarcKey):
    key_type: DarcKeyType = DarcKeyType.inner_shuffle

    @classmethod
    def init_matrix(cls, width: int, height: int):
        """ originally buildRandom """
        matrix = [cls._random_shuffle(list(range(width))) for _ in range(height)]
        return InnerKey(
            matrix=matrix,
        )

    @classmethod
    def init_identity_matrix(cls, width: int, height: int):
        """ originally buildRandom """
        matrix = [list(range(width)) for _ in range(height)]
        return InnerKey(
            matrix=matrix,
        )


class SubstitutionKey(DarcKey):
    key_type: DarcKeyType = DarcKeyType.substitution

    @classmethod
    def init_matrix(cls, width: int, height: int, max_value: int = 255):
        """ originally buildScramble """
        matrix = [cls._random_mask(width, max_value) for _ in range(height)]
        return SubstitutionKey(
            matrix=matrix,
        )

    @classmethod
    def init_identity_matrix(cls, width: int, height: int):
        """ originally buildScramble """
        matrix = [[0 for _ in range(width)] for _ in range(height)]
        return SubstitutionKey(
            matrix=matrix,
        )


class Mask(DarcKey):
    key_type: DarcKeyType = DarcKeyType.substitution

    @classmethod
    def init_matrix(cls, width: int, height: int, col_mask: int):
        """ originally buildScramble """
        max_value = 255
        mask = cls._random_mask(width, max_value)
        matrix = [[0 for _ in range(width)] for _ in range(height)]
        if col_mask < height:
            matrix[col_mask] = mask

        return Mask(
            matrix=matrix,
        )


class ClientEphemeralDataKeys(BaseModel):
    outer_key_1: OuterKey
    outer_key_2: OuterKey
    alpha_key: SubstitutionKey
    inner_key_1: InnerKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ClientEphemeralDataKeys(
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            inner_key_1=InnerKey.init_matrix(width, height),
        )


class ClientEphemeralMediumKeys(BaseModel):
    outer_key_1: OuterKey
    outer_key_2: OuterKey
    inner_key_1: InnerKey
    alpha_key: SubstitutionKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ClientEphemeralMediumKeys(
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
            inner_key_1=InnerKey.init_matrix(width, height),
            alpha_key=SubstitutionKey.init_matrix(width, height, 255)
        )


class ClientPersistentDataKeys(BaseModel):
    alpha_key: SubstitutionKey
    inner_key_1: InnerKey
    outer_key_1: OuterKey
    outer_key_2: OuterKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ClientPersistentDataKeys(
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            inner_key_1=InnerKey.init_matrix(width, height),
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
        )


class ClientPersistentMediumKeys(BaseModel):
    outer_key_1: OuterKey
    outer_key_2: OuterKey
    inner_key_1: InnerKey
    alpha_key: SubstitutionKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ClientPersistentMediumKeys(
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
            inner_key_1=InnerKey.init_matrix(width, height),
            alpha_key=SubstitutionKey.init_matrix(width, height, 255)
        )


class ServerEphemeralDataKeys(BaseModel):
    outer_key_1: OuterKey
    alpha_key: SubstitutionKey
    inner_key_1: InnerKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ServerEphemeralDataKeys(
            outer_key_1=OuterKey.init_matrix(height),
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            inner_key_1=InnerKey.init_matrix(width, height),
        )


class ServerEphemeralMediumKeys(BaseModel):
    alpha_key: SubstitutionKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ServerEphemeralMediumKeys(
            alpha_key=SubstitutionKey.init_matrix(width, height, 255)
        )


class ServerPersistentDataKeys(BaseModel):
    outer_key_1: OuterKey
    alpha_key: SubstitutionKey
    inner_key_1: InnerKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ServerPersistentDataKeys(
            outer_key_1=OuterKey.init_matrix(height),
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            inner_key_1=InnerKey.init_matrix(width, height),
        )


class ServerPersistentMediumKeys(BaseModel):
    alpha_key: SubstitutionKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return ServerPersistentMediumKeys(
            alpha_key=SubstitutionKey.init_matrix(width, height, 255)
        )


class MutualEphemeralDataKeys(BaseModel):
    alpha_key: SubstitutionKey
    inner_key_1: InnerKey
    outer_key_1: OuterKey
    outer_key_2: OuterKey
    outer_key_3: OuterKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return MutualEphemeralDataKeys(
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            inner_key_1=InnerKey.init_matrix(width, height),
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
            outer_key_3=OuterKey.init_matrix(height),
        )


class MutualEphemeralMediumKeys(BaseModel):
    alpha_key: SubstitutionKey

    outer_key_1: OuterKey
    outer_key_2: OuterKey
    outer_key_3: OuterKey
    outer_key_4: OuterKey
    outer_key_5: OuterKey
    outer_key_6: OuterKey

    inner_key_1: InnerKey
    inner_key_4: InnerKey
    inner_key_3: InnerKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return MutualEphemeralMediumKeys(
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
            outer_key_3=OuterKey.init_matrix(height),
            outer_key_4=OuterKey.init_matrix(height),
            outer_key_5=OuterKey.init_matrix(height),
            outer_key_6=OuterKey.init_matrix(height),
            inner_key_1=InnerKey.init_matrix(width, height),
            inner_key_4=InnerKey.init_matrix(width, height),
            inner_key_3=InnerKey.init_matrix(width, height),
        )


class MutualPersistentDataKeys(BaseModel):
    outer_key_1: OuterKey
    outer_key_2: OuterKey
    outer_key_3: OuterKey
    alpha_key: SubstitutionKey
    inner_key_1: InnerKey
    inner_key_2: InnerKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return MutualPersistentDataKeys(
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
            outer_key_3=OuterKey.init_matrix(height),
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            inner_key_1=InnerKey.init_matrix(width, height),
            inner_key_2=InnerKey.init_matrix(width, height),
        )


class MutualPersistentMediumKeys(BaseModel):
    alpha_key: SubstitutionKey

    outer_key_1: OuterKey
    outer_key_2: OuterKey
    outer_key_3: OuterKey
    outer_key_4: OuterKey
    outer_key_5: OuterKey
    outer_key_6: OuterKey

    inner_key_1: InnerKey
    inner_key_2: InnerKey
    inner_key_3: InnerKey
    inner_key_4: InnerKey

    @classmethod
    def random_init(cls, height: int, width: int):
        return MutualPersistentMediumKeys(
            alpha_key=SubstitutionKey.init_matrix(width, height, 255),
            outer_key_1=OuterKey.init_matrix(height),
            outer_key_2=OuterKey.init_matrix(height),
            outer_key_3=OuterKey.init_matrix(height),
            outer_key_4=OuterKey.init_matrix(height),
            outer_key_5=OuterKey.init_matrix(height),
            outer_key_6=OuterKey.init_matrix(height),
            inner_key_1=InnerKey.init_matrix(width, height),
            inner_key_2=InnerKey.init_matrix(width, height),
            inner_key_3=InnerKey.init_matrix(width, height),
            inner_key_4=InnerKey.init_matrix(width, height),
        )

class ServerPersistentKeys(BaseModel):
    data: ServerPersistentDataKeys
    medium: ServerPersistentMediumKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return ServerPersistentKeys(
            data=ServerPersistentDataKeys.random_init(height, width),
            medium=ServerPersistentMediumKeys.random_init(height, width)
        )


class ClientPersistentKeys(BaseModel):
    data: ClientPersistentDataKeys
    medium: ClientPersistentMediumKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return ClientPersistentKeys(
            data=ClientPersistentDataKeys.random_init(height, width),
            medium=ClientPersistentMediumKeys.random_init(height, width)
        )


class MutualPersistentKeys(BaseModel):
    data: MutualPersistentDataKeys
    medium: MutualPersistentMediumKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return MutualPersistentKeys(
            data=MutualPersistentDataKeys.random_init(height, width),
            medium=MutualPersistentMediumKeys.random_init(height, width)
        )


class ServerEphemeralKeys(BaseModel):
    data: ServerEphemeralDataKeys
    medium: ServerEphemeralMediumKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return ServerEphemeralKeys(
            data=ServerEphemeralDataKeys.random_init(height, width),
            medium=ServerEphemeralMediumKeys.random_init(height, width)
        )


class MutualEphemeralKeys(BaseModel):
    data: MutualEphemeralDataKeys
    medium: MutualEphemeralMediumKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return MutualEphemeralKeys(
            data=MutualEphemeralDataKeys.random_init(height, width),
            medium=MutualEphemeralMediumKeys.random_init(height, width)
        )


class ClientEphemeralKeys(BaseModel):
    data: ClientEphemeralDataKeys
    medium: ClientEphemeralMediumKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return ClientEphemeralKeys(
            data=ClientEphemeralDataKeys.random_init(height, width),
            medium=ClientEphemeralMediumKeys.random_init(height, width)
        )


class ServerKeys(BaseModel):
    persistent: ServerPersistentKeys
    ephemeral: ServerEphemeralKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return ServerKeys(
            persistent=ServerPersistentKeys.random_init(height, width),
            ephemeral=ServerEphemeralKeys.random_init(height, width)
        )


class MutualKeys(BaseModel):
    persistent: MutualPersistentKeys
    ephemeral: MutualEphemeralKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return MutualKeys(
            persistent=MutualPersistentKeys.random_init(height, width),
            ephemeral=MutualEphemeralKeys.random_init(height, width)
        )


class ClientKeys(BaseModel):
    persistent: ClientPersistentKeys
    ephemeral: ClientEphemeralKeys

    @classmethod
    def random_init(cls, height: int, width: int):
        return ClientKeys(
            persistent=ClientPersistentKeys.random_init(height, width),
            ephemeral=ClientEphemeralKeys.random_init(height, width)
        )