from typing import Set

import pytest
from dapd_db_schema.schemas.uow_meta import UnitOfWorkStatusEnum

from dapd_helpers.fsm import UnitOfWorkStatusFSM


@pytest.mark.parametrize(
    'to_state,expected', [
        (
            UnitOfWorkStatusEnum.EXPORT_SCHEDULED,
            set(),
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            {UnitOfWorkStatusEnum.EXPORT_SCHEDULED},
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_COMPLETED,
            {
                UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
                UnitOfWorkStatusEnum.IMPORT_FAILED,
            },
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_FAILED,
            {
                UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
                UnitOfWorkStatusEnum.EXPORT_SCHEDULED,
            },
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            {UnitOfWorkStatusEnum.EXPORT_COMPLETED},
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_IN_PROGRESS,
            {UnitOfWorkStatusEnum.IMPORT_SCHEDULED},
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_COMPLETED,
            {UnitOfWorkStatusEnum.IMPORT_IN_PROGRESS},
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_FAILED,
            {
                UnitOfWorkStatusEnum.IMPORT_IN_PROGRESS,
                UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            },
        ),
    ]
)
def test_fsm__get_previous_states(
    fsm: UnitOfWorkStatusFSM, to_state: UnitOfWorkStatusEnum, expected: Set[UnitOfWorkStatusEnum]
):
    got = fsm.get_previous_states(to_state)
    assert got == expected


@pytest.mark.parametrize(
    'from_state, expected', [
        (
            UnitOfWorkStatusEnum.EXPORT_COMPLETED,
            {UnitOfWorkStatusEnum.IMPORT_SCHEDULED},
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            {UnitOfWorkStatusEnum.EXPORT_COMPLETED, UnitOfWorkStatusEnum.EXPORT_FAILED},
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_SCHEDULED,
            {
                UnitOfWorkStatusEnum.EXPORT_FAILED,
                UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            },
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_FAILED,
            {UnitOfWorkStatusEnum.EXPORT_COMPLETED},
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_IN_PROGRESS,
            {
                UnitOfWorkStatusEnum.IMPORT_COMPLETED,
                UnitOfWorkStatusEnum.IMPORT_FAILED,
            },
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            {
                UnitOfWorkStatusEnum.IMPORT_FAILED,
                UnitOfWorkStatusEnum.IMPORT_IN_PROGRESS,
            },
        ),
    ]
)
def test_fsm__get_next_states(
    fsm: UnitOfWorkStatusFSM, from_state: UnitOfWorkStatusEnum, expected: Set[UnitOfWorkStatusEnum]
):
    got = fsm.get_next_states(from_state)
    assert got == expected


@pytest.mark.parametrize(
    'from_state, to_state, expected', [
        (
            UnitOfWorkStatusEnum.IMPORT_IN_PROGRESS,
            UnitOfWorkStatusEnum.IMPORT_FAILED,
            True,
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            UnitOfWorkStatusEnum.IMPORT_FAILED,
            True,
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_COMPLETED,
            UnitOfWorkStatusEnum.IMPORT_FAILED,
            False,
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_COMPLETED,
            UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            False,
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_COMPLETED,
            UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            False,
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_COMPLETED,
            UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            True,
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            True,
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            UnitOfWorkStatusEnum.EXPORT_FAILED,
            True,
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            UnitOfWorkStatusEnum.EXPORT_COMPLETED,
            True,
        ),
        (
            UnitOfWorkStatusEnum.IMPORT_FAILED,
            UnitOfWorkStatusEnum.EXPORT_COMPLETED,
            True,
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_IN_PROGRESS,
            UnitOfWorkStatusEnum.EXPORT_SCHEDULED,
            False,
        ),
        (
            UnitOfWorkStatusEnum.EXPORT_COMPLETED,
            UnitOfWorkStatusEnum.IMPORT_SCHEDULED,
            True,
        ),
    ]
)
def test_fsm__can_make_transition(
    from_state: UnitOfWorkStatusEnum, to_state: UnitOfWorkStatusEnum, expected: bool,
    fsm: UnitOfWorkStatusFSM
):
    got = fsm.can_make_transition(to_state, from_state)
    assert got == expected