import requests
import docker
import os

from ..logger import LoggerFactory

manager_logger = LoggerFactory.get_logger("CustomModuleManager")

def change_module_status(
    module_id: int,
    status: str
):
    _ = requests.post(
        f"http://localhost:{int(os.environ.get('PORT', 5010))}/api/v1.0/interact_with_custom_modules", 
        data = {"request_type": "change_status", "model_id": module_id, "status": status}
    )

def build_image(
    client: docker.DockerClient,
    path: str, 
    dockerfile: str,
    tag: str, 
    buildargs: dict
):
    """Сборка docker контейнеров"""
    _, logs = client.images.build(
        path = path,
        dockerfile = dockerfile,
        tag = tag,
        buildargs = buildargs,
        decode = True
    )
    for chunk in logs:
        if 'stream' in chunk:
            manager_logger.info(chunk['stream'].strip())

def start_container(
    module_id: int,
    image_name: str, 
    build_args: dict,
    container_name: str,
    network_name: str,
    env_vars: dict,
    gpus = True,
    detach = True
):
    """Запуск docker контейнеров"""
    # инициализируем клиента
    client = docker.from_env()

    # Соберем контейнер
    try:
        change_module_status(module_id, status = "работает")
        build_image(client, **build_args)
        manager_logger.info(f"Контейнер '{container_name}' успешно собран")
    except Exception as e:
        change_module_status(module_id, status = "не создан")
        manager_logger.error(f"Ошибка при сборке контейнера: {e}")
        return

    # Укажем устройства исполнения
    device_requests = [docker.types.DeviceRequest(count=-1, capabilities=[['gpu']])] if gpus else None

    # Запустим контейнер
    try:
        container = client.containers.run(
            image = image_name,
            name = container_name,
            network = client.networks.get(network_name),
            environment = env_vars,
            device_requests=device_requests,
            detach = detach,
        )
        manager_logger.info(f"Контейнер '{container_name}' с ID: {container.short_id} запустился")
    except Exception as e:
        change_module_status(module_id, status = "остановлен")
        manager_logger.error(f"Ошибка при запуске контейнера: {e}")

def stop_container(
    module_id: int,
    name: str
):
    """Удаление docker контейнеров по имени"""
    # инициализируем клиента
    client = docker.from_env()

    # Остановим контейнер
    try:
        container = client.containers.get(name)
        container.stop()
        container.remove()
        manager_logger.info(f"Контейнер '{name}' был остановлен")
        change_module_status(module_id, status = "остановлен")
    except docker.errors.NotFound:
        manager_logger.warning(f"Контейнер '{name}' не найден")