from __future__ import annotations

import asyncio
import functools
import hashlib
import logging
import os
import random
import shutil
import signal
import sys
import uuid
from collections.abc import Coroutine, Mapping, MutableMapping, Sequence
from decimal import Decimal
from importlib.resources import files
from io import StringIO
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    override,
)

import aiotools
import cattr
from kubernetes.client.models import V1Service, V1ServicePort
from kubernetes_asyncio import client as kube_client
from kubernetes_asyncio import config as kube_config

from ai.backend.agent.agent import (
    ACTIVE_STATUS_SET,
    AbstractAgent,
    AbstractKernelCreationContext,
    AgentClass,
    ScanImagesResult,
)
from ai.backend.agent.config.unified import AgentUnifiedConfig, ScratchType
from ai.backend.agent.etcd import AgentEtcdClientView
from ai.backend.agent.exception import K8sError, UnsupportedResource
from ai.backend.agent.kernel import AbstractKernel, KernelRegistry
from ai.backend.agent.kernel_registry.recovery.kubernetes_recovery import (
    KubernetesKernelRegistryRecovery,
    KubernetesKernelRegistryRecoveryArgs,
)
from ai.backend.agent.kernel_registry.writer.types import KernelRegistrySaveMetadata
from ai.backend.agent.resources import (
    AbstractComputePlugin,
    ComputerContext,
    KernelResourceSpec,
    Mount,
    known_slot_types,
)
from ai.backend.agent.types import Container, KernelOwnershipData, MountInfo, Port
from ai.backend.common.asyncio import current_loop
from ai.backend.common.docker import ImageRef, KernelFeatures
from ai.backend.common.dto.agent.response import PurgeImagesResp
from ai.backend.common.dto.manager.rpc_request import PurgeImagesReq
from ai.backend.common.events.dispatcher import EventProducer
from ai.backend.common.plugin.monitor import ErrorPluginContext, StatsPluginContext
from ai.backend.common.types import (
    AutoPullBehavior,
    ClusterInfo,
    ClusterSSHPortMapping,
    ContainerId,
    ContainerStatus,
    DeviceId,
    DeviceName,
    ImageConfig,
    ImageRegistry,
    KernelCreationConfig,
    KernelId,
    MountPermission,
    MountTypes,
    ResourceSlot,
    Sentinel,
    SlotName,
    VFolderMount,
    current_resource_slots,
)
from ai.backend.logging import BraceStyleAdapter

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

if TYPE_CHECKING:
    from ai.backend.common.auth import PublicKey

