#!/usr/bin/env python3


import enum
import docker
from typing import Optional


class CPUField(enum.Enum):
    CPU_QUOTA="CpuQuota"
    CPU_PERIOD="CpuPeriod"
    CPUSET_CPUS="CpusetCpus"


def get_docker_client() -> Optional[docker.DockerClient]:
    """
    Create a docker client object
    """
    try:
        client = docker.DockerClient(
            base_url="unix://var/run/docker.sock",
        )
    except Exception as e:
        print(f"Cloud not create docker client: {e}")
        return None
    else:
        return client


def get_container(client:docker.DockerClient, name_or_id:str) -> Optional[docker.models.containers.Container]:
    """
    Get container
    """
    try:
        container = client.containers.get(name_or_id)
    except docker.errors.NotFound:
        return None
    else:
        return container


def get_container_host_config_configuration(container:docker.models.containers.Container) -> dict:
    """
    Get container CPU configuration
    """
    return container.attrs["HostConfig"]


def limit_cpu_consumption(container:docker.models.containers.Container, conf:dict) -> bool:
    """
    Limit the CPU resources to be used
    """
    try:
        container.update(**conf)
    except docker.errors.APIError:
        return False
    else:
        return True
    finally:
        pass


if __name__ == "__main__":
    client = get_docker_client()
    container = get_container(client, "container_name")

    assert container is not None

    host_config = get_container_host_config_configuration(container)
    print(f"{CPUField.CPU_QUOTA.value}: {host_config[CPUField.CPU_QUOTA.value]}")
    print(f"{CPUField.CPU_PERIOD.value}: {host_config[CPUField.CPU_PERIOD.value]}")
    print(f"{CPUField.CPUSET_CPUS.value}: {host_config[CPUField.CPUSET_CPUS.value]}")

    cpu_new_conf = {
        "cpu_quota": 100000,
        "cpu_period": 100000,
        "cpuset_cpus": "0,3"
    }

    limit_cpu_consumption(container, cpu_new_conf)
    container = get_container(client, "container_name")
    host_config = get_container_host_config_configuration(container)
    print(f"{CPUField.CPU_QUOTA.value}: {host_config[CPUField.CPU_QUOTA.value]}")
    print(f"{CPUField.CPU_PERIOD.value}: {host_config[CPUField.CPU_PERIOD.value]}")
    print(f"{CPUField.CPUSET_CPUS.value}: {host_config[CPUField.CPUSET_CPUS.value]}")