import asyncio

try:
    from behave.runner import Context  # type: ignore
    from behave.model import Scenario  # type: ignore
    from cattle_grid.testing.features.environment import (
        before_all,  # noqa: F401
        before_scenario as cg_before_scenario,
        after_scenario as cg_after_scenario,
    )


except ImportError:

    class Context:
        herd_task: None = None

    class Scenario: ...


async def cancel_herd(context):
    try:
        async with asyncio.timeout(1):
            context.herd_task.cancel()
            await context.herd_task

    except asyncio.CancelledError as e:
        print(e)


def before_scenario(context: Context, scenario: Scenario):
    context.herd_task = None

    cg_before_scenario(context, scenario)  # type: ignore


def after_scenario(context: Context, scenario: Scenario):
    if context.herd_task:
        try:
            asyncio.get_event_loop().run_until_complete(cancel_herd(context))
        except Exception as e:
            print(e)

    cg_after_scenario(context, scenario)  # type: ignore
