from contextlib import contextmanager
from typing import Generator, TypeVar, Union
from uuid import uuid1

from docker import DockerClient
from docker.models.containers import Container
from docker.models.networks import Network

from yellowbox.containers import get_aliases, is_removed
from yellowbox.subclasses import ContainerService

__all__ = ['temp_network', 'anonymous_network', 'connect', 'disconnecting']

_T = TypeVar("_T")
_NT = TypeVar("_NT", bound=Network)


def anonymous_network(client: DockerClient, *args, **kwargs) -> Network:
    name = f"yellowbox-{uuid1()}"
    return client.networks.create(name, *args, **kwargs)


@contextmanager
def temp_network(client: DockerClient, name=None, *args, **kwargs):
    """Context manager for creating a temporary Docker network

    Network will be automatically removed upon context manager completion.

    Example:
        >>> with docker_client() as client, temp_network(client) as network:
        ...   print(network.name)
        ...
        yellowbox-...

    Args:
        client: Docker client to create network with.
        name: Network name. Defaults to a randomly generated name, with the
        prefix "yellowbox".
        *args, **kwargs: Extra arguments for docker network creation.

    Returns:
        Context manager for the newly created network
    """
    if name:
        network = client.networks.create(name, *args, **kwargs)
    else:
        network = anonymous_network(client, *args, **kwargs)
    with disconnecting(network, remove=True) as network:
        yield network


@contextmanager
def connect(network: Network, obj: Union[Container, ContainerService], **kwargs):
    """Temporarily connect a container or yellow service into a network.

    Args:
        network: Docker network.
        obj: Container or YellowService to connect to the network.

    Returns:
        Context manager for handling the connection.
    """
    if isinstance(obj, ContainerService):
        ret = obj.connect(network, **kwargs)  # type: ignore[call-arg]
    else:
        network.connect(obj, **kwargs)
        obj.reload()
        ret = get_aliases(obj, network)
    try:
        yield ret
    finally:
        if isinstance(obj, ContainerService):
            obj.disconnect(network)
        elif not is_removed(obj):
            network.disconnect(obj)


@contextmanager
def disconnecting(network: _NT, *, remove: bool = False) -> Generator[_NT, None, None]:
    """A context manager that disconnects a docker network upon completion.

    Example:
        network = DockerNetwork(...)
        with disconnecting(network):
            ...
        # Network is now disconnected from all containers.

    Args:
        network: Network to disconnect upon completion.
        remove: Whether to remove the network upon completion or not.

    Returns:
        A context manager to be used in a 'with' statement.
    """
    try:
        yield network
    finally:
        network.reload()
        for container in network.containers:
            network.disconnect(container)
        if remove:
            network.remove()
