import logging
from collections.abc import Mapping, MutableMapping, Sequence
from decimal import Decimal
from pathlib import Path
from pprint import pformat
from typing import (
    Any,
)

import aiodocker

from . import __version__, lpu_api
from .types import LPUDevice

try:
    from ai.backend.agent.resources import get_resource_spec_from_container  # type: ignore
except ImportError:
    from ai.backend.agent.docker.resources import get_resource_spec_from_container

from ai.backend.agent.resources import (
    AbstractAllocMap,
    AbstractComputePlugin,
    DeviceSlotInfo,
    DiscretePropertyAllocMap,
)
from ai.backend.agent.stats import (
    ContainerMeasurement,
    NodeMeasurement,
    ProcessMeasurement,
    StatContext,
)
from ai.backend.agent.types import Container, MountInfo
from ai.backend.common.logging import BraceStyleAdapter
from ai.backend.common.types import (
    AcceleratorMetadata,
    BinarySize,
    DeviceId,
    DeviceModelInfo,
    DeviceName,
    MountTypes,
    SlotName,
    SlotTypes,
)

PREFIX = "hyperaccel-lpu"
SLOT_NAME = "hyperaccel-lpu.device"


log = BraceStyleAdapter(logging.getLogger(__spec__.name))  # type: ignore


