from src.models import (
    AlphabetKey, ClientKeys, MutualKeys, ClientPersistentDataKeys, ClientEphemeralDataKeys,
    MutualPersistentDataKeys, MutualEphemeralDataKeys, ClientPersistentMediumKeys, ClientEphemeralMediumKeys,
    MutualPersistentMediumKeys, MutualEphemeralMediumKeys, DarcKey, OuterKey, Mask
)


def darc_phase2(
        alphabet_phase1: AlphabetKey,
        medium_phase1: AlphabetKey,
        client_keys: ClientKeys,
        mutual_keys: MutualKeys
):
    alphabet_phase2 = receive_alphabet(
        alphabet_phase1,
        client_keys.persistent.data,
        client_keys.ephemeral.data,
        mutual_keys.persistent.data,
        mutual_keys.ephemeral.data
    )
    medium_phase2 = receive_medium(
        medium_phase1,
        client_keys.persistent.medium,
        client_keys.ephemeral.medium,
        mutual_keys.persistent.medium,
        mutual_keys.ephemeral.medium
    )
    return alphabet_phase2, medium_phase2


def receive_alphabet(
        alphabet_phase1: AlphabetKey,
        client_persistent_data_keys: ClientPersistentDataKeys,
        client_ephemeral_data_keys: ClientEphemeralDataKeys,
        mutual_persistent_data_keys: MutualPersistentDataKeys,
        mutual_ephemeral_data_keys: MutualEphemeralDataKeys
):
    mdi1 = mutual_persistent_data_keys.inner_key_1 << mutual_ephemeral_data_keys.inner_key_1
    mdo1 = mutual_persistent_data_keys.outer_key_1 << mutual_ephemeral_data_keys.outer_key_1
    mdo2 = mutual_persistent_data_keys.outer_key_2 << mutual_ephemeral_data_keys.outer_key_2
    mdo3 = mutual_persistent_data_keys.outer_key_3 << mutual_ephemeral_data_keys.outer_key_3

    cdi1 = client_persistent_data_keys.inner_key_1 << client_ephemeral_data_keys.inner_key_1
    cdo1 = client_persistent_data_keys.outer_key_1 << client_ephemeral_data_keys.outer_key_1
    cdo2 = client_persistent_data_keys.outer_key_2 << client_ephemeral_data_keys.outer_key_2

    ceda = client_ephemeral_data_keys.alpha_key
    cpda = client_persistent_data_keys.alpha_key
    mpda = mutual_persistent_data_keys.alpha_key

    alphabet_phase2 = ((((
                                  alphabet_phase1 ^
                                  (((ceda ^ cpda ^ mpda) < mdo1) << (mdi1 < mdo2))
                          ) < ~mdo3)) << cdi1) < (mdo3 << cdo1 << cdo2)

    return alphabet_phase2


def receive_medium(
        medium_phase1: AlphabetKey,
        client_persistent_medium_keys: ClientPersistentMediumKeys,
        client_ephemeral_medium_keys: ClientEphemeralMediumKeys,
        mutual_persistent_medium_keys: MutualPersistentMediumKeys,
        mutual_ephemeral_medium_keys: MutualEphemeralMediumKeys
):
    mmi1 = mutual_persistent_medium_keys.inner_key_1 << mutual_ephemeral_medium_keys.inner_key_1
    mmo1 = mutual_persistent_medium_keys.outer_key_1 << mutual_ephemeral_medium_keys.outer_key_1
    mmo2 = mutual_persistent_medium_keys.outer_key_2 << mutual_ephemeral_medium_keys.outer_key_2
    mmo3 = mutual_persistent_medium_keys.outer_key_3 << mutual_ephemeral_medium_keys.outer_key_3

    cmi1 = client_persistent_medium_keys.inner_key_1 << client_ephemeral_medium_keys.inner_key_1
    cmo1 = client_persistent_medium_keys.outer_key_1 << client_ephemeral_medium_keys.outer_key_1
    cmo2 = client_persistent_medium_keys.outer_key_2 << client_ephemeral_medium_keys.outer_key_2

    cema = client_ephemeral_medium_keys.alpha_key
    cpma = client_persistent_medium_keys.alpha_key
    mpma = mutual_persistent_medium_keys.alpha_key

    medium_phase2 = (
                            (
                                    medium_phase1 ^
                                    (((cema ^ cpma ^ mpma) < mmo1) << (mmi1 < mmo2))
                            ) << (cmi1 < mmo3)
                    ) < (cmo1 << cmo2)

    return medium_phase2


