#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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 aiohttp
import asyncio
import socket
import sys

from gremlin_python.driver.exceptions import ReadTimeoutError

if sys.version_info >= (3, 11):
    import asyncio as async_timeout
else:
    import async_timeout

__author__ = 'Lyndon Bauto (lyndonb@bitquilltech.com)'

# Default connection option values (canonical TinkerPop 4.x GLV defaults). The millisecond-suffixed
# options are the primary form (mirroring the other GLVs); aiohttp itself works in seconds, so the
# values are converted internally.
DEFAULT_CONNECT_TIMEOUT_MILLIS = 5000
DEFAULT_IDLE_TIMEOUT_MILLIS = 180000
DEFAULT_KEEP_ALIVE_TIME_MILLIS = 30000
# Seconds equivalents retained for internal use / backwards reference.
DEFAULT_CONNECT_TIMEOUT = DEFAULT_CONNECT_TIMEOUT_MILLIS / 1000
DEFAULT_IDLE_TIMEOUT = DEFAULT_IDLE_TIMEOUT_MILLIS / 1000
DEFAULT_KEEP_ALIVE_TIME = DEFAULT_KEEP_ALIVE_TIME_MILLIS / 1000
DEFAULT_COMPRESSION = 'deflate'

def _resolve_timeout_seconds(millis, seconds, default_millis):
    """Resolve a timeout to seconds (aiohttp's unit) from the ``*_millis`` number or the idiomatic
    unsuffixed seconds number (``None`` means not supplied for either). Supplying both raises
    ``ValueError``; if neither is given, ``default_millis`` is used (``None`` leaves it unset).
    """
    if millis is not None and seconds is not None:
        raise ValueError("provide only one of the milliseconds option or the seconds option, not both")
    if seconds is not None:
        return seconds
    if millis is not None:
        return millis / 1000
    if default_millis is None:
        return None
    return default_millis / 1000


def _normalize_compression(compression):
    """Normalize the compression option to a canonical string ('none' or 'deflate').

    Accepts the string forms 'none'/'deflate'.
    """
    if compression is None:
        return DEFAULT_COMPRESSION
    if isinstance(compression, str):
        normalized = compression.lower()
        if normalized in ('none', 'deflate'):
            return normalized
        raise ValueError("compression must be one of 'none', 'deflate', got '%s'" % compression)
    raise TypeError("compression must be a str ('none'|'deflate'), got %s" % type(compression).__name__)


def _run_read(loop, read_timeout, coro):
    """Run a response-read coroutine on ``loop``, normalizing aiohttp's read timeout
    (SocketTimeoutError / ServerTimeoutError) into a ``ReadTimeoutError`` so a read
    timeout always surfaces as one deterministic, transport-agnostic type. It subclasses
    the builtin ``TimeoutError`` so callers can still ``except TimeoutError``."""
    try:
        return loop.run_until_complete(coro)
    except aiohttp.ServerTimeoutError as e:
        raise ReadTimeoutError(
            f"Read timed out after {read_timeout}s waiting for response data.") from e


def _keep_alive_socket_options(keep_alive_time):
    """Build the list of socket options that enable TCP keep-alive with the
    given idle time before probes begin. TCP_KEEPIDLE is platform dependent
    (Linux); macOS exposes the equivalent as TCP_KEEPALIVE. Both are guarded so
    platforms lacking the option (e.g. Windows) simply enable SO_KEEPALIVE."""
    options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
    idle_opt = getattr(socket, 'TCP_KEEPIDLE', None)
    if idle_opt is None:
        # macOS names the idle-before-probe option TCP_KEEPALIVE
        idle_opt = getattr(socket, 'TCP_KEEPALIVE', None)
    if idle_opt is not None:
        options.append((socket.IPPROTO_TCP, idle_opt, int(keep_alive_time)))
    return options


def _keep_alive_socket_factory(keep_alive_time):
    """Return a socket_factory (aiohttp >= 3.11) that applies the keep-alive
    socket options when each connection socket is created."""
    options = _keep_alive_socket_options(keep_alive_time)

    def factory(addr_info):
        family, type_, proto, _, _ = addr_info
        sock = socket.socket(family=family, type=type_, proto=proto)
        for level, optname, value in options:
            try:
                sock.setsockopt(level, optname, value)
            except (OSError, AttributeError):
                pass
        return sock

    return factory


