from __future__ import annotations

import abc
import asyncio
import fcntl
import logging
from collections.abc import Mapping
from io import IOBase
from pathlib import Path
from typing import Any, ClassVar

import trafaret as t
from etcd_client import Client as EtcdClient
from etcd_client import Communicator as EtcdCommunicator
from etcd_client import EtcdLockOption
from redis.asyncio import Redis
from redis.asyncio.lock import Lock as AsyncRedisLock
from redis.exceptions import LockError, LockNotOwnedError
from tenacity import (
    AsyncRetrying,
    RetryError,
    retry_if_exception_type,
    stop_after_delay,
    stop_never,
    wait_exponential,
    wait_random,
)

from ai.backend.logging import BraceStyleAdapter

from .etcd import AsyncEtcd
from .types import RedisConnectionInfo

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


class AbstractDistributedLock(metaclass=abc.ABCMeta):
    default_config: ClassVar[Mapping[str, Any]] = {}
    config_iv: ClassVar[t.Trafaret] = t.Dict().allow_extra("*")

    def __init__(self, *, lifetime: float | None = None) -> None:
        if lifetime is not None and lifetime < 0.0:
            raise ValueError("Lifetime must be non-negative")
        self._lifetime = lifetime

    @abc.abstractmethod
    async def __aenter__(self) -> Any:
        raise NotImplementedError

    @abc.abstractmethod
    async def __aexit__(self, *exc_info: Any) -> bool | None:
        raise NotImplementedError


