from __future__ import annotations

import asyncio
import json
from pathlib import Path, PurePosixPath
from typing import AsyncIterator, FrozenSet, Sequence
from uuid import UUID

from ai.backend.common.types import BinarySize

from ..abc import CAP_FAST_SCAN, CAP_METRIC, CAP_VFOLDER
from ..types import DirEntry, DirEntryType, FSPerfMetric, Stat, VFolderUsage
from ..utils import fstime2datetime
from ..vfs import BaseVolume
from .purity import PurityClient


class FlashBladeVolume(BaseVolume):
    async def init(self) -> None:
        available = True
        try:
            proc = await asyncio.create_subprocess_exec(
                b"pdu",
                b"--version",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
            )
        except FileNotFoundError:
            available = False
        else:
            try:
                stdout, stderr = await proc.communicate()
                if b"RapidFile Toolkit" not in stdout or proc.returncode != 0:
                    available = False
            finally:
                await proc.wait()
        if not available:
            raise RuntimeError(
                "PureStorage RapidFile Toolkit is not installed. "
                "You cannot use the PureStorage backend for the stroage proxy."
            )
        self.purity_client = PurityClient(
            self.config["purity_endpoint"],
            self.config["purity_api_token"],
            api_version=self.config["purity_api_version"],
        )

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

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

    async def copy_tree(
        self,
        src_vfpath: Path,
        dst_vfpath: Path,
    ) -> None:
        proc = await asyncio.create_subprocess_exec(
            b"pcp",
            b"-r",
            b"-p",
            bytes(src_vfpath / "."),
            bytes(dst_vfpath),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await proc.communicate()
        if proc.returncode != 0:
            raise RuntimeError(f'"pcp" command failed: {stderr.decode()}')

    async def get_quota(self, vfid: UUID) -> BinarySize:
        raise NotImplementedError

    async def set_quota(self, vfid: UUID, size_bytes: BinarySize) -> None:
        raise NotImplementedError

    async def get_performance_metric(self) -> FSPerfMetric:
        async with self.purity_client as client:
            async for item in client.get_nfs_metric(self.config["purity_fs_name"]):
                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 RuntimeError(
                    "no metric found for the configured flashblade filesystem"
                )

    async def get_usage(
        self, vfid: UUID, relpath: PurePosixPath = None
    ) -> VFolderUsage:
        target_path = self.sanitize_vfpath(vfid, relpath)
        total_size = 0
        total_count = 0
        raw_target_path = bytes(target_path)
        # Measure the exact file sizes and bytes
        proc = await asyncio.create_subprocess_exec(
            b"pdu",
            b"-0",
            b"-b",
            b"-a",
            b"-s",
            raw_target_path,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
        assert proc.stdout is not None
        try:
            # TODO: check slowdowns when there are millions of files
            while True:
                try:
                    line = await proc.stdout.readuntil(b"\0")
                    line = line.rstrip(b"\0")
                except asyncio.IncompleteReadError:
                    break
                size, name = line.split(maxsplit=1)
                if len(name) != len(raw_target_path) and name != raw_target_path:
                    total_size += int(size)
                    total_count += 1
        finally:
            await proc.wait()
        return VFolderUsage(file_count=total_count, used_bytes=total_size)

    async def get_used_bytes(self, vfid: UUID) -> BinarySize:
        vfpath = self.mangle_vfpath(vfid)
        proc = await asyncio.create_subprocess_exec(
            b"pdu",
            b"-hs",
            bytes(vfpath),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await proc.communicate()
        if proc.returncode != 0:
            raise RuntimeError(f"pdu command failed: {stderr.decode()}")
        used_bytes, _ = stdout.decode().split()
        return BinarySize.finite_from_str(used_bytes)

    # ------ vfolder internal operations -------

    def scandir(self, vfid: UUID, relpath: PurePosixPath) -> AsyncIterator[DirEntry]:
        target_path = self.sanitize_vfpath(vfid, relpath)
        raw_target_path = bytes(target_path)

        async def _aiter() -> AsyncIterator[DirEntry]:
            proc = await asyncio.create_subprocess_exec(
                b"pls",
                b"--json",
                raw_target_path,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
            )
            assert proc.stdout is not None
            try:
                while True:
                    line = await proc.stdout.readline()
                    if not line:
                        break
                    line = line.rstrip(b"\n")
                    item = json.loads(line)
                    item_path = Path(item["path"])
                    entry_type = DirEntryType.FILE
                    if item["filetype"] == 40000:
                        entry_type = DirEntryType.DIRECTORY
                    if item["filetype"] == 120000:
                        entry_type = DirEntryType.SYMLINK
                    yield DirEntry(
                        name=item_path.name,
                        path=item_path,
                        type=entry_type,
                        stat=Stat(
                            size=item["size"],
                            owner=str(item["uid"]),
                            # The integer represents the octal number in decimal
                            # (e.g., 644 which actually means 0o644)
                            mode=int(str(item["mode"]), 8),
                            modified=fstime2datetime(item["mtime"]),
                            created=fstime2datetime(item["ctime"]),
                        ),
                    )
            finally:
                await proc.wait()

        return _aiter()

    async def copy_file(
        self, vfid: UUID, src: PurePosixPath, dst: PurePosixPath
    ) -> None:
        src_path = self.sanitize_vfpath(vfid, src)
        dst_path = self.sanitize_vfpath(vfid, dst)
        proc = await asyncio.create_subprocess_exec(
            b"pcp",
            b"-p",
            bytes(src_path),
            bytes(dst_path),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await proc.communicate()
        if proc.returncode != 0:
            raise RuntimeError(f'"pcp" command failed: {stderr.decode()}')

    async def delete_files(
        self, vfid: UUID, relpaths: Sequence[PurePosixPath], recursive: bool = False
    ) -> None:
        target_paths = [bytes(self.sanitize_vfpath(vfid, p)) for p in relpaths]
        proc = await asyncio.create_subprocess_exec(
            b"prm",
            *target_paths,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        await proc.communicate()
        if proc.returncode != 0:
            raise RuntimeError("'prm' command returned a non-zero exit code.")
