from __future__ import annotations

import asyncio
import contextlib
import logging
import re
from typing import FrozenSet

from ai.backend.common.types import HardwareMetadata
from ai.backend.logging import BraceStyleAdapter

from ...errors import MetricNotFoundError, PureStorageCommandFailedError, VolumeNotInitializedError
from ...types import CapacityUsage, FSPerfMetric
from ..abc import (
    CAP_FAST_FS_SIZE,
    CAP_FAST_SCAN,
    CAP_METRIC,
    CAP_VFOLDER,
    AbstractFSOpModel,
)
from ..vfs import BaseVolume
from .purity import PurityClient
from .rapidfiles import RapidFileToolsFSOpModel
from .rapidfiles_v2 import RapidFileToolsv2FSOpModel

FLASHBLADE_TOOLKIT_V2_VERSION_RE = re.compile(r"version p[a-zA-Z\d]+ \(RapidFile\) (2\..+)")
FLASHBLADE_TOOLKIT_V1_VERSION_RE = re.compile(r"p[a-zA-Z\d]+ \(RapidFile Toolkit\) (1\..+)")

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


class FlashBladeVolume(BaseVolume):
    name = "purestorage"
    _toolkit_version: int | None

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self._toolkit_version = None

    async def create_fsop_model(self) -> AbstractFSOpModel:
        if (await self.get_toolkit_version()) == 2:
            return RapidFileToolsv2FSOpModel(
                self.mount_path,
                self.local_config["storage-proxy"]["scandir-limit"],
            )
        else:
            return RapidFileToolsFSOpModel(
                self.mount_path,
                self.local_config["storage-proxy"]["scandir-limit"],
            )

    async def get_toolkit_version(self) -> int:
        if self._toolkit_version is not None:
            return self._toolkit_version
        try:
            proc = await asyncio.create_subprocess_exec(
                b"pdu",
                b"--version",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
            )
        except FileNotFoundError:
            self._toolkit_version = -1
            return -1
        try:
            stdout, stderr = await proc.communicate()
            if proc.returncode != 0:
                self._toolkit_version = -1
            else:
                version_line = stdout.decode().splitlines()[0]
                if FLASHBLADE_TOOLKIT_V2_VERSION_RE.match(version_line):
                    self._toolkit_version = 2
                    log.info("FlashBlade Toolkit 2 detected")
                elif FLASHBLADE_TOOLKIT_V1_VERSION_RE.match(version_line):
                    self._toolkit_version = 1
                    log.info("FlashBlade Toolkit 1 detected")
                else:
                    log.warning("Unrecogized FlashBlade Toolkit version: {}", version_line)
                    self._toolkit_version = -1
        finally:
            await proc.wait()
            if not self._toolkit_version:
                raise VolumeNotInitializedError("Failed to detect FlashBlade Toolkit version")
            return self._toolkit_version

    async def init(self) -> None:
        toolkit_version = await self.get_toolkit_version()
        if toolkit_version == -1:
            raise PureStorageCommandFailedError(
                "PureStorage RapidFile Toolkit is not installed. "
                "You cannot use the PureStorage backend for the storage proxy."
            )
        self.purity_client = PurityClient(
            self.config["purity_endpoint"],
            self.config["purity_api_token"],
            api_version=self.config["purity_api_version"],
        )
        await super().init()

    async def shutdown(self) -> None:
        await self.purity_client.aclose()

    async def get_capabilities(self) -> FrozenSet[str]:
        return frozenset(
            [
                CAP_FAST_FS_SIZE,
                CAP_VFOLDER,
                CAP_METRIC,
                CAP_FAST_SCAN,
            ],
        )

    async def get_hwinfo(self) -> HardwareMetadata:
        async with self.purity_client as client:
            metadata = await client.get_metadata()
        return {
            "status": "healthy",
            "status_info": None,
            "metadata": {
                **metadata,
            },
        }

    async def get_fs_usage(self) -> CapacityUsage:
        async with self.purity_client as client:
            usage = await client.get_usage(self.config["purity_fs_name"])
        return CapacityUsage(
            capacity_bytes=usage["capacity_bytes"],
            used_bytes=usage["used_bytes"],
        )

    async def get_performance_metric(self) -> FSPerfMetric:
        async with self.purity_client as client:
            async with contextlib.aclosing(
                client.get_nfs_metric(self.config["purity_fs_name"]),
            ) as items:
                async for item in items:
                    return FSPerfMetric(
                        iops_read=item["reads_per_sec"],
                        iops_write=item["writes_per_sec"],
                        io_bytes_read=item["read_bytes_per_sec"],
                        io_bytes_write=item["write_bytes_per_sec"],
                        io_usec_read=item["usec_per_read_op"],
                        io_usec_write=item["usec_per_write_op"],
                    )
                else:
                    raise MetricNotFoundError(
                        "no metric found for the configured flashblade filesystem",
                    )
