from __future__ import annotations

import asyncio
import inspect
from collections.abc import Awaitable, Callable
from contextlib import suppress
from typing import Any

from proto_socket.communicator import Communicator
from proto_socket.packets.message_common_pb2 import HeartBeat


class _HbTimer:
    def __init__(self, delay: float, callback: Callable[[], Awaitable[None]]) -> None:
        self._task = asyncio.create_task(self._run(delay, callback))

    async def _run(self, delay: float, callback: Callable[[], Awaitable[None]]) -> None:
        await asyncio.sleep(delay)
        await callback()

    def cancel(self) -> None:
        self._task.cancel()


class BaseClient:
    def __init__(
        self,
        interval_sec: int,
        wait_sec: int,
        do_close: Callable[[], Awaitable[None]],
    ) -> None:
        self.communicator = Communicator()
        self._interval_sec = interval_sec
        self._wait_sec = wait_sec
        self._do_close = do_close
        self._close_called = False
        self._close_lock = asyncio.Lock()
        self._disconnect_listeners: list[Callable[[Any], Any]] = []
        self._hb_timer: _HbTimer | None = None
        self._waiting_hb_response = False

    def _init_base(self) -> None:
        self._close_lock = asyncio.Lock()

    async def close(self) -> None:
        await self._close_with_grace(True)

    async def _close_with_grace(self, graceful: bool) -> None:
        async with self._close_lock:
            if self._close_called:
                return
            self._close_called = True

        if graceful:
            await self.communicator.close()
        else:
            self.communicator.shutdown()

        self._stop_heartbeat()
        await self._do_close()
        await self._notify_disconnected()

    def add_disconnect_listener(self, fn: Callable[[Any], Any]) -> None:
        self._disconnect_listeners.append(fn)

    def remove_disconnect_listeners(self) -> None:
        self._disconnect_listeners.clear()

    async def _notify_disconnected(self) -> None:
        listeners = list(self._disconnect_listeners)
        self._disconnect_listeners.clear()
        for listener in listeners:
            result = listener(self)
            if inspect.isawaitable(result):
                await result

    async def _on_disconnected(self) -> None:
        await self._close_with_grace(False)

    async def _send_heartbeat(self) -> None:
        if not self.communicator.is_alive() or self._interval_sec <= 0:
            return
        self._stop_heartbeat()
        self._hb_timer = _HbTimer(self._interval_sec, self._heartbeat_interval_elapsed)

    async def _heartbeat_interval_elapsed(self) -> None:
        if not self.communicator.is_alive():
            return
        self._waiting_hb_response = True
        with suppress(Exception):
            await self.communicator.send(HeartBeat())
        self._stop_heartbeat()
        if self._wait_sec > 0:
            self._hb_timer = _HbTimer(self._wait_sec, self._heartbeat_wait_elapsed)

    async def _heartbeat_wait_elapsed(self) -> None:
        if self.communicator.is_alive():
            await self._on_disconnected()

    async def _on_heartbeat(self) -> None:
        if self._waiting_hb_response:
            self._waiting_hb_response = False
            return
        with suppress(Exception):
            await self.communicator.send(HeartBeat())

    def _stop_heartbeat(self) -> None:
        if self._hb_timer is not None:
            self._hb_timer.cancel()
            self._hb_timer = None