import asyncio
from typing import Callable, AsyncGenerator
import grpc


class ConnectionManager:
    def __init__(
            self,
            communicator: BaseCommunicator,
            retry_policy: RetryPolicy,
            circuit_breaker: CircuitBreaker,
            max_reconnect_attempts: int | None = 5,  # None = бесконечные попытки
            reconnect_timeout: float = 30.0  # Таймаут для бесконечного реконнекта
    ):
        self.communicator = communicator
        self.retry_policy = retry_policy
        self.circuit_breaker = circuit_breaker
        self._max_reconnect_attempts = max_reconnect_attempts
        self.reconnect_timeout = reconnect_timeout
        self._reconnect_delay = 1
        self._is_connecting = False
        self._connection_lock = asyncio.Lock()

    async def _check_connection(self) -> bool:
        try:
            await asyncio.wait_for(
                self.communicator.get_status(),
                timeout=1.0
            )
            return True
        except (grpc.RpcError, asyncio.TimeoutError, ConnectionError):
            return False
        except Exception as e:
            pass # todo
            return False

    async def _ensure_connection(self) -> bool:
        async with self._connection_lock:
            if self._is_connecting:
                while self._is_connecting:
                    await asyncio.sleep(0.1)

            if await self._check_connection():
                return True

            return await self._reconnect()

    async def _reconnect(self) -> bool:
        self._is_connecting = True

        try:
            if self._max_reconnect_attempts is None:
                return await self._reconnect_infinite()
            else:
                return await self._reconnect_limited()
        finally:
            self._is_connecting = False

    async def _reconnect_limited(self) -> bool:
        for attempt in range(self._max_reconnect_attempts):
            try:
                await self.communicator.disconnect()
                await self.communicator.connect()

                if await self._check_connection():
                    return True

            except Exception as e:
                print(f"Reconnect attempt {attempt + 1} failed: {e}") #  todo logger

            await asyncio.sleep(self._reconnect_delay)

        return False

    async def _reconnect_infinite(self) -> bool:
        start_time = asyncio.get_event_loop().time()
        attempt = 0

        while True:
            try:
                await self.communicator.disconnect()
                await self.communicator.connect()

                if await self._check_connection():
                    return True

            except Exception as e:
                print(f"Reconnect attempt {attempt + 1} failed: {e}") # todo logger

            elapsed = asyncio.get_event_loop().time() - start_time
            if elapsed > self.reconnect_timeout:
                return False

            attempt += 1
            await asyncio.sleep(self._reconnect_delay)

    async def execute_command(
            self,
            command_fn: Callable,
            *args, **kwargs
    ) -> Any:
        """Выполняет команду с обработкой ошибок и реконнектами"""
        async for attempt in self.retry_policy:
            try:
                if not await self._ensure_connection():
                    print("Failed to establish connection")
                    continue

                return await self.circuit_breaker.call(
                    command_fn, *args, **kwargs
                )

            except (grpc.RpcError, ConnectionError) as e:
                print(f"Command attempt {attempt} failed: {e}")
                await self._handle_error(e, attempt)

        raise ConnectionError("All command attempts failed")

    async def _handle_error(self, error: Exception, attempt: int):
        """Обрабатывает ошибки соединения"""
        # При определенных ошибках сразу разрываем соединение
        if isinstance(error, (
                grpc.aio.AioRpcError,
                ConnectionResetError,
                TimeoutError
        )):
            await self.communicator.disconnect()

        # Ждем перед следующей попыткой (уже учтено в retry policy)
        pass

    async def execute_stream_command(
            self,
            command_fn: Callable,
            *args, **kwargs
    ) -> AsyncGenerator:
        """Выполняет потоковую команду"""
        if not await self._ensure_connection():
            raise ConnectionError("No connection available")

        async for item in self.circuit_breaker.call_stream(
                command_fn, *args, **kwargs
        ):
            yield item