import asyncio
import logging
import lzma
import os
import shutil
import textwrap
from collections.abc import Mapping
from importlib.resources import files
from pathlib import Path, PurePosixPath
from typing import Any, override

import zmq
from aiodocker.docker import Docker
from aiotools import TaskGroup
from kubernetes_asyncio import client as kube_client
from kubernetes_asyncio import config as kube_config
from kubernetes_asyncio import watch

from ai.backend.agent.errors import KernelRunnerNotInitializedError
from ai.backend.agent.kernel import AbstractCodeRunner, AbstractKernel
from ai.backend.agent.resources import KernelResourceSpec
from ai.backend.agent.types import AgentEventData, KernelOwnershipData
from ai.backend.agent.utils import get_arch_name
from ai.backend.common.asyncio import current_loop
from ai.backend.common.docker import ImageRef
from ai.backend.common.dto.agent.response import CodeCompletionResp
from ai.backend.common.events.dispatcher import EventProducer
from ai.backend.common.types import CommitStatus
from ai.backend.logging import BraceStyleAdapter
from ai.backend.plugin.entrypoint import scan_entrypoints

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


class KubernetesKernel(AbstractKernel):
    deployment_name: str

    def __init__(
        self,
        ownership_data: KernelOwnershipData,
        network_id: str,
        image: ImageRef,
        version: int,
        *,
        agent_config: Mapping[str, Any],
        resource_spec: KernelResourceSpec,
        service_ports: Any,  # TODO: type-annotation
        data: dict[str, Any],
        environ: Mapping[str, Any],
    ) -> None:
        super().__init__(
            ownership_data,
            network_id,
            image,
            version,
            agent_config=agent_config,
            resource_spec=resource_spec,
            service_ports=service_ports,
            data=data,
            environ=environ,
        )

        self.deployment_name = f"kernel-{ownership_data.kernel_id}"

    @override
    async def close(self) -> None:
        await self.scale(0)

    @override
    async def create_code_runner(
        self, event_producer: EventProducer, *, client_features: frozenset[str], api_version: int
    ) -> AbstractCodeRunner:
        scale = await self.scale(1)
        if scale.to_dict()["spec"]["replicas"] == 0:
            log.error("Scaling failed! Response body: {0}", scale)
            raise ValueError("Scaling failed!")

        if scale.to_dict()["status"]["replicas"] == 0:
            # Wait for scaling to complete using asyncio.Event pattern
            scaled_event = asyncio.Event()

            async def wait_for_scaling() -> None:
                max_iterations = 120  # 60 seconds timeout at 0.5s per iteration
                for _ in range(max_iterations):
                    if await self.is_scaled():
                        scaled_event.set()
                        return
                    await asyncio.sleep(0.5)
                # Timeout - still set event to avoid hanging
                scaled_event.set()

            wait_task = asyncio.create_task(wait_for_scaling())
            try:
                await scaled_event.wait()
            finally:
                if not wait_task.done():
                    wait_task.cancel()
                    try:
                        await wait_task
                    except asyncio.CancelledError:
                        pass

        # TODO: Find way to detect if kernel runner has started inside container

        runner = await KubernetesCodeRunner.new(
            self.kernel_id,
            self.session_id,
            event_producer,
            kernel_host=self.data["kernel_host"],
            repl_in_port=self.data["repl_in_port"],
            repl_out_port=self.data["repl_out_port"],
            exec_timeout=0,
            client_features=client_features,
        )

        retries = 0
        while True:
            try:
                await runner.feed_and_get_status()
                break
            except zmq.error.ZMQError:
                if retries < 4:
                    retries += 1
                    log.debug("Socket not responding, retrying #{}", retries)
                    await asyncio.sleep(retries**2)
                else:
                    raise

        return runner

    async def scale(self, num: int) -> Any:
        await kube_config.load_kube_config()
        apps_api = kube_client.AppsV1Api()
        try:
            return await apps_api.replace_namespaced_deployment_scale(
                self.deployment_name,
                "backend-ai",
                body={
                    "apiVersion": "autoscaling/v1",
                    "kind": "Scale",
                    "metadata": {
                        "name": self.deployment_name,
                        "namespace": "backend-ai",
                    },
                    "spec": {"replicas": num},
                    "status": {"replicas": num, "selector": f"run={self.deployment_name}"},
                },
            )
        except Exception as e:
            log.exception("scale failed: {}", e)

    async def is_scaled(self) -> bool:
        await kube_config.load_kube_config()
        apps_api = kube_client.AppsV1Api()
        core_api = kube_client.CoreV1Api()
        scale = await apps_api.read_namespaced_deployment(self.deployment_name, "backend-ai")

        if scale.to_dict()["status"]["replicas"] == 0:
            return False
        for condition in scale.to_dict()["status"]["conditions"]:
            if not condition["status"]:
                return False

        pods = await core_api.list_namespaced_pod(
            "backend-ai",
            label_selector=f"run=kernel-{self.kernel_id}",
        )
        pods = pods.to_dict()["items"] or []
        if len(pods) == 0:
            return False
        for pod in pods:
            containers = pod["status"]["container_statuses"] or []
            if len(containers) == 0:
                return False
            for container in containers:
                started = container.get("started")
                if not container["ready"] or (started is not None and not started):
                    return False
        return True

    @override
    async def get_completions(self, text: str, opts: Mapping[str, Any]) -> CodeCompletionResp:
        if self.runner is None:
            raise KernelRunnerNotInitializedError("Kernel runner is not initialized")
        result = await self.runner.feed_and_get_completion(text, opts)
        return CodeCompletionResp(result=result)

    @override
    async def check_status(self) -> dict[str, Any] | None:
        if self.runner is None:
            raise KernelRunnerNotInitializedError("Kernel runner is not initialized")
        return await self.runner.feed_and_get_status()

    @override
    async def get_logs(self) -> dict[str, Any]:
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()

        result = await core_api.read_namespaced_pod_log(self.kernel_id, "backend-ai")
        return {"logs": result.data.decode("utf-8")}

    @override
    async def interrupt_kernel(self) -> dict[str, Any]:
        if self.runner is None:
            raise KernelRunnerNotInitializedError("Kernel runner is not initialized")
        await self.runner.feed_interrupt()
        return {"status": "finished"}

    @override
    async def start_service(self, service: str, opts: Mapping[str, Any]) -> dict[str, Any]:
        if self.runner is None:
            raise KernelRunnerNotInitializedError("Kernel runner is not initialized")
        if self.data.get("block_service_ports", False):
            return {
                "status": "failed",
                "error": "operation blocked",
            }
        for sport in self.service_ports:
            if sport["name"] == service:
                break
        else:
            return {"status": "failed", "error": "invalid service name"}
        return await self.runner.feed_start_service({
            "name": service,
            "port": sport["container_ports"][0],  # primary port
            "ports": sport["container_ports"],
            "protocol": sport["protocol"],
            "options": opts,
        })

    @override
    async def start_model_service(self, model_service: Mapping[str, Any]) -> dict[str, Any]:
        if self.runner is None:
            raise KernelRunnerNotInitializedError("Kernel runner is not initialized")
        return await self.runner.feed_start_model_service(model_service)

    @override
    async def shutdown_service(self, service: str) -> None:
        if self.runner is None:
            raise KernelRunnerNotInitializedError("Kernel runner is not initialized")
        await self.runner.feed_shutdown_service(service)

    @override
    async def get_service_apps(self) -> dict[str, Any]:
        if self.runner is None:
            raise KernelRunnerNotInitializedError("Kernel runner is not initialized")
        return await self.runner.feed_service_apps()

    @override
    async def check_duplicate_commit(self, kernel_id: Any, subdir: Any) -> CommitStatus:
        log.error("Committing in Kubernetes is not supported yet.")
        raise NotImplementedError

    @override
    async def commit(
        self,
        kernel_id: Any,
        subdir: Any,
        *,
        canonical: str | None = None,
        filename: str | None = None,
        extra_labels: dict[str, str] | None = None,
    ) -> None:
        # TODO: Implement container commit on Kubernetes kernel.
        if extra_labels is None:
            extra_labels = {}
        log.error("Committing in Kubernetes is not supported yet.")
        raise NotImplementedError

    @override
    async def accept_file(self, container_path: os.PathLike[str] | str, filedata: bytes) -> None:
        loop = current_loop()
        host_work_dir: Path = (
            self.agent_config["container"]["scratch-root"] / str(self.kernel_id) / "work"
        )
        host_abspath = (host_work_dir / container_path).resolve(strict=False)
        if not host_abspath.is_relative_to(host_work_dir):
            raise PermissionError("Not allowed to upload files outside /home/work")

        def _write_to_disk() -> None:
            host_abspath.parent.mkdir(parents=True, exist_ok=True)
            host_abspath.write_bytes(filedata)

        try:
            await loop.run_in_executor(None, _write_to_disk)
        except OSError as e:
            raise RuntimeError(
                f"{self.kernel_id}: writing uploaded file failed: {container_path} -> {host_abspath} ({e!r})"
            ) from e

    @override
    async def download_file(self, container_path: os.PathLike[str] | str) -> bytes:
        # TODO: Implement file operations with pure Kubernetes API
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()

        container_home_path = PurePosixPath("/home/work")
        container_abspath = PurePosixPath(os.path.normpath(container_home_path / container_path))
        if not container_abspath.is_relative_to(container_home_path):
            raise PermissionError("You cannot download files outside /home/work")

        async with watch.Watch().stream(
            core_api.connect_get_namespaced_pod_exec,
            self.kernel_id,
            "backend-ai",
            command=["tar", "cf", "-", container_abspath],
            stderr=True,
            stdin=True,
            stdout=True,
            tty=False,
            _preload_content=False,
        ) as stream:
            async for event in stream:
                log.debug("stream: {}", event)
                # TODO: retrieve the output stream as a bytes buffer
        return b""

    @override
    async def download_single(self, container_path: os.PathLike[str] | str) -> bytes:
        # TODO: Implement download single file operations with pure Kubernetes API
        log.error("download_single() in the k8s backend is not supported yet.")
        raise NotImplementedError

    @override
    async def list_files(self, container_path: os.PathLike[str] | str) -> dict[str, Any]:
        # TODO: Implement file operations with pure Kubernetes API
        await kube_config.load_kube_config()
        core_api = kube_client.CoreV1Api()

        # Confine the lookable paths in the home directory
        container_home_path = PurePosixPath("/home/work")
        container_abspath = PurePosixPath(os.path.normpath(container_home_path / container_path))
        if not container_abspath.is_relative_to(container_home_path):
            raise PermissionError("You cannot list files outside /home/work")

        # Gather individual file information in the target path.
        code = textwrap.dedent(
            """
        import json
        import os
        import stat
        import sys

        files = []
        for f in os.scandir(sys.argv[1]):
            fstat = f.stat(follow_symlinks=False)

            ctime = fstat.st_ctime  # TODO: way to get concrete create time?
            mtime = fstat.st_mtime
            atime = fstat.st_atime
            files.append({
                'mode': stat.filemode(fstat.st_mode),
                'size': fstat.st_size,
                'ctime': ctime,
                'mtime': mtime,
                'atime': atime,
                'filename': f.name,
            })
        print(json.dumps(files))
        """
        )

        command = ["/opt/backend.ai/bin/python", "-c", code, str(container_path)]
        async with watch.Watch().stream(
            core_api.connect_get_namespaced_pod_exec,
            self.kernel_id,
            "backend-ai",
            command=command,
            stderr=True,
            stdin=True,
            stdout=True,
            tty=False,
            _preload_content=False,
        ) as stream:
            async for event in stream:
                log.debug("stream: {}", event)
                # TODO: retrieve the output stream as a bytes buffer

        return {"files": "", "errors": "", "abspath": str(container_path)}

    @override
    async def notify_event(self, evdata: AgentEventData) -> None:
        raise NotImplementedError


