import functools
from io import StringIO
from pathlib import Path
import uuid

import aiotools
from ai.backend.common.asyncio import current_loop
from ai.backend.common.docker import ImageRef
import asyncio
from decimal import Decimal
import hashlib
import logging
import pkg_resources
import random
import shutil
import signal
from typing import (
    Any,
    FrozenSet,
    List,
    Literal,
    Mapping,
    MutableMapping,
    Optional,
    Sequence,
    Tuple,
    Union,
)
from ai.backend.common.logging_utils import BraceStyleAdapter

import cattr
from kubernetes_asyncio import (
    client as kube_client,
    config as kube_config,
)

from .resources import detect_resources
from .kernel import KubernetesKernel
from .kube_object import (
    ConfigMap, HostPathPersistentVolume, KubernetesAbstractVolume,
    KubernetesConfigMapVolume, KubernetesHostPathVolume,
    KubernetesPVCVolume, KubernetesVolumeMount, NFSPersistentVolume,
    PersistentVolumeClaim, Service,
)

from ..agent import ACTIVE_STATUS_SET, AbstractAgent, AbstractKernelCreationContext, ComputerContext
from ..exception import K8sError, UnsupportedResource
from ..kernel import KernelFeatures
from ..resources import (
    AbstractComputePlugin,
    KernelResourceSpec,
    Mount,
    known_slot_types,
)
from ..types import Container, ContainerStatus, Port

from ai.backend.common.etcd import AsyncEtcd
from ai.backend.common.plugin.monitor import ErrorPluginContext, StatsPluginContext
from ai.backend.common.types import (
    AutoPullBehavior,
    ClusterInfo,
    ContainerId,
    DeviceName,
    ImageRegistry,
    KernelCreationConfig,
    KernelId,
    MountPermission,
    MountTypes,
    ResourceSlot,
    SlotName,
    VFolderMount,
    current_resource_slots,
)

log = BraceStyleAdapter(logging.getLogger('ai.backend.agent.kubernetes.agent'))


