#
# This source file is part of the EdgeDB open source project.
#
# Copyright 2016-present MagicStack Inc. and the EdgeDB authors.
#
# 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.
#


import random
import socket
import ssl
import time
import typing
import warnings

from . import abstract
from . import base_con
from . import con_utils
from . import enums
from . import errors
from . import options
from . import transaction as _transaction
from . import retry as _retry

from .datatypes import datatypes
from .protocol import blocking_proto, protocol
from .protocol.protocol import CodecsRegistry as _CodecsRegistry
from .protocol.protocol import QueryCodecsCache as _QueryCodecsCache


class _BlockingIOConnectionImpl:

    def __init__(self, codecs_registry, query_cache):
        self._addr = None
        self._protocol = None
        self._codecs_registry = codecs_registry
        self._query_cache = query_cache

    def connect(self, addrs, config, params, *,
                single_attempt=False, connection):
        addr = None
        start = time.monotonic()
        if single_attempt:
            max_time = 0
        else:
            max_time = start + config.wait_until_available
        iteration = 1

        while True:
            for addr in addrs:
                try:
                    self._connect_addr(addr, config, params, connection)
                except TimeoutError as e:
                    if iteration == 1 or time.monotonic() < max_time:
                        continue
                    else:
                        raise errors.ClientConnectionTimeoutError(
                            f"connecting to {addr} failed in"
                            f" {config.connect_timeout} sec"
                        ) from e
                except errors.ClientConnectionError as e:
                    if (
                        e.has_tag(errors.SHOULD_RECONNECT) and
                        (iteration == 1 or time.monotonic() < max_time)
                    ):
                        continue
                    nice_err = e.__class__(
                        con_utils.render_client_no_connection_error(
                            e,
                            addr,
                            attempts=iteration,
                            duration=time.monotonic() - start,
                        ))
                    raise nice_err from e.__cause__
                else:
                    assert self._protocol
                    return

            iteration += 1
            time.sleep(0.01 + random.random() * 0.2)

    def _connect_addr(self, addr, config, params, connection):
        timeout = config.connect_timeout
        deadline = time.monotonic() + timeout
        tls_compat = False

        if isinstance(addr, str):
            # UNIX socket
            sock = socket.socket(socket.AF_UNIX)
        else:
            sock = socket.socket(socket.AF_INET)

        try:
            sock.settimeout(timeout)

            try:
                sock.connect(addr)

                if not isinstance(addr, str):
                    time_left = deadline - time.monotonic()
                    if time_left <= 0:
                        raise TimeoutError

                    # Upgrade to TLS
                    if params.ssl_ctx.check_hostname:
                        server_hostname = addr[0]
                    else:
                        server_hostname = None
                    sock.settimeout(time_left)
                    try:
                        sock = params.ssl_ctx.wrap_socket(
                            sock, server_hostname=server_hostname
                        )
                    except ssl.CertificateError as e:
                        raise con_utils.wrap_error(e) from e
                    except ssl.SSLError as e:
                        if e.reason == 'CERTIFICATE_VERIFY_FAILED':
                            raise con_utils.wrap_error(e) from e

                        # Retry in plain text
                        time_left = deadline - time.monotonic()
                        if time_left <= 0:
                            raise TimeoutError
                        sock.close()
                        sock = socket.socket(socket.AF_INET)
                        sock.settimeout(time_left)
                        sock.connect(addr)
                        tls_compat = True
                    else:
                        con_utils.check_alpn_protocol(sock)
            except socket.gaierror as e:
                # All name resolution errors are considered temporary
                err = errors.ClientConnectionFailedTemporarilyError(str(e))
                raise err from e
            except OSError as e:
                raise con_utils.wrap_error(e) from e

            time_left = deadline - time.monotonic()
            if time_left <= 0:
                raise TimeoutError

            if not isinstance(addr, str):
                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

            proto = blocking_proto.BlockingIOProtocol(
                params, sock, tls_compat
            )
            proto.set_connection(connection)

            try:
                sock.settimeout(time_left)
                proto.sync_connect()
                sock.settimeout(None)
            except OSError as e:
                raise con_utils.wrap_error(e) from e

            self._protocol = proto
            self._addr = addr

        except Exception:
            sock.close()
            raise

    def privileged_execute(self, query):
        self._protocol.sync_simple_query(query, enums.Capability.ALL)

    def is_closed(self):
        proto = self._protocol
        return not (proto and proto.sock is not None and
                    proto.sock.fileno() >= 0 and proto.connected)

    def close(self):
        if self._protocol:
            self._protocol.abort()


