import asyncio
import gzip
import logging
import subprocess
from pathlib import Path
from typing import Any, BinaryIO, Mapping, Tuple, cast

import pkg_resources
from aiodocker.docker import Docker
from aiodocker.exceptions import DockerError

from ai.backend.common.logging import BraceStyleAdapter

from ..exception import InitializationError
from ..utils import closing_async, get_arch_name, update_nested_dict

log = BraceStyleAdapter(logging.getLogger(__name__))


class PersistentServiceContainer:
    def __init__(
        self,
        image_ref: str,
        container_config: Mapping[str, Any],
        *,
        name: str = None,
    ) -> None:
        self.image_ref = image_ref
        arch = get_arch_name()
        default_container_name = image_ref.split(":")[0].rsplit("/", maxsplit=1)[-1]
        if name is None:
            self.container_name = default_container_name
        else:
            self.container_name = name
        self.container_config = container_config
        self.img_version = int(
            Path(
                pkg_resources.resource_filename(
                    "ai.backend.agent.docker",
                    f"{default_container_name}.version.txt",
                )
            ).read_text()
        )
        self.img_path = Path(
            pkg_resources.resource_filename(
                "ai.backend.agent.docker",
                f"{default_container_name}.img.{arch}.tar.gz",
            )
        )

    async def get_container_version_and_status(self) -> Tuple[int, bool]:
        async with closing_async(Docker()) as docker:
            try:
                c = docker.containers.container(self.container_name)
                await c.show()
            except DockerError as e:
                if e.status == 404:
                    return 0, False
                else:
                    raise
        if c["Config"].get("Labels", {}).get("ai.backend.system", "0") != "1":
            raise RuntimeError(
                f"An existing container named \"{c['Name'].lstrip('/')}\" is not a system container "
                f"spawned by Backend.AI. Please check and remove it."
            )
        return (
            int(c["Config"].get("Labels", {}).get("ai.backend.version", "0")),
            c["State"]["Status"].lower() == "running",
        )

    async def get_image_version(self) -> int:
        async with closing_async(Docker()) as docker:
            try:
                img = await docker.images.inspect(self.image_ref)
            except DockerError as e:
                if e.status == 404:
                    return 0
                else:
                    raise
        return int((img["Config"].get("Labels") or {}).get("ai.backend.version", "0"))

    async def ensure_running_latest(self) -> None:
        image_version = await self.get_image_version()
        if image_version == 0:
            log.info("PersistentServiceContainer({}): installing...", self.image_ref)
            await self.install_latest()
        elif image_version < self.img_version:
            log.info(
                "PersistentServiceContainer({}): upgrading (v{} -> v{})",
                self.image_ref,
                image_version,
                self.img_version,
            )
            await self.install_latest()
        container_version, is_running = await self.get_container_version_and_status()
        if container_version == 0 or image_version != container_version or not is_running:
            log.info("PersistentServiceContainer({}): recreating...", self.image_ref)
            await self.recreate()
        if not is_running:
            log.info("PersistentServiceContainer({}): starting...", self.image_ref)
            await self.start()

    async def install_latest(self) -> None:
        with gzip.open(self.img_path, "rb") as reader:
            proc = await asyncio.create_subprocess_exec(
                *["docker", "load"],
                stdin=cast(BinaryIO, reader),
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
            if await proc.wait() != 0:
                stderr = b"(unavailable)"
                if proc.stderr is not None:
                    stderr = await proc.stderr.read()
                raise RuntimeError(
                    "loading the image has failed!",
                    self.image_ref,
                    proc.returncode,
                    stderr,
                )

    async def recreate(self) -> None:
        async with closing_async(Docker()) as docker:
            try:
                c = docker.containers.container(self.container_name)
                await c.stop()
                await c.delete(force=True)
            except DockerError as e:
                if e.status == 409 and "is not running" in e.message:
                    pass
                elif e.status == 404:
                    pass
                else:
                    raise
            container_config = {
                "Image": self.image_ref,
                "Tty": True,
                "Privileged": False,
                "AttachStdin": False,
                "AttachStdout": False,
                "AttachStderr": False,
                "HostConfig": {
                    "Init": True,
                    "RestartPolicy": {
                        "Name": "unless-stopped",  # make it persistent
                        "MaximumRetryCount": 0,
                    },
                },
            }
            update_nested_dict(container_config, self.container_config)
            try:
                await docker.containers.create(config=container_config, name=self.container_name)
            except DockerError as e:
                err_msg = e.args[1].get("message", "")
                if (
                    e.args[0] == 400
                    and "bind source path does not exist" in err_msg
                    and "/tmp/backend.ai/ipc" in err_msg
                ):
                    raise InitializationError(
                        f"Could not create persistent service container '{self.container_name}' "
                        f"because it cannot access /tmp/backend.ai/ipc directory. "
                        f"This may occur when Docker is installed with Snap or the agent is configured "
                        f"to use a private tmp directory. "
                        f"To resolve, explicitly configure the 'ipc-base-path' option in agent.toml to "
                        f"indicate a directory under $HOME or a non-virtualized directory.",
                    )
                else:
                    raise

    async def start(self) -> None:
        async with closing_async(Docker()) as docker:
            c = docker.containers.container(self.container_name)
            await c.start()
