from __future__ import annotations

import asyncio
import hashlib
import io
import ipaddress
import logging
import platform
import re
import resource
from collections.abc import Mapping, MutableMapping
from contextlib import AbstractAsyncContextManager
from decimal import Decimal
from pathlib import Path
from typing import (
    Any,
    Final,
    NamedTuple,
    Protocol,
    TypedDict,
    TypeVar,
    overload,
)
from uuid import UUID

import aiodocker
import trafaret as t
from aiodocker.docker import DockerContainer

from ai.backend.common import identity
from ai.backend.common.asyncio import current_loop
from ai.backend.common.cgroup import (
    get_cgroup_of_pid,
    get_container_id_of_cgroup,
    get_container_pids,
)
from ai.backend.common.etcd import AsyncEtcd
from ai.backend.common.json import dump_json_str
from ai.backend.common.types import PID, ContainerId, ContainerPID, HostPID, KernelId
from ai.backend.logging import BraceStyleAdapter

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

IPNetwork = ipaddress.IPv4Network | ipaddress.IPv6Network
IPAddress = ipaddress.IPv4Address | ipaddress.IPv6Address

InOtherContainerPID: Final = ContainerPID(PID(-2))
NotContainerPID: Final = ContainerPID(PID(-1))
NotHostPID: Final = HostPID(PID(-1))


class SupportsAsyncClose(Protocol):
    async def close(self) -> None: ...


_SupportsAsyncCloseT = TypeVar("_SupportsAsyncCloseT", bound=SupportsAsyncClose)


class closing_async(AbstractAsyncContextManager[_SupportsAsyncCloseT]):
    """
    contextlib.closing calls close(), and aiotools.aclosing() calls aclose().
    This context manager calls close() as a coroutine.
    """

    def __init__(self, obj: _SupportsAsyncCloseT) -> None:
        self.obj = obj

    async def __aenter__(self) -> _SupportsAsyncCloseT:
        return self.obj

    async def __aexit__(self, *exc_info: Any) -> None:
        await self.obj.close()


def generate_local_instance_id(hint: str) -> str:
    return hashlib.md5(hint.encode("utf-8"), usedforsecurity=False).hexdigest()[:12]


def get_arch_name() -> str:
    ret = platform.machine().lower()
    aliases = {
        "arm64": "aarch64",  # macOS with LLVM
        "amd64": "x86_64",  # Windows/Linux
        "x64": "x86_64",  # Windows
        "x32": "x86",  # Windows
        "i686": "x86",  # Windows
    }
    return aliases.get(ret, ret)


def update_nested_dict(dest: MutableMapping[str, Any], additions: Mapping[str, Any]) -> None:
    for k, v in additions.items():
        if k not in dest:
            dest[k] = v
        else:
            if isinstance(dest[k], MutableMapping):
                if not isinstance(v, MutableMapping):
                    raise TypeError(
                        f"Cannot merge non-mapping value into mapping for key '{k}': "
                        f"expected MutableMapping, got {type(v).__name__}."
                    )
                update_nested_dict(dest[k], v)
            elif isinstance(dest[k], list):
                if not isinstance(v, list):
                    raise TypeError(
                        f"Cannot extend non-list value into list for key '{k}': "
                        f"expected List, got {type(v).__name__}."
                    )
                dest[k].extend(v)
            else:
                dest[k] = v


def numeric_list(s: str) -> list[int]:
    return [int(p) for p in s.split()]


def remove_exponent(num: Decimal) -> Decimal:
    return num.quantize(Decimal(1)) if num == num.to_integral() else num.normalize()


@overload
def read_sysfs(path: str | Path, type_: type[bool], default: bool | None = None) -> bool: ...


@overload
def read_sysfs(path: str | Path, type_: type[int], default: int | None = None) -> int: ...


@overload
def read_sysfs(path: str | Path, type_: type[float], default: float | None = None) -> float: ...


@overload
def read_sysfs(path: str | Path, type_: type[str], default: str | None = None) -> str: ...


def read_sysfs(path: str | Path, type_: type[Any], default: Any | None = None) -> Any:
    def_vals: Mapping[Any, Any] = {
        bool: False,
        int: 0,
        float: 0.0,
        str: "",
    }
    if type_ not in def_vals:
        raise TypeError("unsupported conversion type from sysfs content")
    if default is None:
        default = def_vals[type_]
    try:
        raw_str = Path(path).read_text().strip()
        if type_ is bool:
            return t.ToBool().check(raw_str)
        return type_(raw_str)
    except OSError:
        return default


async def read_tail(path: Path, nbytes: int) -> bytes:
    file_size = path.stat().st_size

    def _read_tail() -> bytes:
        with path.open("rb") as f:
            f.seek(max(file_size - nbytes, 0), io.SEEK_SET)
            return f.read(nbytes)

    loop = current_loop()
    return await loop.run_in_executor(None, _read_tail)


