from src.client_darc_noclient import merge_message_no_client, darc_phase_2_3
from src.client_darc import merge_message, darc_phase2, darc_phase3
from src.key_generation import generate_keys
from src.models import SubstitutionKey, OuterKey
from src.server_darc import darc_phase1, darc_phase_4, resolve_message
import pytest


@pytest.mark.parametrize(
    "height, width, message_len",
    [(256, 10, 256), (100, 70, 70), (40, 50, 40), (256, 10, 128), (100, 70, 35), (40, 50, 20)]
)
def test_darc(height, width, message_len):
    alphabet = SubstitutionKey.init_matrix(width, height, 255)

    server_keys, mutual_keys, client_keys = generate_keys(height, width)
    original_message = OuterKey.init_matrix(height).matrix[0]
    original_message = original_message[:message_len]

    (alphabet_phase1, medium_phase1) = darc_phase1(alphabet, server_keys, mutual_keys)

    (alphabet_phase2, medium_phase2) = darc_phase2(alphabet_phase1, medium_phase1, client_keys, mutual_keys)
    (phase3_alphabet, phase3_medium) = darc_phase3(alphabet_phase2, medium_phase2, client_keys, mutual_keys)
    darc_message = merge_message(phase3_alphabet, phase3_medium, client_keys, mutual_keys, original_message)

    translated_alphabet, unwrapped_message = darc_phase_4(alphabet, darc_message, server_keys, mutual_keys)

    message = resolve_message(translated_alphabet, unwrapped_message)
    assert original_message == message


@pytest.mark.skip
@pytest.mark.parametrize(
    "height, width, message_len",
    [(256, 10, 256), (100, 70, 70), (40, 50, 40), (256, 10, 128), (100, 70, 35), (40, 50, 20)]
)
def test_darc_no_client(height, width, message_len):
    alphabet = SubstitutionKey.init_matrix(width, height, 255)

    server_keys, mutual_keys, client_keys = generate_keys(height, width)

    original_message = OuterKey.init_matrix(height).matrix[0]
    original_message = original_message[:message_len]
    (alphabet_phase1, medium_phase1) = darc_phase1(alphabet, server_keys, mutual_keys)

    (phase3_alphabet, phase3_medium) = darc_phase_2_3(alphabet_phase1, medium_phase1, mutual_keys)
    darc_message = merge_message_no_client(phase3_alphabet, phase3_medium, mutual_keys, original_message)

    translated_alphabet, unwrapped_message = darc_phase_4(alphabet, darc_message, server_keys, mutual_keys)
    message = resolve_message(translated_alphabet, unwrapped_message)
    assert original_message == message