def transmit_alphabet(
        alphabet_phase2: DarcKey,
        client_persistent_data_keys: ClientPersistentDataKeys,
        client_ephemeral_data_keys: ClientEphemeralDataKeys,
        mutual_persistent_data_keys: MutualPersistentDataKeys,
        mutual_ephemeral_data_keys: MutualEphemeralDataKeys
):
    mdpi2 = mutual_persistent_data_keys.inner_key_2
    mdi1 = mutual_persistent_data_keys.inner_key_1 << mutual_ephemeral_data_keys.inner_key_1
    mdo1 = mutual_persistent_data_keys.outer_key_1 << mutual_ephemeral_data_keys.outer_key_1
    mdo2 = mutual_persistent_data_keys.outer_key_2 << mutual_ephemeral_data_keys.outer_key_2
    mdo3 = mutual_persistent_data_keys.outer_key_3 << mutual_ephemeral_data_keys.outer_key_3

    cdi1 = client_persistent_data_keys.inner_key_1 << client_ephemeral_data_keys.inner_key_1
    cdo1 = client_persistent_data_keys.outer_key_1 << client_ephemeral_data_keys.outer_key_1
    cdo2 = client_persistent_data_keys.outer_key_2 << client_ephemeral_data_keys.outer_key_2

    cdo1_cdo2 = cdo1 << cdo2

    ceda = client_ephemeral_data_keys.alpha_key
    cpda = client_persistent_data_keys.alpha_key
    meda = mutual_ephemeral_data_keys.alpha_key

    alphabet_phase3 = (
            (alphabet_phase2 << ((~cdi1 << mdpi2) < (mdo3 << cdo1_cdo2))) ^
            ((((ceda ^ cpda ^ meda) < mdo1) << ((mdi1 < mdo2) << (mdpi2 < mdo3))) < cdo1_cdo2)
    )


    alphabet_phase3 = (
            (alphabet_phase2 << ((~cdi1 << mdpi2) < (mdo3 << cdo1_cdo2))) ^
            (
                    (
                            ((ceda ^ cpda ^ meda) < (mdo1 << cdo1_cdo2)) << ((mdi1 < (mdo2 << cdo1_cdo2)) << (mdpi2 < (mdo3 << cdo1_cdo2)))
                    )
            )
    )
    return alphabet_phase3


def transmit_medium(
        medium_phase2: DarcKey,
        client_persistent_medium_keys: ClientPersistentMediumKeys,
        client_ephemeral_medium_keys: ClientEphemeralMediumKeys,
        mutual_persistent_medium_keys: MutualPersistentMediumKeys,
        mutual_ephemeral_medium_keys: MutualEphemeralMediumKeys
):
    mmpi1 = mutual_persistent_medium_keys.inner_key_2
    mmi1 = mutual_persistent_medium_keys.inner_key_1 << mutual_ephemeral_medium_keys.inner_key_1
    mmo1 = mutual_persistent_medium_keys.outer_key_1 << mutual_ephemeral_medium_keys.outer_key_1
    mmo2 = mutual_persistent_medium_keys.outer_key_2 << mutual_ephemeral_medium_keys.outer_key_2
    mmo3 = mutual_persistent_medium_keys.outer_key_3 << mutual_ephemeral_medium_keys.outer_key_3

    cmi1 = client_persistent_medium_keys.inner_key_1 << client_ephemeral_medium_keys.inner_key_1
    cmo1 = client_persistent_medium_keys.outer_key_1 << client_ephemeral_medium_keys.outer_key_1
    cmo2 = client_persistent_medium_keys.outer_key_2 << client_ephemeral_medium_keys.outer_key_2

    cmo1_cmo2 = cmo1 << cmo2

    cpma = client_persistent_medium_keys.alpha_key
    cema = client_ephemeral_medium_keys.alpha_key
    mema = mutual_ephemeral_medium_keys.alpha_key

    medium_phase3 = (
                            ((medium_phase2 < ~cmo1_cmo2) << ((~cmi1 << mmpi1) < mmo3)) ^
                            (((cpma ^ cema ^ mema) < mmo1) << ((mmi1 < mmo2) << (mmpi1 < mmo3)))
                    ) < cmo1_cmo2

    return medium_phase3