class FileLock(AbstractDistributedLock):
    default_timeout: float = 3  # not allow infinite timeout for safety

    _file: IOBase | None
    _locked: bool = False

    def __init__(
        self,
        path: Path,
        *,
        timeout: float | None = None,
        lifetime: float | None = None,
        remove_when_unlock: bool = False,
        debug: bool = False,
    ) -> None:
        super().__init__(lifetime=lifetime)
        self._file = None
        self._path = path
        self._timeout = timeout if timeout is not None else self.default_timeout
        self._debug = debug
        self._remove_when_unlock = remove_when_unlock
        self._watchdog_task: asyncio.Task[Any] | None = None

    @property
    def locked(self) -> bool:
        return self._locked

    def __del__(self) -> None:
        if self._file is not None:
            self._debug = False
            self.release()
            log.debug("file lock implicitly released: {}", self._path)

    async def acquire(self) -> None:
        if self._file is not None:
            raise RuntimeError("File is already opened")
        if self._locked:
            raise RuntimeError("Lock is already acquired")
        if not self._path.exists():
            self._path.touch()
        self._file = self._path.open("wb")
        stop_func = stop_never if self._timeout <= 0 else stop_after_delay(self._timeout)
        try:
            async for attempt in AsyncRetrying(
                retry=retry_if_exception_type(BlockingIOError),
                wait=wait_exponential(multiplier=0.02, min=0.02, max=1.0) + wait_random(0, 0.05),
                stop=stop_func,
            ):
                with attempt:
                    fcntl.flock(self._file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
                    self._locked = True
                    if self._lifetime is not None:
                        self._watchdog_task = asyncio.create_task(
                            self._watchdog_timer(ttl=self._lifetime),
                        )
                    if self._debug:
                        log.debug("file lock acquired: {}", self._path)
        except RetryError as e:
            raise TimeoutError(f"failed to lock file: {self._path}") from e

    def release(self) -> None:
        if self._file is None:
            raise RuntimeError("No file to release")
        if task := self._watchdog_task:
            if not task.done():
                task.cancel()
        if self._locked:
            fcntl.flock(self._file.fileno(), fcntl.LOCK_UN)
            self._locked = False
            if self._debug:
                log.debug("file lock explicitly released: {}", self._path)
        self._file.close()
        if self._remove_when_unlock:
            try:
                self._path.unlink()
            except FileNotFoundError:
                pass
        self._file = None

    async def __aenter__(self) -> FileLock:
        await self.acquire()
        return self

    async def __aexit__(self, *exc_info: Any) -> bool | None:
        self.release()
        return None

    async def _watchdog_timer(self, ttl: float) -> None:
        await asyncio.sleep(ttl)
        if self._locked:
            if self._file is None:
                raise RuntimeError("File is not opened but lock is marked as acquired")
            fcntl.flock(self._file.fileno(), fcntl.LOCK_UN)
            self._locked = False
            if self._debug:
                log.debug(f"file lock implicitly released by watchdog: {self._path}")

    @property
    def is_locked(self) -> bool:
        return self._locked


class EtcdLock(AbstractDistributedLock):
    _etcd_client: EtcdClient | None
    _debug: bool

    lock_name: str
    etcd: AsyncEtcd
    timeout: float

    default_timeout: float = 9600  # not allow infinite timeout for safety

    def __init__(
        self,
        lock_name: str,
        etcd: AsyncEtcd,
        *,
        timeout: float | None = None,
        lifetime: float | None = None,
        debug: bool = False,
    ) -> None:
        super().__init__(lifetime=lifetime)
        self.lock_name = lock_name
        self.etcd = etcd
        self._timeout = timeout if timeout is not None else self.default_timeout
        self._debug = debug
        self._etcd_client = None

    async def __aenter__(self) -> EtcdCommunicator:
        self._etcd_client = self.etcd.etcd.with_lock(
            EtcdLockOption(
                lock_name=self.lock_name.encode("utf-8"),
                timeout=self._timeout,
                ttl=int(self._lifetime) if self._lifetime is not None else None,
            ),
        )

        etcd_communicator = await self._etcd_client.__aenter__()

        if self._debug:
            log.debug("etcd lock acquired")

        return etcd_communicator

    async def __aexit__(self, *exc_info: Any) -> bool | None:
        if self._etcd_client is None:
            raise RuntimeError("Etcd client is not initialized")
        await self._etcd_client.__aexit__(*exc_info)

        if self._debug:
            log.debug("etcd lock released")

        self._etcd_client = None
        return None


class RedisLock(AbstractDistributedLock):
    debug: bool
    _redis: Redis[Any]
    _timeout: float | None
    _lock: AsyncRedisLock | None

    default_timeout = 9600
    default_lock_retry_interval = 1.0

    default_config: ClassVar[Mapping[str, Any]] = {
        "lock_retry_interval": default_lock_retry_interval
    }
    config_iv: ClassVar[t.Trafaret] = t.Dict({
        t.Key("lock_retry_interval", default=None): t.Null | t.ToFloat[0:],
    }).allow_extra("*")

    def __init__(
        self,
        lock_name: str,
        redis: RedisConnectionInfo,
        *,
        timeout: float | None = None,
        lifetime: float | None = None,
        debug: bool = False,
        lock_retry_interval: float | None = None,
    ) -> None:
        super().__init__(lifetime=lifetime)
        self.lock_name = lock_name
        self._redis = redis.client
        self._timeout = timeout if timeout is not None else self.default_timeout
        self._debug = debug
        self._lock_retry_interval = (
            lock_retry_interval
            if lock_retry_interval is not None
            else self.default_lock_retry_interval
        )

    async def __aenter__(self) -> None:
        self._lock = AsyncRedisLock(
            self._redis,
            self.lock_name,
            blocking_timeout=self._timeout,
            timeout=self._lifetime,
            thread_local=False,
            sleep=self._lock_retry_interval,
        )
        try:
            await self._lock.__aenter__()
        except LockError as e:
            raise TimeoutError(str(e)) from e
        if self._debug:
            log.debug("RedisLock.__aenter__(): lock acquired")

    async def __aexit__(self, *exc_info: Any) -> bool | None:
        if self._lock is None:
            raise RuntimeError("Lock is not initialized")
        try:
            val = await self._lock.__aexit__(*exc_info)  # type: ignore[func-returns-value]
        except LockNotOwnedError:
            log.exception("Lock no longer owned. Skip.")
            return True
        except LockError:
            log.exception("Already unlocked. Skip.")
            return True
        if self._debug:
            log.debug("RedisLock.__aexit__(): lock released")

        return val
