# Copyright 2020 Oliver Cope
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any
from typing import Literal
from typing import Callable
from typing import Optional
from contextvars import ContextVar
import threading
import queue
import logging

from embrace import exceptions
from embrace import poolvalidators

logger = logging.getLogger(__name__)

ConnType = Any


class ContextManagerWrappedConnection:
    def __init__(self, pool: "ConnectionPoolBase", transaction=False, contextual=True):
        self.conn = None
        self.pool = pool
        self.transaction = transaction
        self.contextual = contextual

    def __enter__(self) -> ConnType:
        self.conn = self.pool.getconn(self.contextual)
        return self.conn

    def __exit__(self, exc_type, exc_value, tb):
        if self.transaction:
            try:
                if exc_type:
                    self.conn.rollback()
                else:
                    self.conn.commit()
            finally:
                self.pool.release(self.conn)
        else:
            self.pool.release(self.conn)


class ConnectionPoolBase:
    """
    Connection pool base class
    """

    #: A callable returning a db-api connection object
    connection_factory: Callable[[], ConnType]

    #: The class instantiated in response to the ``connect`` method
    connection_context_manager = ContextManagerWrappedConnection

    #: How many simultaneous connections to allow. If zero the number will be
    #: unlimited
    limit: int

    #: Return the same connection in the same context (thread)
    contextual: bool = False
    connection_context: ContextVar[ConnType | None]

    #: Callable to release a connection. Must return True if the connection
    #: may be reused, else False.
    before_release: Optional[Callable[[ConnType], bool]]

    #: Callable called every time a connection is acquired to validate
    #: that it is still alive
    validate: Optional[Callable[[ConnType], bool]]

    def __init__(
        self,
        connection_factory,
        limit=0,
        validator: Callable[[ConnType], bool] | Literal["auto"] | None = "auto",
        contextual=False,
    ):
        if isinstance(validator, poolvalidators.ConnectionValidator):
            self.validate = validator.validate
            self.before_release = validator.before_release
        elif validator == "auto":
            self.validate = self.auto_validate  # type: ignore
            self.before_release = None
        else:
            self.validate = None
            self.before_release = None
        self.connection_factory = connection_factory  # type: ignore # noqa
        self.limit = limit
        self.contextual = contextual
        self.max_validation_retries = self.limit + 3
        self.connection_context = ContextVar("connection_context", default=None)

    def _getconn(self) -> ConnType:
        raise NotImplementedError()

    def getconn(self, contextual=True) -> ConnType:
        """
        Get a connection from the pool. The connection must be released by
        calling ``release``.
        """
        raise NotImplementedError

    def release(self, conn: ConnType):
        """
        Release the connection back to the pool
        """
        raise NotImplementedError

    def connect(self, transaction=False, contextual=True):
        """
        Return a context manager that manages acquiring and releasing a
        connection.

        :param transaction:
            If true, ``commit()`` will be called on the connection object at
            the end of the block, or ``rollback()`` in case of exception.
        """
        return self.connection_context_manager(self, transaction, contextual)

    def transaction(self, contextual=True):
        """
        A convenience method equivalent to calling
        :meth:`~embrace.pool.ConnectionPoolBase(transaction=True)`
        """
        return self.connection_context_manager(self, True, contextual)

    def set_validator(self, v):
        self.validate = v.validate
        self.before_release = v.before_release

    def auto_validate(self, conn):
        validator = poolvalidators.ConnectionValidator()

        for cls in poolvalidators.validators:
            if isinstance(conn, cls):
                validator = poolvalidators.validators[cls]()
                break
        self.set_validator(validator)
        return validator.validate(conn)


class ConnectionPool(ConnectionPoolBase):
    """
    A connection pool implementation using a queue to provide thread-safety.
    """

    # Maintain the pool in a queue for thread/process safety
    queue_class = queue.Queue
    lock_class = threading.Lock

    #: How long to wait for a connection to become available
    timeout = 5

    _pool: queue.Queue

    def __init__(
        self,
        connection_factory,
        limit=0,
        validator: Callable[[ConnType], bool] | Literal["auto"] | None = "auto",
        timeout=timeout,
        contextual=False,
    ):
        super(ConnectionPool, self).__init__(
            connection_factory, limit, validator, contextual=contextual
        )

        self._pool = self.queue_class()
        self.lock = self.lock_class()
        self.connections_created = 0
        self.reached_limit = False
        self.timeout = timeout

    def getconn(self, contextual=True) -> ConnType:
        contextual = contextual and self.contextual
        if contextual and (conncount := self.connection_context.get()):
            conn, refcount = conncount
            self.connection_context.set((conn, refcount + 1))
            return conn

        try:
            conn = self._pool.get(block=self.reached_limit, timeout=self.timeout)
        except queue.Empty:
            if self.limit:
                with self.lock:
                    if self.reached_limit:
                        raise exceptions.ConnectionLimitError()
                    else:
                        conn = self._connect()
            else:
                conn = self._connect()

        if self.validate:
            retry = 0
            while True:
                try:
                    if self.validate(conn):
                        break
                except Exception:
                    logger.exception(f"Could not validate {conn=}")

                logger.info(f"Invalid connection {conn}")
                self.release(conn)

                if retry >= self.max_validation_retries:
                    raise Exception(
                        f"Could not validate a connection after "
                        f"{self.max_validation_retries} attempts"
                    )

                conn = self._getconn()
                retry += 1

        if contextual:
            conncount = self.connection_context.get()
            assert conncount is None
            self.connection_context.set((conn, 1))
        return conn

    def _getconn(self):
        """
        Return a connection from the pool.
        """

    def _connect(self):
        self.connections_created += 1
        self.reached_limit = bool(self.limit and self.connections_created >= self.limit)
        return self.connection_factory()  # type: ignore

    def release(self, conn):

        full_release = True
        if self.contextual:
            conncount = self.connection_context.get()
            if conncount:
                context_conn, refcount = conncount
                if conn is context_conn:
                    if refcount > 1:
                        self.connection_context.set((conn, refcount - 1))
                        full_release = False
                    else:
                        self.connection_context.set(None)
                        full_release = True

        if full_release:
            try:
                reuse = self.before_release(conn) if self.before_release else True
            except Exception:
                reuse = False
                logger.exception(f"Could not release {conn=}")

            if reuse:
                self._pool.put(conn)
            else:
                conn.close()
                if self.limit:
                    self.lock.acquire()
                    self.connections_created -= 1
                    self.reached_limit = self.connections_created >= self.limit
                    self.lock.release()


class SQLAlchemyPool(ConnectionPoolBase):
    def __init__(self, connection_factory, limit=0, validator=None):
        if limit:
            raise ValueError(
                "Cannot set limit here, "
                "configure this in SQLAlchemy connection pooling instead"
            )
        if validator:
            raise ValueError(
                "Cannot set validator here, "
                "configure this in SQLAlchemy connection pooling instead"
            )

        def _connection_factory(sessiongetter=connection_factory):
            return sessiongetter().connection().connection.connection

        super(SQLAlchemyPool, self).__init__(_connection_factory, limit, validator)

    def getconn(self, contextual=True) -> ConnType:
        return self.connection_factory()  # type: ignore

    def release(self, conn):
        pass