log = BraceStyleAdapter(logging.getLogger(__spec__.name))


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]
    agent_sockpath: Path
    volume_mounts: list[KubernetesVolumeMount]
    volumes: list[KubernetesAbstractVolume]

    def __init__(
        self,
        ownership_data: KernelOwnershipData,
        event_producer: EventProducer,
        kernel_image: ImageRef,
        kernel_config: KernelCreationConfig,
        distro: str,
        local_config: AgentUnifiedConfig,
        agent_sockpath: Path,
        computers: Mapping[DeviceName, ComputerContext],
        workers: Mapping[str, Mapping[str, str]],
        static_pvc_name: str,
        restarting: bool = False,
    ) -> None:
        super().__init__(
            ownership_data,
            event_producer,
            kernel_image,
            kernel_config,
            distro,
            local_config,
            computers,
            restarting=restarting,
        )
        kernel_id = ownership_data.kernel_id
        scratch_dir = (self.local_config.container.scratch_root / str(kernel_id)).resolve()
        rel_scratch_dir = Path(str(kernel_id))  # need relative path for nfs mount

        self.scratch_dir = scratch_dir
        self.rel_scratch_dir = rel_scratch_dir
        self.work_dir = scratch_dir / "work"
        self.config_dir = scratch_dir / "config"
        self.rel_work_dir = self.rel_scratch_dir / "work"
        self.rel_config_dir = self.rel_scratch_dir / "config"

        self.static_pvc_name = static_pvc_name
        self.workers = workers
        self.agent_sockpath = agent_sockpath

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

        self.config_maps = []

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

    @override
    async def prepare_resource_spec(self) -> tuple[KernelResourceSpec, Mapping[str, Any] | None]:
        loop = current_loop()
        if self.restarting:
            await kube_config.load_kube_config()

            def _kernel_resource_spec_read() -> KernelResourceSpec:
                resource_file = (self.config_dir / "resource.txt").resolve()
                with resource_file.open() as f:
                    return KernelResourceSpec.read_from_file(f)

            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.
            if SlotName("cpu") not in slots:
                raise UnsupportedResource("cpu slot is required")
            if SlotName("mem") not in slots:
                raise UnsupportedResource("mem slot is required")
            # 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(
                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

    @override
    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() -> None:
            self.work_dir.resolve().mkdir(parents=True, exist_ok=True)
            self.work_dir.chmod(0o755)
            self.config_dir.resolve().mkdir(parents=True, exist_ok=True)
            self.config_dir.chmod(0o755)

        # 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() -> None:
                jupyter_custom_css_path = Path(
                    str(files("ai.backend.runner").joinpath("jupyter-custom.css"))
                )
                logo_path = Path(str(files("ai.backend.runner").joinpath("logo.svg")))
                font_path = Path(str(files("ai.backend.runner").joinpath("roboto.ttf")))
                font_italic_path = Path(
                    str(files("ai.backend.runner").joinpath("roboto-italic.ttf"))
                )
                bashrc_path = Path(str(files("ai.backend.runner").joinpath(".bashrc")))
                bash_profile_path = Path(str(files("ai.backend.runner").joinpath(".bash_profile")))
                zshrc_path = Path(str(files("ai.backend.runner").joinpath(".zshrc")))
                vimrc_path = Path(str(files("ai.backend.runner").joinpath(".vimrc")))
                tmux_conf_path = Path(str(files("ai.backend.runner").joinpath(".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(zshrc_path.resolve(), self.work_dir / ".zshrc")
                shutil.copy(vimrc_path.resolve(), self.work_dir / ".vimrc")
                shutil.copy(tmux_conf_path.resolve(), self.work_dir / ".tmux.conf")
                if KernelFeatures.UID_MATCH in self.kernel_features:
                    uid = self.local_config.container.kernel_uid
                    gid = self.local_config.container.kernel_gid
                    if os.geteuid() == 0:  # only possible when I am root.
                        os.chown(self.work_dir, uid, gid)
                        os.chown(self.work_dir / ".jupyter", uid, gid)
                        os.chown(self.work_dir / ".jupyter" / "custom", uid, gid)
                        os.chown(self.work_dir / ".bashrc", uid, gid)
                        os.chown(self.work_dir / ".bash_profile", uid, gid)
                        os.chown(self.work_dir / ".zshrc", uid, gid)
                        os.chown(self.work_dir / ".vimrc", uid, gid)
                        os.chown(self.work_dir / ".tmux.conf", uid, gid)

            await loop.run_in_executor(None, _clone_dotfiles)

    @override
    async def get_intrinsic_mounts(self) -> Sequence[Mount]:
        mounts: list[Mount] = [
            # Mount scratch directory
            Mount(
                MountTypes.K8S_GENERIC,
                self.rel_config_dir,
                Path("/home/config"),
                MountPermission.READ_ONLY,
                opts={
                    "name": f"kernel-{self.kernel_id}-scratches",
                },
            ),
            Mount(
                MountTypes.K8S_GENERIC,
                self.rel_work_dir,
                Path("/home/work"),
                MountPermission.READ_WRITE,
                opts={
                    "name": f"kernel-{self.kernel_id}-scratches",
                },
            ),
        ]

        rel_agent_sockpath = Path(str(self.agent_sockpath).split("/")[-1])
        # agent-socket mount
        if sys.platform != "darwin":
            mounts.append(
                Mount(
                    MountTypes.K8S_GENERIC,
                    rel_agent_sockpath,
                    Path("/opt/kernel/agent.sock"),
                    MountPermission.READ_WRITE,
                    opts={
                        "name": f"kernel-{self.kernel_id}-scratches",
                    },
                )
            )

        # TODO: Find way to mount extra volumes

        return mounts

    @property
    @override
    def repl_ports(self) -> Sequence[int]:
        return (2000, 2001)

    @property
    @override
    def protected_services(self) -> Sequence[str]:
        # NOTE: Currently K8s does not support binding container ports to 127.0.0.1 when using NodePort.
        return ()

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

    @override
    async def prepare_ssh(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(usedforsecurity=False)
        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 Exception:
            # 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)

        await self.process_volumes([
            KubernetesConfigMapVolume(
                name=f"kernel-{self.kernel_id}-ssh-keypair",
                configMap={
                    "name": "ssh-keypair-hash",
                },
            ),
        ])
        await 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",
                },
            ),
        ])

    @override
    async def process_mounts(self, mounts: Sequence[Mount]) -> None:
        for i, mount in enumerate(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.warning(
                    "Mount {}:{} -> Mount type {} it not supported on K8s Agent. Skipping mount",
                    mount.source,
                    mount.target,
                    mount.type,
                )

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

    @override
    def get_runner_mount(
        self,
        type: MountTypes,
        src: str | Path,
        target: str | Path,
        perm: MountPermission = MountPermission.READ_ONLY,
        opts: Mapping[str, Any] | None = 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

    @override
    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}",
                },
            )
            await 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)

    @override
    async def apply_accelerator_allocation(
        self,
        computer: AbstractComputePlugin,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> None:
        # update_nested_dict(
        #     self.computer_docker_args,
        #     await computer.generate_docker_args(self.docker, device_alloc))
        # TODO: add support for accelerator allocation
        pass

    @override
    async def generate_accelerator_mounts(
        self,
        computer: AbstractComputePlugin,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> list[MountInfo]:
        return []

    async def generate_deployment_object(
        self,
        image: str,
        environ: Mapping[str, Any],
        ports: list[int],
        command: list[str],
        labels: Mapping[str, Any] = {},
    ) -> dict[str, Any]:
        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],
                                "securityContext": {"privileged": True},
                            }
                        ],
                        "volumes": [cattr.unstructure(v) for v in self.volumes],
                        "securityContext": {"privileged": True},
                    },
                },
            },
        }

    @override
    async def prepare_container(
        self,
        resource_spec: KernelResourceSpec,
        environ: Mapping[str, str],
        service_ports: Any,
        cluster_info: ClusterInfo,
    ) -> KubernetesKernel:
        loop = current_loop()
        if self.restarting:
            pass
        else:
            if bootstrap := self.kernel_config.get("bootstrap_script"):

                def _write_user_bootstrap_script() -> None:
                    (self.work_dir / "bootstrap.sh").write_text(bootstrap)
                    if KernelFeatures.UID_MATCH in self.kernel_features:
                        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)

                await loop.run_in_executor(None, _write_user_bootstrap_script)

            def _write_config(file_name: str, content: str) -> None:
                file_path = self.config_dir / file_name
                file_path.write_text(content)
                if KernelFeatures.UID_MATCH in self.kernel_features:
                    uid = self.local_config.container.kernel_uid
                    gid = self.local_config.container.kernel_gid
                    if os.geteuid() == 0:
                        os.chown(str(file_path), uid, gid)

            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)
                await self.process_volumes([
                    KubernetesConfigMapVolume(
                        name="ssh-config",
                        configMap={
                            "name": f"kernel-{self.kernel_id}-ssh-config",
                        },
                    ),
                ])
                await 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"])

            tmp = Path(file_path)
            while tmp != self.work_dir:
                tmp.chmod(int(dotfile["perm"], 8))
                if KernelFeatures.UID_MATCH in self.kernel_features and os.geteuid() == 0:
                    uid = self.local_config.container.kernel_uid
                    gid = self.local_config.container.kernel_gid
                    os.chown(tmp, uid, gid)
                tmp = tmp.parent

        # TODO: Mark shmem feature as unsupported when advertising agent

        return KubernetesKernel(
            self.ownership_data,
            self.kernel_config["network_id"],
            self.image_ref,
            self.kspec_version,
            agent_config=self.local_config.model_dump(by_alias=True),
            service_ports=service_ports,
            resource_spec=resource_spec,
            environ=environ,
            data={},
        )

    @override
    async def start_container(
        self,
        kernel_obj: AbstractKernel,
        cmdargs: list[str],
        resource_opts: Any,
        preopen_ports: Any,
        cluster_info: ClusterInfo,
    ) -> Mapping[str, Any]:
        image_labels = self.kernel_config["image"]["labels"]
        service_ports = kernel_obj.service_ports
        environ = kernel_obj.environ

        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()
        apps_api = kube_client.AppsV1Api()
        exposed_ports = [*self.repl_ports]
        for sport in service_ports:
            exposed_ports.extend(sport["container_ports"])

        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 enumerate(exposed_ports)
            ],
        )

        async def rollup(
            functions: list[
                tuple[
                    functools.partial[Coroutine[Any, Any, None]] | None,
                    functools.partial[Coroutine[Any, Any, None]] | None,
                ]
            ],
        ) -> None:
            rollback_functions: list[functools.partial[Coroutine[Any, Any, None]] | None] = []

            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[
                functools.partial[Coroutine[Any, Any, None]] | None,
                functools.partial[Coroutine[Any, Any, None]] | None,
            ]
        ] = []

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

        if expose_service_api_response.spec is None:
            raise K8sError("expose_service_api_response.spec is None")
        node_ports: list[V1ServicePort] = 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[int, 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
        repl_in_port: int
        repl_out_port: int
        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()])

        return {
            "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),
        }


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

    def __init__(
        self,
        etcd: AgentEtcdClientView,
        local_config: AgentUnifiedConfig,
        *,
        stats_monitor: StatsPluginContext,
        error_monitor: ErrorPluginContext,
        skip_initial_scan: bool = False,
        agent_public_key: PublicKey | None,
        kernel_registry: KernelRegistry,
        computers: Mapping[DeviceName, ComputerContext],
        slots: Mapping[SlotName, Decimal],
        agent_class: AgentClass,
    ) -> None:
        super().__init__(
            etcd,
            local_config,
            stats_monitor=stats_monitor,
            error_monitor=error_monitor,
            skip_initial_scan=skip_initial_scan,
            agent_public_key=agent_public_key,
            kernel_registry=kernel_registry,
            computers=computers,
            slots=slots,
            agent_class=agent_class,
        )
        self._kernel_recovery = KubernetesKernelRegistryRecovery.create(
            KubernetesKernelRegistryRecoveryArgs(
                scratch_root=self.local_config.container.scratch_root,
                ipc_base_path=self.local_config.agent.ipc_base_path,
                var_base_path=self.local_config.agent.var_base_path,
                agent_id=self.id,
                local_instance_id=self.local_instance_id,
            )
        )

    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.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) -> None:
        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 == ScratchType.K8S_NFS:
                scratch_nfs_address = self.local_config.container.scratch_nfs_address
                if scratch_nfs_address is None:
                    raise K8sError("scratch_nfs_address must be set when using K8S_NFS")
                new_pv: NFSPersistentVolume | HostPathPersistentVolume = NFSPersistentVolume(
                    scratch_nfs_address,
                    "backend-ai-static-pv",
                    capacity,
                )
                new_pv.label(
                    "backend.ai/backend-ai-scratch-volume",
                    scratch_nfs_address,
                )
                scratch_nfs_options = self.local_config.container.scratch_nfs_options
                if scratch_nfs_options is not None:
                    new_pv.options = [x.strip() for x in scratch_nfs_options.split(",")]
            elif self.local_config.container.scratch_type == ScratchType.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 Exception:
                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 == ScratchType.K8S_NFS:
                scratch_nfs_address = self.local_config.container.scratch_nfs_address
                if scratch_nfs_address is None:
                    raise K8sError("scratch_nfs_address must be set when using K8S_NFS")
                new_pvc.label(
                    "backend.ai/backend-ai-scratch-volume",
                    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) -> None:
        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) -> None:
        # 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

    @override
    async def _load_kernel_registry_from_recovery(self) -> MutableMapping[KernelId, AbstractKernel]:
        return await self._kernel_recovery.load_kernel_registry()

    @override
    async def _write_kernel_registry_to_recovery(
        self,
        kernel_registry: MutableMapping[KernelId, AbstractKernel],
        metadata: KernelRegistrySaveMetadata,
    ) -> None:
        await self._kernel_recovery.save_kernel_registry(kernel_registry, metadata)

    @override
    def get_cgroup_path(self, controller: str, container_id: str) -> Path:
        # Not implemented yet for K8s Agent
        return Path()

    @override
    def get_cgroup_version(self) -> str:
        # Not implemented yet for K8s Agent
        return ""

    @override
    async def extract_image_command(self, image: str) -> str | None:
        raise NotImplementedError

    @override
    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) -> None:
                kernel_id: KernelId | str | None = "(unknown)"
                try:
                    kernel_id = await get_kernel_id_from_deployment(pod)
                    if kernel_id is None or kernel_id not in self.kernel_registry:
                        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(
                            (
                                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

    @override
    async def resolve_image_distro(self, image: ImageConfig) -> str:
        image_labels = image["labels"]
        distro = image_labels.get("ai.backend.base-distro")
        if distro:
            return distro
        raise NotImplementedError

    @override
    async def scan_images(self) -> ScanImagesResult:
        # Retrieving image label from registry api is not possible
        return ScanImagesResult(scanned_images={}, removed_images={})

    @override
    async def push_image(
        self,
        image_ref: ImageRef,
        registry_conf: ImageRegistry,
        *,
        timeout_seconds: float | None | Sentinel = Sentinel.TOKEN,
    ) -> None:
        raise NotImplementedError

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

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

    @override
    async def purge_images(self, request: PurgeImagesReq) -> PurgeImagesResp:
        # TODO: Add support for appropriate image purging mechanism on K8s
        return PurgeImagesResp([])

    @override
    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

    @override
    async def init_kernel_context(
        self,
        ownership_data: KernelOwnershipData,
        kernel_image: ImageRef,
        kernel_config: KernelCreationConfig,
        *,
        restarting: bool = False,
        cluster_ssh_port_mapping: ClusterSSHPortMapping | None = None,
    ) -> KubernetesKernelCreationContext:
        distro = await self.resolve_image_distro(kernel_config["image"])
        return KubernetesKernelCreationContext(
            ownership_data,
            self.event_producer,
            kernel_image,
            kernel_config,
            distro,
            self.local_config,
            self.agent_sockpath,
            self.computers,
            self.workers,
            "backend-ai-static-pvc",
            restarting=restarting,
        )

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

        async def force_cleanup(_reason: str = "self-terminated") -> None:
            await self.send_event("kernel_terminated", kernel_id, "self-terminated", None)  # type: ignore[attr-defined]

        try:
            kernel = self.kernel_registry[kernel_id]
        except Exception:
            log.warning("_destroy_kernel({0}) kernel missing (already dead?)", kernel_id)
            await asyncio.shield(self.k8s_ptask_group.create_task(force_cleanup()))
            return
        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 Exception:
            log.warning("_destroy({0}) kernel missing (already dead?)", kernel_id)

    @override
    async def clean_kernel(
        self,
        kernel_id: KernelId,
        container_id: ContainerId | None,
        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))

    @override
    async def create_local_network(self, network_name: str) -> None:
        raise NotImplementedError

    @override
    async def destroy_local_network(self, network_name: str) -> None:
        raise NotImplementedError

    @override
    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,
        )

    @override
    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"

        def _write_bytes(data: bytes) -> None:
            (config_dir / name).write_bytes(data)

        return await loop.run_in_executor(
            None,
            _write_bytes,
            data,
        )


async def get_kernel_id_from_deployment(pod: Any) -> KernelId | None:
    # TODO: create function which extracts kernel id from pod object
    if (kernel_id := pod.get("metadata", {}).get("name")) is not None:
        return KernelId(uuid.UUID(kernel_id))
    return None


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,
    )