class BlockingIOConnection(
    base_con.BaseConnection,
    abstract.Executor,
    options._OptionsMixin,
):

    def __init__(self, addrs, config, params, *,
                 codecs_registry, query_cache):
        self._inner = base_con._InnerConnection(
            addrs, config, params,
            codecs_registry=codecs_registry,
            query_cache=query_cache)
        super().__init__()

    def _shallow_clone(self):
        if self._inner._borrowed_for:
            raise base_con.borrow_error(self._inner._borrowed_for)
        new_conn = self.__class__.__new__(self.__class__)
        new_conn._inner = self._inner
        return new_conn

    def ensure_connected(self, single_attempt=False):
        inner = self._inner
        if inner._borrowed_for:
            raise base_con.borrow_error(inner._borrowed_for)
        if not inner._impl or inner._impl.is_closed():
            self._reconnect(single_attempt=single_attempt)

    def _reconnect(self, single_attempt=False):
        inner = self._inner
        inner._impl = _BlockingIOConnectionImpl(
            inner._codecs_registry, inner._query_cache)
        inner._impl.connect(inner._addrs, inner._config, inner._params,
                            single_attempt=single_attempt, connection=inner)
        assert inner._impl._protocol

    def _get_protocol(self):
        inner = self._inner
        if inner._borrowed_for:
            raise base_con.borrow_error(inner._borrowed_for)
        if not inner._impl or inner._impl.is_closed():
            self._reconnect()
        return inner._impl._protocol

    def _dump(
        self,
        *,
        on_header: typing.Callable[[bytes], None],
        on_data: typing.Callable[[bytes], None],
    ) -> None:
        self._get_protocol().sync_dump(
            header_callback=on_header,
            block_callback=on_data)

    def _restore(
        self,
        *,
        header: bytes,
        data_gen: typing.Iterable[bytes],
    ) -> None:
        self._get_protocol().sync_restore(
            header=header,
            data_gen=data_gen
        )

    def _dispatch_log_message(self, msg):
        for cb in self._inner._log_listeners:
            cb(self, msg)

    def _fetchall(
        self,
        query: str,
        *args,
        __limit__: int=0,
        __typenames__: bool=False,
        **kwargs,
    ) -> datatypes.Set:
        inner = self._inner
        return self._get_protocol().sync_execute_anonymous(
            query=query,
            args=args,
            kwargs=kwargs,
            reg=inner._codecs_registry,
            qc=inner._query_cache,
            implicit_limit=__limit__,
            inline_typenames=__typenames__,
            io_format=protocol.IoFormat.BINARY,
            allow_capabilities=enums.Capability.EXECUTE,
        )

    def _fetchall_json(
        self,
        query: str,
        *args,
        __limit__: int=0,
        **kwargs,
    ) -> datatypes.Set:
        inner = self._inner
        return self._get_protocol().sync_execute_anonymous(
            query=query,
            args=args,
            kwargs=kwargs,
            reg=inner._codecs_registry,
            qc=inner._query_cache,
            implicit_limit=__limit__,
            inline_typenames=False,
            io_format=protocol.IoFormat.JSON,
            allow_capabilities=enums.Capability.EXECUTE,
        )

    def _execute(
        self,
        *,
        query: str,
        args,
        kwargs,
        io_format,
        expect_one=False,
        required_one=False,
    ):
        inner = self._inner
        reconnect = False
        capabilities = None
        i = 0
        while True:
            try:
                if reconnect:
                    self._reconnect(single_attempt=True)
                return self._get_protocol().sync_execute_anonymous(
                    query=query,
                    args=args,
                    kwargs=kwargs,
                    reg=inner._codecs_registry,
                    qc=inner._query_cache,
                    expect_one=expect_one,
                    required_one=required_one,
                    io_format=io_format,
                    allow_capabilities=enums.Capability.EXECUTE,
                )
            except errors.EdgeDBError as e:
                if not e.has_tag(errors.SHOULD_RETRY):
                    raise e
                if capabilities is None:
                    _, _, _, capabilities = inner._query_cache.get(
                        query=query,
                        io_format=io_format,
                        implicit_limit=0,
                        inline_typenames=False,
                        inline_typeids=False,
                        expect_one=expect_one,
                    )
                # A query is read-only if it has no capabilities i.e.
                # capabilities == 0. Read-only queries are safe to retry.
                if capabilities != 0:
                    raise e
                rule = self._options.retry_options.get_rule_for_exception(e)
                if i >= rule.attempts:
                    raise e
                time.sleep(rule.backoff(i))
                reconnect = self.is_closed()

    def query(self, query: str, *args, **kwargs) -> datatypes.Set:
        return self._execute(
            query=query,
            args=args,
            kwargs=kwargs,
            io_format=protocol.IoFormat.BINARY,
        )

    def query_single(
        self, query: str, *args, **kwargs
    ) -> typing.Union[typing.Any, None]:
        return self._execute(
            query=query,
            args=args,
            kwargs=kwargs,
            expect_one=True,
            io_format=protocol.IoFormat.BINARY,
        )

    def query_required_single(self, query: str, *args, **kwargs) -> typing.Any:
        return self._execute(
            query=query,
            args=args,
            kwargs=kwargs,
            expect_one=True,
            required_one=True,
            io_format=protocol.IoFormat.BINARY,
        )

    def query_json(self, query: str, *args, **kwargs) -> str:
        return self._execute(
            query=query,
            args=args,
            kwargs=kwargs,
            io_format=protocol.IoFormat.JSON,
        )

    def _fetchall_json_elements(
            self, query: str, *args, **kwargs) -> typing.List[str]:
        inner = self._inner
        return self._get_protocol().sync_execute_anonymous(
            query=query,
            args=args,
            kwargs=kwargs,
            reg=inner._codecs_registry,
            qc=inner._query_cache,
            io_format=protocol.IoFormat.JSON_ELEMENTS,
            allow_capabilities=enums.Capability.EXECUTE,
        )

    def query_single_json(self, query: str, *args, **kwargs) -> str:
        inner = self._inner
        return self._get_protocol().sync_execute_anonymous(
            query=query,
            args=args,
            kwargs=kwargs,
            reg=inner._codecs_registry,
            qc=inner._query_cache,
            expect_one=True,
            io_format=protocol.IoFormat.JSON,
            allow_capabilities=enums.Capability.EXECUTE,
        )

    def query_required_single_json(self, query: str, *args, **kwargs) -> str:
        inner = self._inner
        return self._get_protocol().sync_execute_anonymous(
            query=query,
            args=args,
            kwargs=kwargs,
            reg=inner._codecs_registry,
            qc=inner._query_cache,
            expect_one=True,
            required_one=True,
            io_format=protocol.IoFormat.JSON,
            allow_capabilities=enums.Capability.EXECUTE,
        )

    def fetchall(self, query: str, *args, **kwargs) -> datatypes.Set:
        warnings.warn(
            'The "fetchall()" method is deprecated and is scheduled to be '
            'removed. Use the "query()" method instead.',
            DeprecationWarning, 2)
        return self.query(query, *args, **kwargs)

    def query_one(self, query: str, *args, **kwargs) -> typing.Any:
        warnings.warn(
            'The "query_one()" method is deprecated and is scheduled to be '
            'removed. Use the "query_single()" method instead.',
            DeprecationWarning, 2)
        return self.query_single(query, *args, **kwargs)

    def fetchone(self, query: str, *args, **kwargs) -> typing.Any:
        warnings.warn(
            'The "fetchone()" method is deprecated and is scheduled to be '
            'removed. Use the "query_single()" method instead.',
            DeprecationWarning, 2)
        return self.query_single(query, *args, **kwargs)

    def fetchall_json(self, query: str, *args, **kwargs) -> str:
        warnings.warn(
            'The "fetchall_json()" method is deprecated and is scheduled to '
            'be removed. Use the "query_json()" method instead.',
            DeprecationWarning, 2)
        return self.query_json(query, *args, **kwargs)

    def query_one_json(self, query: str, *args, **kwargs) -> str:
        warnings.warn(
            'The "query_one_json()" method is deprecated and is scheduled to '
            'be removed. Use the "query_single_json()" method instead.',
            DeprecationWarning, 2)
        return self.query_single_json(query, *args, **kwargs)

    def fetchone_json(self, query: str, *args, **kwargs) -> str:
        warnings.warn(
            'The "fetchone_json()" method is deprecated and is scheduled to '
            'be removed. Use the "query_single_json()" method instead.',
            DeprecationWarning, 2)
        return self.query_single_json(query, *args, **kwargs)

    def execute(self, query: str) -> None:
        self._get_protocol().sync_simple_query(query, enums.Capability.EXECUTE)

    def transaction(self) -> _retry.Retry:
        return _retry.Retry(self)

    def raw_transaction(self) -> _transaction.Transaction:
        warnings.warn(
            'The "raw_transaction()" method is deprecated and is scheduled '
            'to be removed. Use the "transaction()" method with '
            'retry attempts=1 instead',
            DeprecationWarning, 2)
        return _transaction.Transaction(
            self,
            self._options.transaction_options,
        )

    def retrying_transaction(self) -> _retry.Retry:
        warnings.warn(
            'The "retrying_transaction()" method has been renamed to '
            '"transaction()"',
            DeprecationWarning, 2)
        return _retry.Retry(self)

    def close(self) -> None:
        if not self.is_closed():
            self._inner._impl.close()

    def is_closed(self) -> bool:
        return self._inner._impl is None or self._inner._impl.is_closed()


def connect(dsn: str = None, *,
            credentials_file: str = None,
            host: str = None, port: int = None,
            user: str = None, password: str = None,
            database: str = None,
            tls_ca_file: str = None,
            tls_security: str = None,
            timeout: int = 10,
            wait_until_available: int = 30) -> BlockingIOConnection:

    connect_config, client_config = con_utils.parse_connect_arguments(
        dsn=dsn, credentials_file=credentials_file, host=host, port=port,
        user=user, password=password, database=database,
        timeout=timeout,
        wait_until_available=wait_until_available,
        tls_ca_file=tls_ca_file, tls_security=tls_security,

        # ToDos
        command_timeout=None,
        server_settings=None)

    conn = BlockingIOConnection(
        addrs=[connect_config.address], params=connect_config,
        config=client_config,
        codecs_registry=_CodecsRegistry(),
        query_cache=_QueryCodecsCache())
    conn.ensure_connected()
    return conn