class KubernetesCodeRunner(AbstractCodeRunner):
    kernel_host: str
    repl_in_port: int
    repl_out_port: int

    def __init__(
        self,
        kernel_id: Any,
        session_id: Any,
        event_producer: EventProducer,
        *,
        kernel_host: str,
        repl_in_port: int,
        repl_out_port: int,
        exec_timeout: int = 0,
        client_features: frozenset[str] | None = None,
    ) -> None:
        super().__init__(
            kernel_id,
            session_id,
            event_producer,
            exec_timeout=exec_timeout,
            client_features=client_features,
        )
        self.kernel_host = kernel_host
        self.repl_in_port = repl_in_port
        self.repl_out_port = repl_out_port

    @override
    async def get_repl_in_addr(self) -> str:
        return f"tcp://{self.kernel_host}:{self.repl_in_port}"

    @override
    async def get_repl_out_addr(self) -> str:
        return f"tcp://{self.kernel_host}:{self.repl_out_port}"


async def prepare_krunner_env_impl(
    distro: str, entrypoint_name: str, root_path: str
) -> tuple[str, str | None]:
    docker = Docker()
    arch = get_arch_name()
    current_version = int(
        Path(
            str(
                files(f"ai.backend.krunner.{entrypoint_name}").joinpath(
                    f"./krunner-version.{distro}.txt"
                )
            )
        )
        .read_text()
        .strip()
    )
    krunner_folder_name = f"backendai-krunner.v{current_version}.{arch}.{distro}"
    target_path = Path(root_path) / krunner_folder_name
    extractor_image = "backendai-krunner-extractor:latest"

    try:
        for item in await docker.images.list():
            if item["RepoTags"] is None or len(item["RepoTags"]) == 0:
                continue
            if item["RepoTags"][0] == extractor_image:
                break
        else:
            log.info("preparing the Docker image for krunner extractor...")
            extractor_archive = str(
                files("ai.backend.runner").joinpath(f"krunner-extractor.img.{arch}.tar.xz")
            )
            with lzma.open(extractor_archive, "rb") as reader:
                proc = await asyncio.create_subprocess_exec(*["docker", "load"], stdin=reader)
                if await proc.wait() != 0:
                    raise RuntimeError("loading krunner extractor image has failed!")

        log.info("checking krunner-env for {}.{}...", distro, arch)

        if not target_path.exists():
            log.info("populating {} volume version {}", krunner_folder_name, current_version)
            target_path.mkdir(exist_ok=False)
            archive_path = Path(
                str(
                    files(f"ai.backend.krunner.{entrypoint_name}").joinpath(
                        f"krunner-env.{distro}.{arch}.tar.xz"
                    )
                )
            ).resolve()
            extractor_path = Path(
                str(files("ai.backend.runner").joinpath("krunner-extractor.sh"))
            ).resolve()

            log.debug(
                "Executing {}",
                " ".join([
                    "docker",
                    "run",
                    "--rm",
                    "-i",
                    "-v",
                    f"{archive_path}:/root/archive.tar.xz",
                    "-v",
                    f"{extractor_path}:/root/krunner-extractor.sh",
                    "-v",
                    f"{target_path.absolute().as_posix()}:/root/volume",
                    "-e",
                    f"KRUNNER_VERSION={current_version}",
                    extractor_image,
                    "/root/krunner-extractor.sh",
                ]),
            )

            proc = await asyncio.create_subprocess_exec(*[
                "docker",
                "run",
                "--rm",
                "-i",
                "-v",
                f"{archive_path}:/root/archive.tar.xz",
                "-v",
                f"{extractor_path}:/root/krunner-extractor.sh",
                "-v",
                f"{target_path.absolute().as_posix()}:/root/volume",
                "-e",
                f"KRUNNER_VERSION={current_version}",
                extractor_image,
                "/root/krunner-extractor.sh",
            ])
            if await proc.wait() != 0:
                raise RuntimeError("extracting krunner environment has failed!")
    except Exception:
        log.exception("unexpected error")
        return distro, None
    finally:
        await docker.close()
    return distro, krunner_folder_name