async def get_kernel_id_from_container(val: str | DockerContainer) -> KernelId | None:
    if isinstance(val, DockerContainer):
        if "Name" not in val._container:
            await val.show()
        name = val["Name"]
    elif isinstance(val, str):
        name = val
    name = name.lstrip("/")
    if not name.startswith("kernel."):
        return None
    try:
        return KernelId(UUID(name.rsplit(".", 2)[-1]))
    except (IndexError, ValueError):
        return None


async def get_subnet_ip(etcd: AsyncEtcd, network: str, fallback_addr: str = "0.0.0.0") -> str:
    raw_subnet = await etcd.get(f"config/network/subnet/{network}")
    if raw_subnet is None:
        addr = fallback_addr
    else:
        subnet = ipaddress.ip_network(raw_subnet)
        if subnet.prefixlen == 0:
            addr = fallback_addr
        else:
            local_ipaddrs = [*identity.fetch_local_ipaddrs(subnet)]
            log.debug("get_subnet_ip(): subnet {} candidates: {}", subnet, local_ipaddrs)
            if local_ipaddrs:
                addr = str(local_ipaddrs[0])
            else:
                addr = fallback_addr
    return addr


class ProcessInfo(NamedTuple):
    pid: int
    command: str


async def get_host_process_table(
    docker: aiodocker.Docker,
    container_id: str,
) -> list[ProcessInfo]:
    result = await docker._query_json(f"containers/{container_id}/top", method="GET")
    procs = result["Processes"]
    return [ProcessInfo(int(x[1]), x[7]) for x in procs]


async def get_container_process_table(
    docker: aiodocker.Docker,
    container_id: str,
) -> list[ProcessInfo]:
    # Get process table from inside container (execute 'ps -aux' command from container).
    # Filter processes which have exactly the same COMMAND like above.
    result = await docker._query_json(
        f"containers/{container_id}/exec",
        method="POST",
        data={
            "AttachStdin": False,
            "AttachStdout": True,
            "AttachStderr": True,
            "Cmd": ["ps", "aux"],
        },
    )
    exec_id = result["Id"]
    async with docker._query(
        f"exec/{exec_id}/start",
        method="POST",
        headers={"content-type": "application/json"},
        data=dump_json_str({
            "Stream": False,  # get response immediately
            "Detach": False,
            "Tty": False,
        }),
    ) as resp:
        result = await resp.read()
        result = result.decode("latin-1").strip().split("\n")
    result = list(map(lambda x: x.split(), result))
    head = result[0]
    procs = result[1:]
    pid_idx, cmd_idx = head.index("PID"), head.index("COMMAND")
    return [ProcessInfo(int(r[pid_idx]), " ".join(r[cmd_idx:])) for r in procs]


async def host_pid_to_container_pid(container_id: str, host_pid: HostPID) -> ContainerPID:
    kernel_ver = Path("/proc/version").read_text()
    if m := re.match(r"Linux version (\d+)\.(\d+)\..*", kernel_ver):
        kernel_ver_tuple: tuple[str, str] = m.groups()  # type: ignore
        if kernel_ver_tuple < ("4", "1"):
            # TODO: this should be deprecated when the minimun supported Linux kernel will be 4.1.
            #
            # In CentOs 7, NSPid is not accesible since it is supported from Linux kernel >=4.1.
            # We provide alternative, although messy, way for older Linux kernels. Below describes
            # the logic briefly:
            #   * Obtain information on all the processes inside the target container,
            #     which contains host PID, by docker top API (containers/<container-id>/top).
            #     - Get the COMMAND of the target process (by using host_pid).
            #     - Filter host processes which have the exact same COMMAND.
            #   * Obtain information on all the processes inside the target container,
            #     which contains container PID, by executing "ps -aux" command from inside the container.
            #     - Filter container processes which have the exact same COMMAND.
            #   * Get the index of the target process from the host process table.
            #   * Use the index to get the target process from the container process table, and get PID.
            #     - Since docker top and ps -aux both displays processes in the order of PID, we
            #       can safely assume that the order of the processes from both tables are the same.
            #
            # Example host and container process table:
            #
            # [
            #   ['devops', '15454', '12942', '99', '15:36', 'pts/1', '00:00:08', 'python mnist.py'],
            #   ... (processes with the same COMMAND)
            # ]
            #
            # [
            #   ['work', '227', '121', '4.6', '22408680', '1525428', 'pts/1', 'Rl+', '06:36', '0:08',
            #    'python', 'mnist.py'],
            #   ... (processes with the same COMMAND)
            # ]
            docker: aiodocker.Docker | None = None
            try:
                docker = aiodocker.Docker()
                # Get process table from host (docker top information). Filter processes which have
                # exactly the same COMMAND as with target host process.
                host_procs = await get_host_process_table(docker, container_id)
                cmd = [x.command for x in host_procs if x.pid == int(host_pid)][0]
                host_pids = [x.pid for x in host_procs if x.command == cmd]

                # When there are multiple processes which have the same COMMAND, just get the index of
                # the target host process and apply it with the container table. Since ps and docker top
                # both displays processes ordered by PID, we can expect those two tables have same
                # order of processes.
                process_idx = host_pids.index(int(host_pid))

                container_procs = await get_container_process_table(docker, container_id)
                container_pids = [x.pid for x in container_procs if x.command == cmd]

                container_pid = ContainerPID(PID(container_pids[process_idx]))
                log.debug("host pid {} is mapped to container pid {}", host_pid, container_pid)
                return container_pid
            except asyncio.CancelledError:
                raise
            except (IndexError, KeyError, aiodocker.exceptions.DockerError):
                return NotContainerPID
            finally:
                if docker is not None:
                    await docker.close()

    try:
        cgroup = get_cgroup_of_pid("pids", host_pid)
        cgroup_container_id = get_container_id_of_cgroup(cgroup)
        if cgroup_container_id is None:
            return NotContainerPID
        if cgroup_container_id == container_id:
            for line in Path(f"/proc/{host_pid}/status").read_text().splitlines():
                key, value = line.split(":\t", 1)
                if key == "NSpid":
                    pid = value.split()[1]
                    return ContainerPID(PID(int(pid)))
        return InOtherContainerPID
    except OSError:
        return NotContainerPID