class KubernetesKernelCreationContext(AbstractKernelCreationContext[KubernetesKernel]):
    config_map_name: str
    deployment_name: str
    scratch_dir: Path
    work_dir: Path
    config_dir: Path
    internal_mounts: List[Mount] = []
    static_pvc_name: str
    workers: Mapping[str, Mapping[str, str]]
    config_maps: List[ConfigMap]

    volume_mounts: List[KubernetesVolumeMount]
    volumes: List[KubernetesAbstractVolume]

    def __init__(
        self,
        kernel_id: KernelId,
        kernel_config: KernelCreationConfig,
        local_config: Mapping[str, Any],
        computers: MutableMapping[str, ComputerContext],
        workers: Mapping[str, Mapping[str, str]],
        static_pvc_name: str,
        restarting: bool = False,
    ) -> None:
        super().__init__(kernel_id, kernel_config, local_config, computers, restarting=restarting)
        scratch_dir = (self.local_config['container']['scratch-root'] / str(kernel_id)).resolve()

        self.scratch_dir = scratch_dir
        self.work_dir = scratch_dir / 'work'
        self.config_dir = scratch_dir / 'config'
        self.static_pvc_name = static_pvc_name
        self.workers = workers

        self.volume_mounts = []
        self.volumes = [
            KubernetesPVCVolume(
                name=f'kernel-{self.kernel_id}-scratches',
                persistentVolumeClaim={
                    'claimName': self.static_pvc_name,
                },
            ),
        ]

        self.config_maps = []

    async def get_extra_envs(self) -> Mapping[str, str]:
        return {}

    async def prepare_resource_spec(self) -> Tuple[KernelResourceSpec, Optional[Mapping[str, Any]]]:
        loop = current_loop()
        if self.restarting:
            await kube_config.load_kube_config()

            def _kernel_resource_spec_read():
                with open((self.config_dir / 'resource.txt').resolve(), 'r') as f:
                    resource_spec = KernelResourceSpec.read_from_file(f)
                return resource_spec

            resource_spec = await loop.run_in_executor(None, _kernel_resource_spec_read)
            resource_opts = None
        else:
            slots = ResourceSlot.from_json(self.kernel_config['resource_slots'])
            # Ensure that we have intrinsic slots.
            assert SlotName('cpu') in slots
            assert SlotName('mem') in slots
            # accept unknown slot type with zero values
            # but reject if they have non-zero values.
            for st, sv in slots.items():
                if st not in known_slot_types and sv != Decimal(0):
                    raise UnsupportedResource(st)
            # sanitize the slots
            current_resource_slots.set(known_slot_types)
            slots = slots.normalize_slots(ignore_unknown=True)
            resource_spec = KernelResourceSpec(
                container_id='',
                allocations={},
                slots={**slots},  # copy
                mounts=[],
                scratch_disk_size=0,  # TODO: implement (#70)
            )
            resource_opts = self.kernel_config.get('resource_opts', {})
        return resource_spec, resource_opts

    async def prepare_scratch(self) -> None:
        loop = current_loop()
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()

        # Unlike Docker, static files will be mounted directly to blank folder
        # Check if NFS PVC for static files exists and bound
        nfs_pvc = await core_api.list_persistent_volume_claim_for_all_namespaces(
            label_selector='backend.ai/backend-ai-scratch-volume',
        )
        if len(nfs_pvc.items) == 0:
            raise K8sError('No PVC for backend.ai static files')
        pvc = nfs_pvc.items[0]
        if pvc.status.phase != 'Bound':
            raise K8sError('PVC not Bound')
        self.static_pvc_name = pvc.metadata.name

        def _create_scratch_dirs():
            self.work_dir.resolve().mkdir(parents=True, exist_ok=True)
            self.config_dir.resolve().mkdir(parents=True, exist_ok=True)

        # Mount scratch directory as PV
        # Config files can be mounted via ConfigMap
        await loop.run_in_executor(None, _create_scratch_dirs)

        if not self.restarting:
            # Since these files are bind-mounted inside a bind-mounted directory,
            # we need to touch them first to avoid their "ghost" files are created
            # as root in the host-side filesystem, which prevents deletion of scratch
            # directories when the agent is running as non-root.
            def _clone_dotfiles():
                jupyter_custom_css_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', 'jupyter-custom.css'))
                logo_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', 'logo.svg'))
                font_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', 'roboto.ttf'))
                font_italic_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', 'roboto-italic.ttf'))
                bashrc_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', '.bashrc'))
                bash_profile_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', '.bash_profile'))
                vimrc_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', '.vimrc'))
                tmux_conf_path = Path(pkg_resources.resource_filename(
                    'ai.backend.runner', '.tmux.conf'))
                jupyter_custom_dir = (self.work_dir / '.jupyter' / 'custom')
                jupyter_custom_dir.mkdir(parents=True, exist_ok=True)
                shutil.copy(jupyter_custom_css_path.resolve(), jupyter_custom_dir / 'custom.css')
                shutil.copy(logo_path.resolve(), jupyter_custom_dir / 'logo.svg')
                shutil.copy(font_path.resolve(), jupyter_custom_dir / 'roboto.ttf')
                shutil.copy(font_italic_path.resolve(), jupyter_custom_dir / 'roboto-italic.ttf')
                shutil.copy(bashrc_path.resolve(), self.work_dir / '.bashrc')
                shutil.copy(bash_profile_path.resolve(), self.work_dir / '.bash_profile')
                shutil.copy(vimrc_path.resolve(), self.work_dir / '.vimrc')
                shutil.copy(tmux_conf_path.resolve(), self.work_dir / '.tmux.conf')

            await loop.run_in_executor(None, _clone_dotfiles)

    async def get_intrinsic_mounts(self) -> Sequence[Mount]:
        mounts: List[Mount] = [
            # Mount scratch directory
            Mount(
                MountTypes.K8S_GENERIC,
                Path(str(self.kernel_id)),
                Path('/home/work'),
                MountPermission.READ_WRITE,
                opts={
                    'name': f'kernel-{self.kernel_id}-scratches',
                },
            ),
        ]

        # TODO: Find way to mount extra volumes

        return mounts

    async def apply_network(self, cluster_info: ClusterInfo) -> None:
        pass

    async def install_ssh_keypair(self, cluster_info: ClusterInfo) -> None:
        sshkey = cluster_info['ssh_keypair']
        if sshkey is None:
            return

        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()

        # Get hash of public key
        enc = hashlib.md5()
        enc.update(sshkey['public_key'].encode('ascii'))
        hash = enc.digest().decode('utf-8')

        try:
            await core_api.read_namespaced_config_map(f'ssh-keypair-{hash}', 'backend-ai')
        except:
            # Keypair not stored on ConfigMap, create one
            cm = ConfigMap('', f'kernel-{self.kernel_id}-ssh-keypair-{hash}')
            cm.put('public', sshkey['public_key'])
            cm.put('private', sshkey['private_key'])

            self.config_maps.append(cm)

        self.process_volumes([
            KubernetesConfigMapVolume(
                name=f'kernel-{self.kernel_id}-ssh-keypair',
                configMap={
                    'name': 'ssh-keypair-hash',
                },
            ),
        ])
        self.process_mounts([
            Mount(
                MountTypes.K8S_GENERIC,
                Path('public'),
                Path('/home/config/ssh/id_cluster.pub'),
                permission=MountPermission.READ_ONLY,
                opts={
                    'name': f'kernel-{self.kernel_id}-ssh-keypair',
                },
            ),
            Mount(
                MountTypes.K8S_GENERIC,
                Path('private'),
                Path('/home/config/ssh/id_cluster.pub'),
                permission=MountPermission.READ_ONLY,
                opts={
                    'name': f'kernel-{self.kernel_id}-ssh-keypair',
                },
            ),
        ])

    async def process_mounts(self, mounts: Sequence[Mount]):
        for i, mount in zip(range(len(mounts)), mounts):
            if mount.type == MountTypes.K8S_GENERIC:
                name = (mount.opts or {})['name']
                self.volume_mounts.append(
                    KubernetesVolumeMount(
                        name=name,
                        mountPath=mount.target.as_posix(),
                        subPath=mount.source.as_posix() if mount.source is not None else None,
                        readOnly=mount.permission == MountPermission.READ_ONLY,
                    ),
                )
            elif mount.type == MountTypes.K8S_HOSTPATH:
                name = (mount.opts or {})['name']
                self.volume_mounts.append(
                    KubernetesVolumeMount(
                        name=name,
                        mountPath=mount.target.as_posix(),
                        subPath=None,
                        readOnly=mount.permission == MountPermission.READ_ONLY,
                    ),
                )
            else:
                log.warn(
                    'Mount {}:{} -> Mount type {} it not supported on K8s Agent. Skipping mount',
                    mount.source,
                    mount.target,
                    mount.type,
                )

    def resolve_krunner_filepath(self, filename: str) -> Path:
        return Path(filename)

    def get_runner_mount(
        self,
        type: MountTypes,
        src: Union[str, Path],
        target: Union[str, Path],
        perm: Literal['ro', 'rw'] = 'ro',
        opts: Mapping[str, Any] = None,
    ) -> Mount:
        return Mount(
            MountTypes.K8S_GENERIC,
            Path(src),
            Path(target),
            MountPermission(perm),
            opts={
                **(opts or {}),
                'name': f'kernel-{self.kernel_id}-scratches',
            },
        )

    async def process_volumes(
        self,
        volumes: List[KubernetesAbstractVolume],
    ) -> None:
        self.volumes += volumes

    async def mount_vfolders(
        self,
        vfolders: Sequence[VFolderMount],
        resource_spec: KernelResourceSpec,
    ) -> None:
        # We can't mount vFolder backed by storage proxy
        for idx, vfolder in enumerate(vfolders):
            if self.internal_data.get('prevent_vfolder_mounts', False):
                # Only allow mount of ".logs" directory to prevent expose
                # internal-only information, such as Docker credentials to user's ".docker" vfolder
                # in image importer kernels.
                if vfolder.name != '.logs':
                    continue
            mount = Mount(
                MountTypes.K8S_HOSTPATH,
                Path(vfolder.host_path),
                Path(vfolder.kernel_path),
                vfolder.mount_perm,
                opts={
                    'name': f'kernel-{self.kernel_id}-hostPath-{idx}',
                },
            )
            self.process_volumes([
                KubernetesHostPathVolume(
                    name=f'kernel-{self.kernel_id}-hostPath-{idx}',
                    hostPath={
                        'path': vfolder.host_path.as_posix(),
                        'type': 'Directory',
                    },
                ),
            ])
            resource_spec.mounts.append(mount)

    async def apply_accelerator_allocation(self, computer, device_alloc) -> None:
        # update_nested_dict(
        #     self.computer_docker_args,
        #     await computer.generate_docker_args(self.docker, device_alloc))
        # TODO: add support for accelerator allocation
        pass

    async def generate_deployment_object(
        self, image: str, environ: Mapping[str, Any],
        ports: List[int], command: List[str],
        labels: Mapping[str, Any] = {},
    ) -> dict:
        return {
            'apiVersion': 'apps/v1',
            'kind': 'Deployment',
            'metadata': {
                'name': f'kernel-{self.kernel_id}',
                'labels': labels,
            },
            'spec': {
                'replicas': 0,
                'selector': {'matchLabels': {'run': f'kernel-{self.kernel_id}'}},
                'template': {
                    'metadata': {'labels': {'run': f'kernel-{self.kernel_id}'}},
                    'spec': {
                        'containers': [{
                            'name': f'kernel-{self.kernel_id}-session',
                            'image': image,
                            'imagePullPolicy': 'IfNotPresent',
                            'command': ['sh', '/opt/kernel/entrypoint.sh'],
                            'args': command,
                            'env': [{'name': k, 'value': v} for k, v in environ.items()],
                            'volumeMounts': [cattr.unstructure(v) for v in self.volume_mounts],
                            'ports': [{'containerPort': x} for x in ports],
                        }],
                        'volumes': [cattr.unstructure(v) for v in self.volumes],
                    },
                },
            },
        }

    async def spawn(
        self,
        resource_spec: KernelResourceSpec,
        resource_opts,
        environ: Mapping[str, str],
        service_ports,
        preopen_ports,
        cmdargs: List[str],
    ) -> KubernetesKernel:
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()
        apps_api = kube_client.AppsV1Api()

        loop = current_loop()
        image_labels = self.kernel_config['image']['labels']
        exposed_ports = [2000, 2001]
        for sport in service_ports:
            exposed_ports.extend(sport['container_ports'])

        if self.restarting:
            pass
        else:
            if bootstrap := self.kernel_config.get('bootstrap_script'):

                def _write_user_bootstrap_script():
                    (self.work_dir / 'bootstrap.sh').write_text(bootstrap)
                    if KernelFeatures.UID_MATCH in self.kernel_features:  # UID Match won't work on K8s
                        # uid = self.local_config['container']['kernel-uid']
                        # gid = self.local_config['container']['kernel-gid']
                        # if os.geteuid() == 0:
                        #     os.chown(self.work_dir / 'bootstrap.sh', uid, gid)
                        pass

                await loop.run_in_executor(None, _write_user_bootstrap_script)

            def _write_config(file_name: str, content: str):
                (self.config_dir / file_name).write_text(content)

            with StringIO() as buf:
                for k, v in environ.items():
                    buf.write(f'{k}={v}\n')
                # accel_envs = self.computer_docker_args.get('Env', [])
                # for env in accel_envs:
                #     buf.write(f'{env}\n')
                await loop.run_in_executor(
                    None, functools.partial(_write_config, 'environ.txt', buf.getvalue()),
                )

            with StringIO() as buf:
                resource_spec.write_to_file(buf)
                for dev_type, device_alloc in resource_spec.allocations.items():
                    computer_self = self.computers[dev_type]
                    kvpairs = \
                        await computer_self.instance.generate_resource_data(device_alloc)
                    for k, v in kvpairs.items():
                        buf.write(f'{k}={v}\n')
                await loop.run_in_executor(
                    None, functools.partial(_write_config, 'resource.txt', buf.getvalue()),
                )

            docker_creds = self.internal_data.get('docker_credentials')
            if docker_creds:
                await loop.run_in_executor(
                    None, functools.partial(_write_config, 'docker-creds.json', docker_creds),
                )

        if keypair := self.internal_data.get('ssh_keypair'):
            for mount in resource_spec.mounts:
                container_path = str(mount).split(':')[1]
                if container_path == '/home/work/.ssh':
                    break
            else:
                pubkey = keypair['public_key'].encode('ascii')
                privkey = keypair['private_key'].encode('ascii')
                ssh_config_map = ConfigMap(self.kernel_id, f'kernel-{self.kernel_id}-ssh-config')
                ssh_config_map.put('authorized_keys', pubkey)
                ssh_config_map.put('id_container', privkey)
                self.process_volumes([
                    KubernetesConfigMapVolume(
                        name='ssh-config',
                        configMap={
                            'name': f'kernel-{self.kernel_id}-ssh-config',
                        },
                    ),
                ])
                self.process_mounts([
                    Mount(
                        MountTypes.K8S_GENERIC,
                        Path('authorized_keys'),
                        Path('/home/work/.ssh/authorized_keys'),
                        opts={
                            'name': 'ssh-config',
                        },
                    ),
                    Mount(
                        MountTypes.K8S_GENERIC,
                        Path('id_container'),
                        Path('/home/work/.ssh/id_container'),
                        opts={
                            'name': 'ssh-config',
                        },
                    ),
                ])

        # higher priority dotfiles are stored last to support overwriting
        for dotfile in self.internal_data.get('dotfiles', []):
            if dotfile['path'].startswith('/'):
                if dotfile['path'].startswith('/home/'):
                    path_arr = dotfile['path'].split('/')
                    file_path: Path = self.scratch_dir / '/'.join(path_arr[2:])
                else:
                    file_path = Path(dotfile['path'])
            else:
                file_path = self.work_dir / dotfile['path']
            file_path.parent.mkdir(parents=True, exist_ok=True)
            await loop.run_in_executor(
                None,
                file_path.write_text,
                dotfile['data'])

        # TODO: Mark shmem feature as unsupported when advertising agent

        # PHASE 4: Run!

        encoded_preopen_ports = ','.join(f'{port_no}:preopen:{port_no}' for port_no in preopen_ports)

        service_port_label = image_labels['ai.backend.service-ports']
        if len(encoded_preopen_ports) > 0:
            service_port_label += f',{encoded_preopen_ports}'

        deployment = await self.generate_deployment_object(
            self.image_ref.canonical,
            environ,
            exposed_ports,
            cmdargs,
            labels={
                'ai.backend.service-ports': service_port_label.replace(':', '-').replace(',', '.'),
            },
        )

        if self.local_config['debug']['log-kernel-config']:
            log.debug('Initial container config: {0}', deployment)

        expose_service = Service(
            str(self.kernel_id),
            f'kernel-{self.kernel_id}-expose',
            [
                (port, f'kernel-{self.kernel_id}-svc-{index}')
                for index, port in zip(range(len(exposed_ports)), exposed_ports)
            ],
        )

        async def rollup(
            functions: List[Tuple[Optional[functools.partial],
            Optional[functools.partial]]],
        ):
            rollback_functions: List[Optional[functools.partial]] = []

            for (rollup_function, future_rollback_function) in functions:
                try:
                    if rollup_function:
                        await rollup_function()
                    rollback_functions.append(future_rollback_function)
                except Exception as e:
                    for rollback_function in rollback_functions[::-1]:
                        if rollback_function:
                            await rollback_function()
                    log.exception('Error while rollup: {}', e)
                    raise

        arguments: List[Tuple[Optional[functools.partial], Optional[functools.partial]]] = []
        node_ports = []

        try:
            expose_service_api_response = \
                await core_api.create_namespaced_service('backend-ai', body=expose_service.to_dict())
        except Exception as e:
            log.exception('Error while rollup: {}', e)
            raise

        node_ports = expose_service_api_response.spec.ports
        arguments.append((
            None,
            functools.partial(core_api.delete_namespaced_service, expose_service.name, 'backend-ai'),
        ))
        for cm in self.config_maps:
            arguments.append((
                functools.partial(
                    core_api.create_namespaced_config_map, 'backend-ai',
                    body=cm.to_dict(),
                ),
                functools.partial(core_api.delete_namespaced_config_map, cm.name, 'backend-ai'),
            ))

        arguments.append((
            functools.partial(
                apps_api.create_namespaced_deployment, 'backend-ai',
                body=deployment, pretty='pretty-example',
            ),
            None,
        ))

        await rollup(arguments)

        assigned_ports: MutableMapping[str, int] = {}
        for port in node_ports:
            assigned_ports[port.port] = port.node_port

        ctnr_host_port_map: MutableMapping[int, int] = {}
        stdin_port = 0
        stdout_port = 0
        for idx, port in enumerate(exposed_ports):
            host_port = assigned_ports[port]

            if port == 2000:     # intrinsic
                repl_in_port = host_port
            elif port == 2001:   # intrinsic
                repl_out_port = host_port
            elif port == 2002:   # legacy
                stdin_port = host_port
            elif port == 2003:   # legacy
                stdout_port = host_port
            else:
                ctnr_host_port_map[port] = host_port
        for sport in service_ports:
            sport['host_ports'] = tuple(
                ctnr_host_port_map[cport] for cport in sport['container_ports']
            )

        target_node_ip = random.choice([x['InternalIP'] for x in self.workers.values()])

        kernel_obj = await KubernetesKernel.new(
            self.kernel_id,
            self.image_ref,
            self.kspec_version,
            agent_config=self.local_config,
            service_ports=service_ports,
            resource_spec=resource_spec,
            data={
                'container_id': '',
                'kernel_host': target_node_ip,
                'repl_in_port': repl_in_port,
                'repl_out_port': repl_out_port,
                'stdin_port': stdin_port,    # legacy
                'stdout_port': stdout_port,  # legacy
                'assigned_ports': assigned_ports,
                # 'domain_socket_proxies': self.domain_socket_proxies,
                'block_service_ports': self.internal_data.get('block_service_ports', False),
            })
        return kernel_obj


