from __future__ import annotations

import asyncio
import logging
import os
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, FrozenSet, Mapping, Optional

import aiofiles.os

from ai.backend.common.etcd import AsyncEtcd
from ai.backend.common.events.dispatcher import EventDispatcher, EventProducer
from ai.backend.common.json import dump_json_str
from ai.backend.common.types import HardwareMetadata, QuotaConfig, QuotaScopeID
from ai.backend.logging import BraceStyleAdapter

from ...errors import VolumeNotInitializedError
from ...types import CapacityUsage, FSPerfMetric, QuotaUsage
from ...watcher import WatcherClient
from ..abc import CAP_FAST_FS_SIZE, CAP_METRIC, CAP_QUOTA, CAP_VFOLDER, AbstractQuotaModel
from ..vfs import BaseQuotaModel, BaseVolume
from .exceptions import WekaAPIError, WekaInitError, WekaNoMetricError, WekaNotFoundError
from .weka_client import WekaAPIClient

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


class WekaQuotaModel(BaseQuotaModel):
    def __init__(
        self,
        mount_path: Path,
        fs_uid: str,
        api_client: WekaAPIClient,
    ) -> None:
        super().__init__(mount_path)
        self.fs_uid = fs_uid
        self.api_client = api_client

    async def _get_inode_id(self, path: Path) -> int:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None,
            lambda: os.stat(path).st_ino,
        )

    async def create_quota_scope(
        self,
        quota_scope_id: QuotaScopeID,
        options: Optional[QuotaConfig] = None,
        extra_args: Optional[dict[str, Any]] = None,
    ) -> None:
        qspath = self.mangle_qspath(quota_scope_id)
        await aiofiles.os.makedirs(qspath)
        if self.fs_uid is None:
            raise VolumeNotInitializedError("Weka filesystem UID is not initialized")
        if options is not None:
            await self.update_quota_scope(quota_scope_id, options)

    async def update_quota_scope(self, quota_scope_id: QuotaScopeID, config: QuotaConfig) -> None:
        qspath = self.mangle_qspath(quota_scope_id)
        inode_id = await self._get_inode_id(qspath)
        qs_relpath = qspath.relative_to(self.mount_path).as_posix()
        if not qs_relpath.startswith("/"):
            qs_relpath = "/" + qs_relpath
        await self.api_client.set_quota_v1(
            qs_relpath, inode_id, soft_limit=config.limit_bytes, hard_limit=config.limit_bytes
        )

    async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> Optional[QuotaUsage]:
        qspath = self.mangle_qspath(quota_scope_id)
        if not qspath.exists():
            return None

        inode_id = await self._get_inode_id(qspath)
        quota = await self.api_client.get_quota(self.fs_uid, inode_id)
        used_bytes = quota.used_bytes if quota.used_bytes is not None else -1
        limit_bytes = quota.hard_limit if quota.hard_limit is not None else -1
        if used_bytes < 0 or limit_bytes < 0:
            log.warning(
                "Data from Weka API negative values in used_bytes({}) or limit_bytes({}) for quota scope {}: response from Weka API = {}",
                used_bytes,
                limit_bytes,
                quota_scope_id,
                quota.to_json(),
            )
        return QuotaUsage(
            used_bytes=used_bytes,
            limit_bytes=limit_bytes,
        )

    async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None:
        qspath = self.mangle_qspath(quota_scope_id)
        inode_id = await self._get_inode_id(qspath)
        try:
            await self.api_client.remove_quota(self.fs_uid, inode_id)
        except WekaNotFoundError:
            pass

    async def delete_quota_scope(self, quota_scope_id: QuotaScopeID) -> None:
        qspath = self.mangle_qspath(quota_scope_id)
        inode_id = await self._get_inode_id(qspath)
        try:
            await self.api_client.remove_quota(self.fs_uid, inode_id)
        except WekaNotFoundError:
            pass
        await aiofiles.os.rmdir(qspath)


class WekaVolume(BaseVolume):
    api_client: WekaAPIClient

    name = "weka"

    _fs_uid: str

    def __init__(
        self,
        local_config: Mapping[str, Any],
        mount_path: Path,
        *,
        etcd: AsyncEtcd,
        event_dispatcher: EventDispatcher,
        event_producer: EventProducer,
        watcher: Optional[WatcherClient] = None,
        options: Optional[Mapping[str, Any]] = None,
    ) -> None:
        super().__init__(
            local_config,
            mount_path,
            etcd=etcd,
            options=options,
            event_dispatcher=event_dispatcher,
            event_producer=event_producer,
            watcher=watcher,
        )
        ssl_verify = self.config.get("weka_verify_ssl", False)
        self.api_client = WekaAPIClient(
            self.config["weka_endpoint"],
            self.config["weka_username"],
            self.config["weka_password"],
            self.config["weka_organization"],
            ssl=ssl_verify,
        )

    async def init(self) -> None:
        for fs in await self.api_client.list_fs():
            if fs.name == self.config["weka_fs_name"]:
                self._fs_uid = fs.uid
                break
        else:
            raise WekaInitError(f"FileSystem {self.config['weka_fs_name']} not found")
        await super().init()

    async def create_quota_model(self) -> AbstractQuotaModel:
        return WekaQuotaModel(self.mount_path, self._fs_uid, self.api_client)

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

    async def get_hwinfo(self) -> HardwareMetadata:
        if self._fs_uid is None:
            raise VolumeNotInitializedError("Weka filesystem UID is not initialized")
        health_status = (await self.api_client.check_health()).lower()
        if health_status == "ok":
            health_status = "healthy"
        try:
            cluster_info = await self.api_client.get_cluster_info()
            quotas = await self.api_client.list_quotas(self._fs_uid)
            return {
                "status": health_status,
                "status_info": None,
                "metadata": {
                    "quota": dump_json_str([q.to_json() for q in quotas]),
                    "cluster_info": dump_json_str(cluster_info),
                },
            }
        except WekaAPIError:
            return {
                "status": health_status,
                "status_info": None,
                "metadata": {},
            }

    async def get_fs_usage(self) -> CapacityUsage:
        if self._fs_uid is None:
            raise VolumeNotInitializedError("Weka filesystem UID is not initialized")
        fs = await self.api_client.get_fs(self._fs_uid)
        return CapacityUsage(
            capacity_bytes=fs.total_budget,
            used_bytes=fs.used_total,
        )

    async def get_performance_metric(self) -> FSPerfMetric:
        start_time = datetime.now().replace(second=0, microsecond=0) - timedelta(
            minutes=1,
        )

        try:
            metrics = await self.api_client.get_metric(
                [
                    "ops.READS",
                    "ops.WRITES",
                    "ops.READ_BYTES",
                    "ops.WRITE_BYTES",
                    "ops.READ_LATENCY",
                    "ops.WRITE_LATENCY",
                ],
                start_time,
            )
            latest_metric = metrics["ops"][-1]["stats"]
        except KeyError:
            raise WekaNoMetricError
        except IndexError:
            raise WekaNoMetricError
        return FSPerfMetric(
            iops_read=latest_metric["READS"],
            iops_write=latest_metric["WRITES"],
            io_bytes_read=latest_metric["READ_BYTES"],
            io_bytes_write=latest_metric["WRITE_BYTES"],
            io_usec_read=latest_metric["READ_LATENCY"] or 0,
            io_usec_write=latest_metric["WRITE_LATENCY"] or 0,
        )
