import asyncio
import os
import shutil
from pathlib import Path
from typing import Dict, FrozenSet, List

import aiofiles.os

from ai.backend.common.types import BinarySize

from ..abc import CAP_FAST_SIZE, CAP_QUOTA, CAP_VFOLDER, AbstractFSOpModel, AbstractQuotaModel
from ..subproc import run
from ..types import CapacityUsage, Optional, QuotaConfig, QuotaUsage, TreeUsage
from ..vfs import BaseFSOpModel, BaseQuotaModel, BaseVolume


class CephDirQuotaModel(BaseQuotaModel):
    async def create_quota_scope(
        self,
        quota_scope_id: str,
        config: Optional[QuotaConfig] = None,
    ) -> None:
        qspath = self.mangle_qspath(quota_scope_id)
        await aiofiles.os.makedirs(qspath)
        if config is not None:
            await self.update_quota_scope(quota_scope_id, config)

    async def describe_quota_scope(self, quota_scope_id: str) -> QuotaUsage:
        qspath = self.mangle_qspath(quota_scope_id)
        loop = asyncio.get_running_loop()

        def read_attrs() -> tuple[int, int]:
            used_bytes = int(os.getxattr(qspath, "ceph.dir.rbytes").decode())  # type: ignore[attr-defined]
            try:
                limit_bytes = int(os.getxattr(qspath, "ceph.quota.max_bytes").decode())  # type: ignore[attr-defined]
            except OSError as e:
                match e.errno:
                    case 61:
                        limit_bytes = 0
                    case _:
                        limit_bytes = -1  # unset
            return used_bytes, limit_bytes

        # without type: ignore mypy will raise error when trying to run on macOS
        # because os.getxattr() exists only for linux
        used_bytes, limit_bytes = await loop.run_in_executor(
            None,
            read_attrs,
        )
        return QuotaUsage(used_bytes=used_bytes, limit_bytes=limit_bytes)

    async def update_quota_scope(
        self,
        quota_scope_id: str,
        config: QuotaConfig,
    ) -> None:
        qspath = self.mangle_qspath(quota_scope_id)
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(
            None,
            # without type: ignore mypy will raise error when trying to run on macOS
            # because os.setxattr() exists only for linux
            lambda: os.setxattr(qspath, "ceph.quota.max_bytes", str(int(config.limit_bytes)).encode()),  # type: ignore[attr-defined]
        )


class CephFSOpModel(BaseFSOpModel):
    async def scan_tree_usage(self, target_path: Path) -> TreeUsage:
        loop = asyncio.get_running_loop()
        raw_reports = await loop.run_in_executor(
            None,
            lambda: (
                os.getxattr(target_path, "ceph.dir.rentries"),  # type: ignore[attr-defined]
                os.getxattr(target_path, "ceph.dir.rbytes"),  # type: ignore[attr-defined]
            ),
        )
        file_count = int(raw_reports[0].strip().decode())
        used_bytes = int(raw_reports[1].strip().decode())
        return TreeUsage(file_count=file_count, used_bytes=used_bytes)

    async def scan_tree_size(self, target_path: Path) -> BinarySize:
        loop = asyncio.get_running_loop()
        raw_report = await loop.run_in_executor(
            None,
            lambda: os.getxattr(target_path, "ceph.dir.rbytes"),  # type: ignore[attr-defined]
        )
        return BinarySize(raw_report.strip().decode())


class CephFSVolume(BaseVolume):
    loop: asyncio.AbstractEventLoop
    registry: Dict[str, int]
    project_id_pool: List[int]

    async def init(self) -> None:
        try:
            await run([b"ceph", b"--version"])
        except FileNotFoundError:
            raise RuntimeError("Ceph is not installed. ")
        await super().init()

    async def create_quota_model(self) -> AbstractQuotaModel:
        return CephDirQuotaModel(self.mount_path)

    async def create_fsop_model(self) -> AbstractFSOpModel:
        return CephFSOpModel(
            self.mount_path,
            self.local_config["storage-proxy"]["scandir-limit"],
        )

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

    async def get_fs_usage(self) -> CapacityUsage:
        (total, used, _) = await asyncio.get_running_loop().run_in_executor(
            None,
            shutil.disk_usage,
            self.mount_path,
        )
        return CapacityUsage(
            used_bytes=used,
            capacity_bytes=total,
        )