class AiohttpSyncStream:
    """Wraps aiohttp's async StreamReader as a synchronous file-like object.
    read(n) blocks until exactly n bytes are available from the HTTP response.

    Maintains an internal byte buffer that is refilled one HTTP chunk at a time
    so the deserializer's many small read(n) calls don't each pay the cost of a
    full asyncio event-loop turn."""

    # Max bytes pulled from the response per underlying read. Matches
    # aiohttp.StreamReader's default 64 KB limit, which is the per-connection
    # high-water mark, asking for more in one read() never returns more.
    _FILL_SIZE = 64 * 1024

    def __init__(self, response, loop, read_timeout):
        self._response = response
        self._loop = loop
        self._read_timeout = read_timeout
        self._buf = bytearray()
        self._pos = 0

    def read(self, n):
        if n <= 0:
            return b''
        while len(self._buf) - self._pos < n:
            data = self._read_chunk()
            if not data:
                partial = bytes(self._buf[self._pos:])
                self._buf.clear()
                self._pos = 0
                raise asyncio.IncompleteReadError(partial=partial, expected=n)
            self._buf.extend(data)
        end = self._pos + n
        out = bytes(self._buf[self._pos:end])
        self._pos = end
        # Reclaim memory once the buffer is fully drained
        if self._pos == len(self._buf):
            self._buf.clear()
            self._pos = 0
        return out

    def _read_chunk(self):
        return _run_read(self._loop, self._read_timeout,
                         self._response.content.read(self._FILL_SIZE))


