"""
Protected module containing the AsyncPoolClient class
"""

import asyncio
import logging
import random
import weakref
from asyncio import Queue, QueueEmpty, Task
from contextlib import asynccontextmanager
from typing import Optional

from ecmind_blue_client.pool._abstract_pool_client import AbstractPoolClient
from ecmind_blue_client.pool._server_connection_settings import ServerConnectionSettings
from ecmind_blue_client.rpc import (
    BlueException,
    JobParameter,
    JobRequestFile,
    JobResult,
    Jobs,
    JobValueTypes,
)
from ecmind_blue_client.rpc.api import append_inferred_params
from ecmind_blue_client.rpc.async_api import (
    RpcAsyncConnection,
    async_call_job,
    async_connect,
)

logger = logging.getLogger(__name__)


class AsyncPoolClient(AbstractPoolClient):
    """
    asyncio Implementation of the AbstractPoolClient
    """

    def __init__(
        self,
        *,
        servers: list[ServerConnectionSettings] | str,
        username: str,
        password: str,
        name: Optional[str] = None,
        use_ssl: bool = True,
        file_cache_byte_limit: int = 33554432,
        pool_size: int = 10,
        connect_timeout: int = 10,
        pool_timeout: int = -1,
        cadata: str | None = None,
    ):
        super().__init__(
            servers=servers,
            username=username,
            password=password,
            name=name,
            use_ssl=use_ssl,
            file_cache_byte_limit=file_cache_byte_limit,
            pool_size=pool_size,
            connect_timeout=connect_timeout,
            pool_timeout=pool_timeout,
            cadata=cadata,
        )

        # pool of available connections
        self._connection_pool = Queue[RpcAsyncConnection]()
        self._lock = asyncio.Lock()

        self._cleanup_task: Task | None = None

    async def _borrow(self) -> RpcAsyncConnection:
        """
        Returns available connection from the pool or creates a new one if the pool is empty.

        Returns:
            AsyncConnection

        Exceptions:
            queue.Empty: If the pool is empty and a new connection cannot be created.
        """
        # Prevent concurrent calls
        async with self._lock:
            try:
                # try to get a element out of the queue without blocking
                connection = self._connection_pool.get_nowait()
                self._connections_in_use += 1
                return connection

            # If the pool is empty
            except QueueEmpty:
                # If the pool already reach its maximum size
                logger.debug(
                    "Pool is empty, creating a new connection use: %d pool: %d",
                    self._connections_in_use,
                    self._connection_pool_size,
                )

                if self._connections_in_use >= self._connection_pool_size:
                    # If timeout is defined, wait for defined time
                    if self._connection_pool_timout >= 0:
                        logger.debug("Waiting for a connection for %d seconds", self._connection_pool_timout)
                        connection_future = self._connection_pool.get()
                        connection = await asyncio.wait_for(connection_future, self._connection_pool_timout)
                    # Wait for a available for a invinity time
                    else:
                        logger.debug("Waiting for a connection")
                        connection = await self._connection_pool.get()
                else:
                    logger.debug("Create a new connection")
                    connection = await self._create()

                self._connections_in_use += 1
                return connection

    async def _create(self) -> RpcAsyncConnection:
        if self._cleanup_task is None:
            self._cleanup_task = asyncio.create_task(self._cleanup(weakref.ref(self)))

        servers = self._servers.copy()

        last_error: ConnectionError | TimeoutError | None = None
        connection: RpcAsyncConnection | None = None

        while len(servers) > 0:
            server = random.choices(servers, weights=[server.weight for server in servers], k=1)[0]
            servers.remove(server)

            try:
                logger.debug("Connecting to %s:%d", server.hostname, server.port)
                connection = await async_connect(
                    server.hostname, port=server.port, use_ssl=self._use_ssl, connect_timeout=self._connect_timeout, cadata=self._cadata
                )
                logger.debug("Successfully connected  to %s:%d", server.hostname, server.port)
                break
            except (ConnectionError, TimeoutError) as err:
                last_error = err
                logger.error("Failed to connect to %s:%d", server.hostname, server.port)

        # If no connection was established, raise an error
        if connection is None:
            logger.error("No valid enaio server found")
            raise last_error or ConnectionError("No valid enaio server found")

        logger.debug("Attaching new session")
        session_guid = self._session_attach_parse(
            connection, await self.__execute(connection, "krn.SessionAttach", self._session_attach_prepare())
        )
        logger.debug("Successfully attached new session with guid %s", session_guid)

        logger.debug("Set connection properties")
        self._session_properties_set_parse(
            await self.__execute(connection, "krn.SessionPropertiesSet", self._session_properties_set_prepare(self._name))
        )
        logger.debug("Successfully set connection properties")

        logger.debug("Login as %s", self._password)
        self._session_login_parse(
            connection,
            await self.__execute(connection, "krn.SessionLogin", self._session_login_prepare(self._username, self._password)),
            self._username,
        )
        logger.debug("Successfully logged as as %s", self._username)

        return connection

    async def _invalidate(self, connection: RpcAsyncConnection):
        try:
            connection.writer.close()
            await connection.writer.wait_closed()
        except OSError as ex:
            logger.warning(ex)
        finally:
            self._connections_in_use -= 1

    async def _return(self, connection: RpcAsyncConnection):
        await self._connection_pool.put(connection)
        self._connections_in_use -= 1

    async def _cleanup(self, pool_ref):
        while True:
            if pool_ref() is None:  # Check if the pool has been deleted
                # Small delay to ensure the loop is not too busy
                while self._connection_pool.qsize() > 0:

                    connection = await self._connection_pool.get()
                    try:
                        connection.writer.close()
                        await connection.writer.wait_closed()
                    except (RuntimeError, ConnectionError) as e:
                        logger.error("Error while closing connection: %s", e)
                break
            await asyncio.sleep(0.1)

    async def execute(
        self,
        job: Jobs | str,
        parameter: Optional[list[JobParameter]] = None,
        files: Optional[list[JobRequestFile]] = None,
        **params: JobValueTypes,
    ) -> JobResult:
        """
        Execute a job asynchronously.

        Args:
            job_name (str): The name of the job to execute.
            parameter (list[JobParameter]): The parameters for the job.
            files (Optional[list[JobRequestFile]], optional): The files for the job. Defaults to None.
            **params (JobValueTypes): Inferred parameters for the job. See parameters

        Returns:
            JobResult: The result of the job execution.

        Raises:
            BlueException: If the job fails to execute.

        """

        if parameter is None:
            parameter = []

        job_name = job.value if isinstance(job, Jobs) else job

        # append params to parameter
        append_inferred_params(parameter, params)

        async with self.connection as connection:
            return await self.__execute(connection, job_name, parameter, files)

    async def __execute(
        self, connection: RpcAsyncConnection, job_name: str, parameter: list[JobParameter], files: Optional[list[JobRequestFile]] = None
    ) -> JobResult:

        result = await async_call_job(connection, job_name, parameter, files)

        if result.return_code != 0:
            raise BlueException(result.return_code, f"{job_name} failed with result code {result.return_code}", result.errors)
        return result

    @property
    @asynccontextmanager
    async def connection(self):
        """
        Returns a connection as a context manager.

        The connection is borrowed from the pool and returned after the with block is exited.
        On ``BlueException`` (server-side error response) the TCP connection is still valid
        and is returned to the pool unchanged. On any other exception the connection is
        invalidated (socket closed, removed from the pool).
        """
        connection = await self._borrow()
        try:
            yield connection
        except BlueException:
            await self._return(connection)
            raise
        except Exception as e:
            logger.error("Exception in pool context manager: %s", e)
            await self._invalidate(connection)
            raise e
        await self._return(connection)