async def copy_runner_files(scratch_path: Path) -> None:
    artifact_path = Path(str(files("ai.backend.agent").joinpath("../runner")))
    kernel_path = Path(str(files("ai.backend.agent").joinpath("../kernel")))
    helpers_path = Path(str(files("ai.backend.agent").joinpath("../helpers")))

    destination_path = scratch_path

    if (destination_path / "runner").exists():
        shutil.rmtree(destination_path / "runner", ignore_errors=True)
    (destination_path / "runner").mkdir(parents=True)

    target_files = [
        "entrypoint.sh",
        "*.bin",
        "*.so",
        "DO_NOT_STORE_PERSISTENT_FILES_HERE.md",
        "extract_dotfiles.py",
        "fantompass.py",
        "hash_phrase.py",
        "words.json",
    ]

    for target_glob in target_files:
        for matched_path in artifact_path.glob(target_glob):
            shutil.copy(matched_path.resolve(), destination_path / "runner")

    if (destination_path / "kernel").exists():
        shutil.rmtree(destination_path / "kernel", ignore_errors=True)
    shutil.copytree(kernel_path.resolve(), destination_path / "kernel")

    if (destination_path / "helpers").exists():
        shutil.rmtree(destination_path / "helpers", ignore_errors=True)
    shutil.copytree(helpers_path.resolve(), destination_path / "helpers")