class AiohttpHTTPTransport:
    nest_asyncio_applied = False

    def __init__(self, call_from_event_loop=None, write_timeout=None,
                 connect_timeout_millis=None, connect_timeout=None,
                 idle_timeout_millis=None, idle_timeout=None,
                 read_timeout_millis=None, read_timeout=None,
                 keep_alive_time_millis=None, keep_alive_time=None,
                 compression=DEFAULT_COMPRESSION, proxy=None, trust_env=False,
                 max_connections=None, **kwargs):
        if call_from_event_loop is not None and call_from_event_loop and not AiohttpHTTPTransport.nest_asyncio_applied:
            """ 
                The AiohttpTransport implementation uses the asyncio event loop. Because of this, it cannot be called 
                within an event loop without nest_asyncio. If the code is ever refactored so that it can be called 
                within an event loop this import and call can be removed. Without this, applications which use the 
                event loop to call gremlin-python (such as Jupyter) will not work.
            """
            import nest_asyncio
            nest_asyncio.apply()
            AiohttpHTTPTransport.nest_asyncio_applied = True

        # Start event loop and initialize client session and response to None
        self._loop = asyncio.new_event_loop()
        self._client_session = None
        self._http_req_resp = None
        self._enable_ssl = False
        self._ssl_context = None
        self._url = None

        # Set all inner variables to parameters passed in.
        self._aiohttp_kwargs = kwargs
        self._write_timeout = write_timeout

        # Timeouts accept a millisecond number (*_millis, primary form) or the idiomatic seconds number
        # (unsuffixed). read_timeout defaults off; the others use the canonical millisecond defaults.
        self._read_timeout = _resolve_timeout_seconds(read_timeout_millis, read_timeout, None)

        # Connection-level pooling / lifecycle options.
        self._connect_timeout = _resolve_timeout_seconds(connect_timeout_millis, connect_timeout, DEFAULT_CONNECT_TIMEOUT_MILLIS)
        self._idle_timeout = _resolve_timeout_seconds(idle_timeout_millis, idle_timeout, DEFAULT_IDLE_TIMEOUT_MILLIS)
        self._keep_alive_time = _resolve_timeout_seconds(keep_alive_time_millis, keep_alive_time, DEFAULT_KEEP_ALIVE_TIME_MILLIS)
        # Caps the aiohttp connector's simultaneous connections per Connection
        # (the Client also sizes its Connection pool by this value).
        self._max_connections = max_connections

        # Compression negotiation. Default 'deflate' (on); advertises
        # Accept-Encoding: deflate. Set 'none' to opt out.
        self._compression = _normalize_compression(compression)

        # HTTP proxy support routed through the ClientSession.
        self._proxy = proxy
        self._trust_env = trust_env

        # ssl: canonical name accepting an ssl.SSLContext.
        if "ssl" in self._aiohttp_kwargs:
            self._ssl_context = self._aiohttp_kwargs.pop("ssl")
            self._enable_ssl = True

    def __del__(self):
        # Close will only actually close if things are left open, so this is safe to call.
        # Clean up any connection resources and close the event loop.
        self.close()

    def connect(self, url, headers=None):
        self._url = url
        # Inner function to perform async connect.
        async def async_connect():
            # Build the TCP connector with the standardized pooling / lifecycle
            # options. keepalive_timeout maps to the idle connection timeout;
            # the socket factory enables TCP keep-alive probes after
            # keep_alive_time idle.
            connector_kwargs = {}
            if self._enable_ssl:
                # ssl context is established through the tcp connector
                connector_kwargs['ssl_context'] = self._ssl_context
            if self._idle_timeout is not None:
                connector_kwargs['keepalive_timeout'] = self._idle_timeout
            if self._keep_alive_time is not None:
                self._apply_keep_alive(connector_kwargs)
            # Reflect max_connections at the aiohttp layer so the connector's
            # simultaneous-connection limit matches the driver option.
            if self._max_connections is not None:
                connector_kwargs['limit'] = self._max_connections

            session_kwargs = {'headers': headers, 'loop': self._loop,
                              'trust_env': self._trust_env}
            # Use the per-socket timeouts (sock_connect/sock_read) rather than a
            # whole-request total, which would abort long but legitimate streaming
            # responses. sock_read bounds idle time between chunks so a stalled
            # server cannot hang forever.
            timeout_kwargs = {}
            if self._connect_timeout is not None:
                timeout_kwargs['sock_connect'] = self._connect_timeout
            if self._read_timeout is not None:
                timeout_kwargs['sock_read'] = self._read_timeout
            if timeout_kwargs:
                session_kwargs['timeout'] = aiohttp.ClientTimeout(**timeout_kwargs)
            if connector_kwargs:
                session_kwargs['connector'] = aiohttp.TCPConnector(**connector_kwargs)

            self._client_session = aiohttp.ClientSession(**session_kwargs)

        # Execute the async connect synchronously.
        self._loop.run_until_complete(async_connect())

    def _apply_keep_alive(self, connector_kwargs):
        """Wire TCP keep-alive into the connector via the aiohttp socket_factory.
        The factory sets SO_KEEPALIVE plus the per-socket idle time; unsupported
        platforms degrade gracefully inside the factory."""
        connector_kwargs['socket_factory'] = _keep_alive_socket_factory(self._keep_alive_time)

    def write(self, message):
        # Negotiate compression unless the caller already set Accept-Encoding:
        # deflate advertises Accept-Encoding: deflate; none suppresses aiohttp's
        # auto-injected Accept-Encoding so compression is not silently negotiated.
        headers = message['headers']
        has_accept_encoding = any(k.lower() == 'accept-encoding' for k in headers)
        skip_auto_headers = None
        if self._compression == 'deflate':
            if not has_accept_encoding:
                headers['accept-encoding'] = 'deflate'
        elif not has_accept_encoding:
            skip_auto_headers = ['Accept-Encoding']

        # Inner function to perform async write.
        async def async_write():
            post_kwargs = dict(self._aiohttp_kwargs)
            if self._proxy is not None:
                post_kwargs['proxy'] = self._proxy
            if skip_auto_headers is not None:
                post_kwargs['skip_auto_headers'] = skip_auto_headers
            async with async_timeout.timeout(self._write_timeout):
                self._http_req_resp = await self._client_session.post(url=self._url,
                                                                      data=message['payload'],
                                                                      headers=headers,
                                                                      **post_kwargs)

        # Execute the async write synchronously.
        self._loop.run_until_complete(async_write())

    def get_stream(self):
        """Returns a synchronous file-like object for the HTTP response body."""
        return AiohttpSyncStream(self._http_req_resp, self._loop, self._read_timeout)

    @property
    def content_type(self):
        """Returns the Content-Type header of the HTTP response."""
        if self._http_req_resp is not None:
            return self._http_req_resp.headers.get('content-type', '')
        return ''

    @property
    def status_code(self):
        """Returns the HTTP status code of the response."""
        if self._http_req_resp is not None:
            return self._http_req_resp.status
        return None

    def read_body(self):
        """Read the entire HTTP response body as bytes."""
        return _run_read(self._loop, self._read_timeout, self._http_req_resp.read())

    def evict_response(self):
        """Close the current HTTP response and its underlying connection, evicting it from
        aiohttp's connection pool. Used after a transport failure so a dead/half-closed
        connection is discarded rather than returned to the pool for reuse."""
        if self._http_req_resp is not None:
            self._http_req_resp.close()
            self._http_req_resp = None

    def close(self):
        # Inner function to perform async close.
        async def async_close():
            if self._client_session is not None and not self._client_session.closed:
                await self._client_session.close()
                self._client_session = None

        # If the loop is not closed (connection hasn't already been closed)
        if not self._loop.is_closed():
            # Execute the async close synchronously.
            self._loop.run_until_complete(async_close())

            # Close the event loop.
            self._loop.close()

    @property
    def closed(self):
        # Connection is closed when client session is closed (or not yet created).
        return self._client_session is None or self._client_session.closed
