from __future__ import annotations

import asyncio
import inspect
import socket
import sys
from contextlib import asynccontextmanager, contextmanager
from functools import wraps
from typing import Any

import async_timeout

from .bfunc import to_any
from .btype import AsyncFuncType, Func

# ------------------------------------------------------------------------------------------------------------------------
# 根据端口限制多开


@contextmanager
def port_lock(port: int, err_msg: str = '程序禁止多开'):
    "占用端口方式实现多开限制 port=0 表示不限制"
    if port == 0:
        yield
    else:
        assert 60000 <= port <= 65535, '端口范围为 60000 ~ 65535'
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            sock.bind(('localhost', port))
            sock.listen(1)
        except Exception:
            raise Exception(err_msg)
        try:
            yield
        finally:
            sock.close()


# ------------------------------------------------------------------------------------------------------------------------
# 限制并执行数量


def limit(value: int = 1):
    def wraperfun(func: AsyncFuncType) -> AsyncFuncType:
        @wraps(func)
        async def wraper(*args: Any, **kwargs: Any):
            funid = id(inspect.unwrap(func))
            if funid not in _limit_dict:
                _limit_dict[funid] = _Limit(value)
            try:
                await _limit_dict[funid].wait()
                return await func(*args, **kwargs)
            finally:
                await _limit_dict[funid].release()

        return to_any(wraper)

    return wraperfun


async def set_limit(func: Func, limit: int):
    funid = id(inspect.unwrap(func))
    if funid not in _limit_dict:
        _limit_dict[funid] = _Limit(limit)
    await _limit_dict[funid].set_limit(limit)


_limit_dict: dict[int, _Limit] = {}


class _Limit:
    _queue: asyncio.Queue[Any]
    _running: int

    def __init__(self, limit: int):
        self._limit = limit
        self._queue = asyncio.Queue()
        self._running = 0
        while self._queue.qsize() < self._limit:
            self._queue.put_nowait(True)

    async def wait(self):
        await self._queue.get()
        self._running += 1

    async def release(self):
        if self._queue.qsize() < self._limit:
            await self._queue.put(True)
        self._running -= 1

    async def set_limit(self, limit: int):
        self._limit = limit
        while self._running + self._queue.qsize() < self._limit:
            await self._queue.put(True)
        while self._running + self._queue.qsize() > self._limit:
            if self._queue.empty():
                break
            await self._queue.get()


# ------------------------------------------------------------------------------------------------------------------------
# 超时锁


class TimeoutLock:
    def __init__(self, timeout: float | None = None) -> None:
        self._lock: asyncio.Lock = asyncio.Lock()
        self._timeout: float = sys.float_info.max if timeout is None else timeout

    async def acquire(self) -> None:
        async with async_timeout.timeout(self._timeout):
            await self._lock.acquire()

    def release(self) -> None:
        self._lock.release()

    def locked(self) -> bool:
        return self._lock.locked()

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

    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        self.release()


# ------------------------------------------------------------------------------------------------------------------------
# 读写锁


class RWLock:
    def __init__(self, timeout: float | None = None) -> None:
        timeout = sys.float_info.max if timeout is None else timeout
        self._lock: TimeoutLock = TimeoutLock(timeout)
        self._read_num: int = 0

    async def acquire_read(self) -> None:
        if self._read_num:
            self._read_num += 1
        else:
            await self._lock.acquire()
            self._read_num += 1

    async def release_read(self) -> None:
        self._read_num -= 1
        if self._read_num == 0:
            self._lock.release()

    async def acquire_write(self) -> None:
        await self._lock.acquire()

    async def release_write(self) -> None:
        self._lock.release()

    @asynccontextmanager
    async def use_read(self):
        await self.acquire_read()
        try:
            yield
        finally:
            await self.release_read()

    @asynccontextmanager
    async def use_write(self):
        await self.acquire_write()
        try:
            yield
        finally:
            await self.release_write()