async def prepare_krunner_env(local_config: Mapping[str, Any]) -> Mapping[str, str]:
    """
    Check if the volume "backendai-krunner.{distro}.{arch}" exists and is up-to-date.
    If not, automatically create it and update its content from the packaged pre-built krunner
    tar archives.
    """

    all_distros: list[tuple[str, str]] = []
    entry_prefix = "backendai_krunner_v10"
    for entrypoint in scan_entrypoints(entry_prefix):
        log.debug("loading krunner pkg: {}", entrypoint.module)
        plugin = entrypoint.load()
        await plugin.init({})  # currently does nothing
        provided_versions = (
            Path(str(files(f"ai.backend.krunner.{entrypoint.name}").joinpath("versions.txt")))
            .read_text()
            .splitlines()
        )
        all_distros.extend((distro, entrypoint.name) for distro in provided_versions)

    scratch_mount = local_config["container"]["scratch-root"]
    await copy_runner_files(Path(scratch_mount))

    tasks = []
    async with TaskGroup() as tg:
        for distro, entrypoint_name in all_distros:
            tasks.append(
                tg.create_task(prepare_krunner_env_impl(distro, entrypoint_name, scratch_mount))
            )
    distro_volumes = [t.result() for t in tasks if not t.cancelled()]
    result = {}
    for distro_name_and_version, volume_name in distro_volumes:
        if volume_name is None:
            continue
        result[distro_name_and_version] = volume_name
    return result
