# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/init.ipynb (unless otherwise specified).

__all__ = ['check_if_port_is_bound', 'get_random_port_numbers', 'logger', 'generate_port_pairings', 'lookup_host',
           'lookup_host_port', 'lookup_socket_name', 'NamedPorts', 'set_port_name', 'block_for_startup',
           'block_for_command', 'run_container']

# Cell
import socket
import atexit
import docker
import uuid
import logging
import random
from numbers import Number

import requests
from retrying import retry
logger = logging.getLogger()

def check_if_port_is_bound(port_number):
    '''
    Attempts to bind a socket.  One of two things will happen
    - Nothing.  The socket is open
    - OSError.  The socket is in use
    '''
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(('', port_number))
    s.close()

@retry(stop_max_attempt_number=50)
def get_random_port_numbers(start, end, quantity=1):
    '''
    gives you sequential port numbers.
    Randomly tries sequences of ports until it finds an open sequence.
    '''
    starting_number = random.randint(start, end-quantity)
    port_numbers = list(range(starting_number, starting_number+quantity))
    for port_number in port_numbers:
        check_if_port_is_bound(port_number)
    return port_numbers


# Cell
def generate_port_pairings(ports, min_port, max_port):
    '''
    Create a dictionary that is compatible with the docker interface
    that describes the ports
    {container_port: host_port}  e.g.
    {2222:3333} exposes port 2222 in the container to the host's 3333
    https://docker-py.readthedocs.io/en/stable/containers.html
    '''
    port_pairs = {}
    ports = [] if ports is None else ports

    if isinstance(ports, Number):
        quantity = 1
        ports = [ports]
    else:
        quantity = len(ports)

    host_ports = get_random_port_numbers(min_port, max_port, quantity=quantity)
    for index, container_port in enumerate(ports):
        host_port = host_ports[index]
        port_pairs[container_port] = host_port

    return port_pairs


# Cell
def lookup_host(container, container_port, protocol='tcp'):
    port_str = '{}/{}'.format(container_port, protocol)
    return container.ports[port_str][0]

def lookup_host_port(container, container_port, protocol='tcp'):
    hostport = lookup_host(container, container_port, protocol='tcp')['HostPort']
    return int(hostport)

def lookup_socket_name(container, container_port, protocol='tcp'):
    hostport = lookup_host(container, container_port, protocol=protocol)
    return '{HostIp}:{HostPort}'.format(**hostport)

# Cell
class NamedPorts:
    def __init__(self):
        self.mapping = {}
    def add_port(self, name, port_num):
        self.mapping[name] = port_num
        #setattr(self, name, port_num)
    def __getattr__(self, name):
        try:
            return self.mapping[name]
        except KeyError:
            raise AttributeError

def set_port_name(container, container_port, name):
    hostport = lookup_host_port(container, container_port)
    try:
        named_ports = container.port
    except AttributeError:
        named_ports = NamedPorts()

    #set the name of the named port
    named_ports.add_port(name, hostport)
    container.port = named_ports
    return container


# Cell

@retry(wait_fixed=1000, stop_max_attempt_number=30)
def block_for_startup(container, container_port):
    '''
    Keep trying http get until the container is available and serves content
    '''
    host = lookup_host(container, container_port)
    logger.debug("Waiting for avilability of {host}".format(host=host))
    resp = requests.get("http://{HostIp}:{HostPort}".format(**host))
    return resp

@retry(wait_fixed=1000, stop_max_attempt_number=30)
def block_for_command(container, cmd, exit_code=0, **kwargs):
    '''
    Keep trying the command in the docker container until you get the desired exit code
    '''
    resp = container.exec_run(cmd, **kwargs)
    assert resp.exit_code==exit_code

def run_container(image, container_name=None,
                  ports=None, min_port=6000, max_port=10000,
                  persist=False, docker_client=None,
                  block_until_port_available=None,
                  environment=None,
                  **kwargs):
    '''
    ports is either
        - a port number
        - list of port numbers you need inside the container.  The ports on the host will be randomized
        - a dictionary of the port number mappings you want in the form of container:host
    environment: dict
        key:value pairs of environment variables for the container

    block_until_port_available: int
        try to requests.get this port until it is successful or times out multiple times

    any additional kwargs will be passed directly into the docker.run command
    https://docker-py.readthedocs.io/en/stable/containers.html
    '''
    client = docker_client if docker_client is not None else docker.from_env()
    container_name = container_name or uuid.uuid4().hex
    try:
        container = client.containers.get(container_name)
    except docker.errors.NotFound:
        docker_run_kwargs = dict(image=image,
                             remove=True, detach=True, auto_remove=True,
                             name=container_name)
        if environment:
            docker_run_kwargs['environment'] = environment

        if isinstance(ports, dict):
            docker_run_kwargs['ports'] = ports
        else:
            port_pairs = generate_port_pairings(ports, min_port, max_port)
            docker_run_kwargs['ports'] = port_pairs

        docker_run_kwargs.update(kwargs)
        container = client.containers.run(**docker_run_kwargs)

    if not persist:
        atexit.register(container.stop)

    container.reload()

    if block_until_port_available:
        block_for_startup(container, container_port=block_until_port_available)
    return container