async def container_pid_to_host_pid(container_id: str, container_pid: ContainerPID) -> HostPID:
    kernel_ver = Path("/proc/version").read_text()
    if m := re.match(r"Linux version (\d+)\.(\d+)\..*", kernel_ver):
        kernel_ver_tuple: tuple[str, str] = m.groups()  # type: ignore
        if kernel_ver_tuple < ("4", "1"):
            # reverse implementation of host_pid_to_container_pid().
            docker: aiodocker.Docker | None = None
            try:
                docker = aiodocker.Docker()
                container_procs = await get_container_process_table(docker, container_id)
                cmd = [x.command for x in container_procs if x.pid == int(container_pid)][0]
                container_pids = [x.pid for x in container_procs if x.command == cmd]
                process_idx = container_pids.index(int(container_pid))

                host_procs = await get_host_process_table(docker, container_id)
                host_pids = [x.pid for x in host_procs if x.command == cmd]

                host_pid = HostPID(PID(host_pids[process_idx]))
                log.debug("container pid {} is mapped to host pid {}", container_pid, host_pid)
                return host_pid
            except asyncio.CancelledError:
                raise
            except (IndexError, KeyError, aiodocker.exceptions.DockerError):
                return NotHostPID
            finally:
                if docker is not None:
                    await docker.close()

    try:
        cgtasks = await get_container_pids(ContainerId(container_id))
        for pid in cgtasks:
            proc_path = Path(f"/proc/{pid}/status")
            proc_status = {
                k: v
                for k, v in map(lambda line: line.split(":\t"), proc_path.read_text().splitlines())
            }
            nspids = proc_status["NSpid"].split()
            if nspids[1] == str(container_pid):
                return HostPID(PID(pid))
        return NotHostPID
    except (OSError, ValueError, KeyError):
        return NotHostPID


class DockerUlimitEntry(TypedDict):
    Name: str
    Soft: int
    Hard: int


def get_safe_ulimit(name: str, desired_soft: int, desired_hard: int) -> DockerUlimitEntry:
    """
    Get ulimit values that don't exceed the host's system limits to avoid
    'operation not permitted' errors from Docker.
    """
    try:
        # Map ulimit names to resource constants
        ulimit_map = {
            "nofile": resource.RLIMIT_NOFILE,
            "memlock": resource.RLIMIT_MEMLOCK,
            "nproc": resource.RLIMIT_NPROC,
            "fsize": resource.RLIMIT_FSIZE,
            "core": resource.RLIMIT_CORE,
            "stack": resource.RLIMIT_STACK,
            "data": resource.RLIMIT_DATA,
            "rss": resource.RLIMIT_RSS,
        }

        if name not in ulimit_map:
            # If we don't know how to check this limit, use desired values
            return {"Name": name, "Soft": desired_soft, "Hard": desired_hard}

        current_soft_limit, current_hard_limit = resource.getrlimit(ulimit_map[name])
        infinity_val = resource.RLIM_INFINITY

        if desired_soft == -1:
            safe_soft = current_soft_limit if current_soft_limit != infinity_val else -1
        else:
            if current_soft_limit == infinity_val:
                safe_soft = desired_soft
            else:
                safe_soft = min(desired_soft, current_soft_limit)

        if desired_hard == -1:
            safe_hard = current_hard_limit if current_hard_limit != infinity_val else -1
        else:
            if current_hard_limit == infinity_val:
                safe_hard = desired_hard
            else:
                safe_hard = min(desired_hard, current_hard_limit)

        # Log if we had to adjust the limits
        if safe_soft != desired_soft or safe_hard != desired_hard:
            log.debug(
                f"Adjusted ulimit {name}: desired=({desired_soft}, {desired_hard}) "
                f"-> safe=({safe_soft}, {safe_hard}) (system limits: {current_soft_limit}, {current_hard_limit})"
            )

        return {"Name": name, "Soft": safe_soft, "Hard": safe_hard}

    except (OSError, ValueError, AttributeError) as e:
        # If we can't get system limits, log and use desired values
        log.warning(f"Could not get system ulimit for {name}: {e}")
        return {"Name": name, "Soft": desired_soft, "Hard": desired_hard}