def darc_phase3(
        alphabet_phase2: AlphabetKey,
        medium_phase2: AlphabetKey,
        client_keys: ClientKeys,
        mutual_keys: MutualKeys
):
    alphabet_phase3 = transmit_alphabet(
        alphabet_phase2,
        client_keys.persistent.data,
        client_keys.ephemeral.data,
        mutual_keys.persistent.data,
        mutual_keys.ephemeral.data
    )
    medium_phase3 = transmit_medium(
        medium_phase2,
        client_keys.persistent.medium,
        client_keys.ephemeral.medium,
        mutual_keys.persistent.medium,
        mutual_keys.ephemeral.medium
    )
    return alphabet_phase3, medium_phase3


def merge_message(
        alphabet_phase3: AlphabetKey,
        medium_phase3: AlphabetKey,
        client_keys: ClientKeys,
        mutual_keys: MutualKeys,
        input_sequence: list[int]
):
    height = len(alphabet_phase3.matrix)
    width = len(alphabet_phase3.matrix[0])
    msg_len = len(input_sequence)

    eof_msg_mask = Mask.init_matrix(width, height, msg_len)

    if msg_len < height:
        padded_input_sequence = OuterKey.init_matrix(height)
        padded_input_sequence.matrix = [input_sequence + padded_input_sequence.matrix[0][msg_len:]]
    else:
        padded_input_sequence = OuterKey(matrix=[input_sequence])

    cdo1 = client_keys.persistent.data.outer_key_1 << client_keys.ephemeral.data.outer_key_1
    cdo2 = client_keys.persistent.data.outer_key_2 << client_keys.ephemeral.data.outer_key_2

    mmi3 = mutual_keys.persistent.medium.inner_key_3 << mutual_keys.ephemeral.medium.inner_key_3
    mmi4 = mutual_keys.persistent.medium.inner_key_4 << mutual_keys.ephemeral.medium.inner_key_4
    mmo4 = mutual_keys.persistent.medium.outer_key_4 << mutual_keys.ephemeral.medium.outer_key_4
    mmo5 = mutual_keys.persistent.medium.outer_key_5 << mutual_keys.ephemeral.medium.outer_key_5
    mmo6 = mutual_keys.persistent.medium.outer_key_6 << mutual_keys.ephemeral.medium.outer_key_6

    cmo1 = client_keys.persistent.medium.outer_key_1 << client_keys.ephemeral.medium.outer_key_1
    cmo2 = client_keys.persistent.medium.outer_key_2 << client_keys.ephemeral.medium.outer_key_2

    darc_message = (
                           (
                                   (medium_phase3 < ~(cmo1 << cmo2)) ^
                                   (((alphabet_phase3 < (~(cdo1 << cdo2) << padded_input_sequence)) ^ eof_msg_mask) << (mmi3 < ~(mmo4 << mmo6)))
                           ) < mmo4
                   ) << (mmi4 < (mmo5 << ~mmo6))
    return darc_message