class KubernetesAgent(AbstractAgent[KubernetesKernel, KubernetesKernelCreationContext]):
    workers: MutableMapping[str, MutableMapping[str, str]] = {}
    k8s_ptask_group: aiotools.PersistentTaskGroup

    def __init__(
        self,
        etcd: AsyncEtcd,
        local_config: Mapping[str, Any],
        *,
        stats_monitor: StatsPluginContext,
        error_monitor: ErrorPluginContext,
        skip_initial_scan: bool = False,
    ) -> None:
        super().__init__(
            etcd,
            local_config,
            stats_monitor=stats_monitor,
            error_monitor=error_monitor,
            skip_initial_scan=skip_initial_scan,
        )

    async def __ainit__(self) -> None:
        await super().__ainit__()
        ipc_base_path = self.local_config['agent']['ipc-base-path']
        self.agent_sockpath = ipc_base_path / 'container' / f'agent.{self.local_instance_id}.sock'

        await self.check_krunner_pv_status()
        await self.fetch_workers()
        self.k8s_ptask_group = aiotools.PersistentTaskGroup()
        # Socket Relay initialization
        # Agent socket handler initialization
        # K8s event monitor task initialization

    async def check_krunner_pv_status(self):
        capacity = format(self.local_config['container']['scratch-size'], 'g')[:-1]

        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()

        namespaces = await core_api.list_namespace()
        if len(list(filter(lambda ns: ns.metadata.name == 'backend-ai', namespaces.items))) == 0:
            await core_api.create_namespace({
                'apiVersion': 'v1',
                'kind': 'Namespace',
                'metadata': {
                    'name': 'backend-ai',
                },
            })

        pv = await core_api.list_persistent_volume(label_selector='backend.ai/backend-ai-scratch-volume')

        if len(pv.items) == 0:
            # PV does not exists; create one
            if self.local_config['container']['scratch-type'] == 'k8s-nfs':
                new_pv = NFSPersistentVolume(
                    self.local_config['container']['scratch-nfs-address'],
                    'backend-ai-static-pv', capacity,
                )
                new_pv.label(
                    'backend.ai/backend-ai-scratch-volume',
                    self.local_config['container']['scratch-nfs-address'],
                )
                new_pv.options = [
                    x.strip()
                    for x in self.local_config['container']['scratch-nfs-options'].split(',')
                ]
            elif self.local_config['container']['scratch-type'] == 'hostdir':
                new_pv = HostPathPersistentVolume(
                    self.local_config['container']['scratch-root'].as_posix(),
                    'backend-ai-static-pv', capacity,
                )
                new_pv.label('backend.ai/backend-ai-scratch-volume', 'hostPath')
            else:
                raise NotImplementedError(
                    f'Scratch type {self.local_config["container"]["scratch-type"]} is not supported',
                )

            try:
                await core_api.create_persistent_volume(body=new_pv.to_dict())
            except:
                raise

        pvc = await core_api.list_persistent_volume_claim_for_all_namespaces(
            label_selector='backend.ai/backend-ai-scratch-volume',
        )

        if len(pvc.items) == 0:
            # PV does not exists; create one
            new_pvc = PersistentVolumeClaim(
                'backend-ai-static-pvc', 'backend-ai-static-pv', capacity,
            )
            if self.local_config['container']['scratch-type'] == 'k8s-nfs':
                new_pvc.label(
                    'backend.ai/backend-ai-scratch-volume',
                    self.local_config['container']['scratch-nfs-address'],
                )
            else:
                new_pvc.label('backend.ai/backend-ai-scratch-volume', 'hostPath')
            try:
                await core_api.create_namespaced_persistent_volume_claim(
                    'backend-ai', body=new_pvc.to_dict(),
                )
            except Exception as e:
                log.exception('Error: {}', e)
                raise

    async def fetch_workers(self):
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()
        nodes = await core_api.list_node()
        for node in nodes.items:
            # if 'node-role.kubernetes.io/master' in node.metadata.labels.keys():
            #     continue
            self.workers[node.metadata.name] = node.status.capacity
            for addr in node.status.addresses:
                if addr.type == 'InternalIP':
                    self.workers[node.metadata.name]['InternalIP'] = addr.address
                if addr.type == 'ExternalIP':
                    self.workers[node.metadata.name]['ExternalIP'] = addr.address

    async def shutdown(self, stop_signal: signal.Signals):
        # Stop agent socket handler task

        try:
            if self.k8s_ptask_group is not None:
                await self.k8s_ptask_group.shutdown()
            await super().shutdown(stop_signal)
        finally:
            # Stop k8s event monitoring.
            pass

    async def detect_resources(self) -> Tuple[
        Mapping[DeviceName, AbstractComputePlugin],
        Mapping[SlotName, Decimal],
    ]:
        return await detect_resources(self.etcd, self.local_config)

    async def enumerate_containers(
        self,
        status_filter: FrozenSet[ContainerStatus] = ACTIVE_STATUS_SET,
    ) -> Sequence[Tuple[KernelId, Container]]:
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()

        result = []
        fetch_tasks = []
        for deployment in (await core_api.list_namespaced_pod('backend-ai')).items:
            # Additional check to filter out real worker pods only?

            async def _fetch_container_info(pod: Any):
                kernel_id: Union[str, None] = "(unknown)"
                try:
                    kernel_id = await get_kernel_id_from_deployment(pod)
                    if kernel_id is None:
                        return
                    # Is it okay to assume that only one container resides per pod?
                    if pod['status']['containerStatuses'][0]['stats'].keys()[0] in status_filter:
                        result.append(
                            (
                                KernelId(uuid.UUID(kernel_id)),
                                await container_from_pod(pod),
                            ),
                        )
                except asyncio.CancelledError:
                    pass
                except Exception:
                    log.exception(
                        "error while fetching container information (cid:{}, k:{})",
                        pod['metadata']['uid'], kernel_id,
                    )

            fetch_tasks.append(_fetch_container_info(deployment))

        await asyncio.gather(*fetch_tasks, return_exceptions=True)
        return result

    async def scan_images(self) -> Mapping[str, str]:
        # Retrieving image label from registry api is not possible
        return {}

    async def handle_agent_socket(self):
        # TODO: Add support for remote agent socket mechanism
        pass

    async def pull_image(self, image_ref: ImageRef, registry_conf: ImageRegistry) -> None:
        # TODO: Add support for appropriate image pulling mechanism on K8s
        pass

    async def check_image(self, image_ref: ImageRef, image_id: str, auto_pull: AutoPullBehavior) -> bool:
        # TODO: Add support for appropriate image checking mechanism on K8s
        # Just mark all images as 'pulled' since we can't manually initiate image pull on each kube node
        return True

    async def init_kernel_context(
        self,
        kernel_id: KernelId,
        kernel_config: KernelCreationConfig,
        *,
        restarting: bool = False,
    ) -> KubernetesKernelCreationContext:
        return KubernetesKernelCreationContext(
            kernel_id,
            kernel_config,
            self.local_config,
            self.computers,
            self.workers,
            'backend-ai-static-pvc',
            restarting=restarting,
        )

    async def destroy_kernel(self, kernel_id: KernelId, container_id: Optional[ContainerId]) -> None:
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()
        apps_api = kube_client.AppsV1Api()

        async def force_cleanup(reason='self-terminated'):
            await self.send_event('kernel_terminated',
                                  kernel_id, 'self-terminated',
                                  None)
        try:
            kernel = self.kernel_registry[kernel_id]
        except:
            log.warning('_destroy_kernel({0}) kernel missing (already dead?)',
                        kernel_id)
            await asyncio.shield(self.k8s_ptask_group.create_task(force_cleanup()))
            return None
        deployment_name = kernel['deployment_name']
        try:
            await core_api.delete_namespaced_service(f'{deployment_name}-service', 'backend-ai')
            await core_api.delete_namespaced_service(f'{deployment_name}-nodeport', 'backend-ai')
            await apps_api.delete_namespaced_deployment(f'{deployment_name}', 'backend-ai')
        except:
            log.warning('_destroy({0}) kernel missing (already dead?)', kernel_id)

    async def clean_kernel(
        self,
        kernel_id: KernelId,
        container_id: Optional[ContainerId],
        restarting: bool,
    ) -> None:
        loop = current_loop()
        if not restarting:
            scratch_dir = self.local_config['container']['scratch-root'] / str(kernel_id)
            await loop.run_in_executor(None, shutil.rmtree, str(scratch_dir))

    async def create_overlay_network(self, network_name: str) -> None:
        return await super().create_overlay_network(network_name)

    async def destroy_overlay_network(self, network_name: str) -> None:
        return await super().destroy_overlay_network(network_name)

    async def create_local_network(self, network_name: str) -> None:
        return await super().create_local_network(network_name)

    async def destroy_local_network(self, network_name: str) -> None:
        return await super().destroy_local_network(network_name)

    async def restart_kernel__load_config(
        self,
        kernel_id: KernelId,
        name: str,
    ) -> bytes:
        loop = current_loop()
        scratch_dir = (self.local_config['container']['scratch-root'] / str(kernel_id)).resolve()
        config_dir = scratch_dir / 'config'
        return await loop.run_in_executor(
            None,
            (config_dir / name).read_bytes,
        )

    async def restart_kernel__store_config(
        self,
        kernel_id: KernelId,
        name: str,
        data: bytes,
    ) -> None:
        loop = current_loop()
        scratch_dir = (self.local_config['container']['scratch-root'] / str(kernel_id)).resolve()
        config_dir = scratch_dir / 'config'
        return await loop.run_in_executor(
            None,
            (config_dir / name).write_bytes,
            data,
        )


async def get_kernel_id_from_deployment(pod: Any) -> Optional[str]:
    # TODO: create function which extracts kernel id from pod object
    return pod.get('metadata', {}).get('name')


async def container_from_pod(pod: Any) -> Container:
    status: ContainerStatus = ContainerStatus.RUNNING
    phase = pod['status']['phase']
    if phase == 'Pending' or phase == 'Running':
        status = ContainerStatus.RUNNING
    elif phase == 'Succeeded':
        status = ContainerStatus.EXITED
    elif phase == 'Failed' or phase == 'Unknown':
        status = ContainerStatus.DEAD

    # TODO: Create Container object from K8s Pod definition
    return Container(
        id=ContainerId(''),
        status=status,
        image=pod['spec']['containers'][0]['image'],
        labels=pod['metadata']['labels'],
        ports=[Port(
            port['host_ip'], port['container_port'], port['host_port'],
        ) for port in pod['spec']['containers'][0]['ports']],
        backend_obj=pod,
    )