class LPUPlugin(AbstractComputePlugin):
    key = DeviceName(PREFIX)
    slot_types: Sequence[tuple[SlotName, SlotTypes]] = ((SlotName(SLOT_NAME), SlotTypes("count")),)
    exclusive_slot_types: set[str] = {SLOT_NAME}

    device_mask: Sequence[DeviceId] = []
    enabled: bool = True

    _all_devices: list[LPUDevice] | None

    async def init(self, context: Any = None) -> None:
        self._all_devices = None

        raw_device_mask = self.plugin_config.get("device_mask")
        if raw_device_mask is not None:
            self.device_mask = [
                *map(lambda dev_id: DeviceId(dev_id), raw_device_mask.split(",")),
            ]

        try:
            detected_devices = await self.list_devices()
            log.info("detected devices:\n" + pformat(detected_devices))
            log.info("LPU acceleration is enabled.")
        except ImportError:
            log.warning("could not find LPU.")
            self.enabled = False

    async def list_devices(self) -> list[LPUDevice]:
        if self._all_devices is not None:
            return self._all_devices
        self._all_devices = await lpu_api.list_devices()
        return self._all_devices

    async def available_slots(self) -> Mapping[SlotName, Decimal]:
        devices = await self.list_devices()
        log.debug("available devices: {}", Decimal(len(devices)))
        return {
            SlotName(SLOT_NAME): Decimal(len(devices)),
        }

    def get_version(self) -> str:
        return __version__

    async def extra_info(self) -> Mapping[str, Any]:
        if self.enabled:
            return {
                "lpu_support": True,
            }
        return {}

    async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]:
        # TODO: implement
        # return [
        #     NodeMeasurement(
        #         MetricKey("lpu_mem"),
        #         MetricTypes.USAGE,
        #         unit_hint="bytes",
        #         stats_filter=frozenset({"max"}),
        #         per_node=Measurement(Decimal(mem_used_total), Decimal(mem_avail_total)),
        #         per_device=mem_stats,
        #     ),
        #     NodeMeasurement(
        #         MetricKey("lpu_util"),
        #         MetricTypes.USAGE,
        #         unit_hint="percent",
        #         stats_filter=frozenset({"avg", "max"}),
        #         per_node=Measurement(Decimal(util_total), Decimal(dev_count * 100)),
        #         per_device=util_stats,
        #     ),
        # ]
        return []

    async def gather_container_measures(
        self,
        ctx: StatContext,
        container_ids: Sequence[str],
    ) -> Sequence[ContainerMeasurement]:
        # TODO: implement
        # return [
        #     ContainerMeasurement(
        #         MetricKey("lpu_mem"),
        #         MetricTypes.USAGE,
        #         unit_hint="bytes",
        #         stats_filter=frozenset({"max"}),
        #         per_container={
        #             cid: Measurement(
        #                 Decimal(usage),
        #                 Decimal(mem_sizes[cid]),
        #             )
        #             for cid, usage in mem_stats.items()
        #         },
        #     ),
        #     ContainerMeasurement(
        #         MetricKey("lpu_util"),
        #         MetricTypes.USAGE,
        #         unit_hint="percent",
        #         stats_filter=frozenset({"avg", "max"}),
        #         per_container={
        #             cid: Measurement(
        #                 util,
        #                 Decimal(number_of_devices_per_container[cid] * 100),
        #             )
        #             for cid, util in util_stats.items()
        #         },
        #     ),
        # ]
        return []

    async def create_alloc_map(self) -> DiscretePropertyAllocMap:
        devices = await self.list_devices()
        return DiscretePropertyAllocMap(
            device_slots={
                dev.device_id: (DeviceSlotInfo(SlotTypes.COUNT, SlotName(SLOT_NAME), Decimal(1)))
                for dev in devices
            },
            exclusive_slot_types=self.exclusive_slot_types,
        )

    async def generate_mounts(
        self,
        source_path: Path,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> list[MountInfo]:
        libpath = Path("/opt/hyperdex")
        return [
            MountInfo(MountTypes.BIND, libpath, libpath),
        ]

    async def generate_docker_args(
        self,
        docker: aiodocker.docker.Docker,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> Mapping[str, Any]:
        assigned_devices: list[str] = []
        for dev in await self.list_devices():
            if dev.device_id in device_alloc.get(SlotName(SLOT_NAME), {}).keys():
                assigned_devices.extend(
                    [dev.xvc_pri_path.as_posix(), dev.renderD_path.as_posix()],
                )

        return {
            "HostConfig": {
                "CapAdd": ["IPC_LOCK"],
                "IpcMode": "host",
                "Ulimits": [
                    {
                        "Name": "memlock",
                        "Hard": -1,
                        "Soft": -1,
                    },
                ],
                "Sysctls": {
                    "net.ipv6.conf.all.disable_ipv6": "0",
                },
                "Devices": [
                    {
                        "PathOnHost": dev,
                        "PathInContainer": dev,
                        "CgroupPermissions": "rwm",
                    }
                    for dev in assigned_devices
                ],
            },
        }

    async def get_attached_devices(
        self,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> Sequence[DeviceModelInfo]:
        device_ids: list[DeviceId] = []
        if SlotName(SLOT_NAME) in device_alloc:
            device_ids.extend(device_alloc[SlotName(SLOT_NAME)].keys())
        available_devices = await self.list_devices()
        attached_devices: list[DeviceModelInfo] = []
        for device in available_devices:
            if device.device_id in device_ids:
                proc = device.processing_units
                mem = BinarySize(device.memory_size)
                attached_devices.append(
                    {  # TODO: update common.types.DeviceModelInfo
                        "device_id": device.device_id,
                        "model_name": device.model_name,
                        "data": {
                            "smp": proc,
                            "mem": mem,
                        },
                    },
                )
        return attached_devices

    async def restore_from_container(
        self,
        container: Container,
        alloc_map: AbstractAllocMap,
    ) -> None:
        if not self.enabled:
            return
        resource_spec = await get_resource_spec_from_container(container.backend_obj)
        if resource_spec is None:
            return
        if hasattr(alloc_map, "apply_allocation"):
            for slot_name, _ in self.slot_types:
                alloc_map.apply_allocation(
                    {
                        slot_name: resource_spec.allocations.get(
                            DeviceName("lpu"),
                            {},
                        ).get(
                            slot_name,
                            {
                                dev_id: Decimal(0)
                                for dev_id, dev_slot_info in alloc_map.device_slots.items()
                                if dev_slot_info.slot_name == slot_name
                            },
                        ),
                    },
                )
        else:  # older agents without lablup/backend.ai-agent#180
            alloc_map.allocations[SlotName(SLOT_NAME)].update(
                resource_spec.allocations.get(
                    DeviceName("lpu"),
                    {},
                ).get(
                    SlotName(SLOT_NAME),
                    {},
                ),
            )

    async def generate_resource_data(
        self,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> Mapping[str, str]:
        data: MutableMapping[str, str] = {}
        if not self.enabled:
            return data

        active_device_id_set: set[DeviceId] = set()
        for slot_type, per_device_alloc in device_alloc.items():
            for dev_id, alloc in per_device_alloc.items():
                if alloc > 0:
                    active_device_id_set.add(dev_id)
        active_device_ids = sorted(active_device_id_set, key=lambda v: int(v))
        data["LPU_GLOBAL_DEVICE_IDS"] = ",".join(
            f"{local_idx}:{global_id}" for local_idx, global_id in enumerate(active_device_ids)
        )
        return data

    async def cleanup(self) -> None:
        pass

    async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None:
        pass

    async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]:
        return []

    async def get_docker_networks(
        self,
        device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]],
    ) -> list[str]:
        return []

    async def gather_process_measures(
        self,
        ctx: StatContext,
        pid_map: Mapping[int, str],
    ) -> Sequence[ProcessMeasurement]:
        return []

    def get_metadata(self) -> AcceleratorMetadata:
        return {
            "slot_name": PREFIX,
            "description": "Hyperaccel LPU",
            "human_readable_name": "Hyperaccel LPU Device",
            "display_unit": "Core",
            "number_format": {"binary": False, "round_length": 0},
            "display_icon": "npu5",
        }
