"""
Watching and streaming watch-events.

Kubernetes client's watching streams are synchronous. To make them asynchronous,
we put them into a `concurrent.futures.ThreadPoolExecutor`,
and yield from there asynchronously.

However, async/await coroutines misbehave with `StopIteration` exceptions
raised by the `next` method: see `PEP-479`_.

As a workaround, we replace `StopIteration` with our custom `StopStreaming`
inherited from `RuntimeError` (as suggested by `PEP-479`_),
and re-implement the generators to make them async.

All of this is a workaround for the standard Kubernetes client's limitations.
They would not be needed if the client library were natively asynchronous.

.. _PEP-479: https://www.python.org/dev/peps/pep-0479/
"""

import asyncio
import contextlib
import json
import logging
from typing import AsyncIterator, Dict, Optional, cast

import aiohttp

from kopf.clients import auth, discovery, errors, fetching
from kopf.structs import bodies, configuration, primitives, resources
from kopf.utilities import aiotasks

logger = logging.getLogger(__name__)


class WatchingError(Exception):
    """
    Raised when an unexpected error happens in the watch-stream API.
    """


async def infinite_watch(
        *,
        settings: configuration.OperatorSettings,
        resource: resources.Resource,
        namespace: Optional[str],
        freeze_checker: Optional[primitives.ToggleSet] = None,
        _iterations: Optional[int] = None,  # used in tests/mocks/fixtures
) -> AsyncIterator[bodies.RawEvent]:
    """
    Stream the watch-events infinitely.

    This routine is extracted because it is difficult to test infinite loops.
    It is made as simple as possible, and is assumed to work without testing.

    This routine never ends gracefully. If a watcher's stream fails,
    a new one is recreated, and the stream continues.
    It only exits with unrecoverable exceptions.
    """
    how = ' (frozen)' if freeze_checker is not None and freeze_checker.is_on() else ''
    where = f'in {namespace!r}' if namespace is not None else 'cluster-wide'
    logger.debug(f"Starting the watch-stream for {resource} {where}{how}.")
    try:
        while _iterations is None or _iterations > 0:  # equivalent to `while True` in non-test mode
            _iterations = None if _iterations is None else _iterations - 1
            async with streaming_block(
                namespace=namespace,
                resource=resource,
                freeze_checker=freeze_checker,
            ) as freeze_waiter:
                stream = continuous_watch(
                    settings=settings,
                    resource=resource,
                    namespace=namespace,
                    freeze_waiter=freeze_waiter,
                )
                async for raw_event in stream:
                    yield raw_event
            await asyncio.sleep(settings.watching.reconnect_backoff)
    finally:
        logger.debug(f"Stopping the watch-stream for {resource} {where}.")


@contextlib.asynccontextmanager
async def streaming_block(
        *,
        resource: resources.Resource,
        namespace: Optional[str],
        freeze_checker: Optional[primitives.ToggleSet],
) -> AsyncIterator[aiotasks.Future]:
    """
    Block the execution until the freeze is off; signal when it is on again.

    This prevents both watching and listing while the freeze mode is on,
    until it is off. Specifically, the watch-stream closes its connection
    once the freeze mode is on, so the while-true & for-event-in-stream cycles
    exit, and the streaming coroutine is started again by `infinite_stream()`
    (the watcher timeout is swallowed by the freeze time).

    Returns a future (or a task) that is set when the freeze is turned on again.

    A stop-future is a client-specific way of terminating the streaming HTTPS
    connections when the freeze is turned back on. The low-level streaming API
    call attaches its `response.close()` to the future's "done" callback,
    so that the stream is closed once the freeze is turned back on.

    Note: this routine belongs to watching and does not belong to peering.
    The freeze can be managed in any other ways: as an imaginary edge case,
    imagine a operator with UI with a "pause" button that freezes the operator.
    """
    where = f'in {namespace!r}' if namespace is not None else 'cluster-wide'

    # Block until unfrozen before even starting the API communication.
    if freeze_checker is not None and freeze_checker.is_on():
        names = {toggle.name for toggle in freeze_checker if toggle.is_on() and toggle.name}
        freezing_reason = f" (blockers: {', '.join(names)})" if names else ""
        logger.debug(f"Freezing the watch-stream for {resource} {where}{freezing_reason}.")

        await freeze_checker.wait_for(False)

        names = {toggle.name for toggle in freeze_checker if toggle.is_on() and toggle.name}
        resuming_reason = f" (resolved: {', '.join(names)})" if names else ""
        logger.debug(f"Resuming the watch-stream for {resource} {where}{resuming_reason}.")

    # Create the signalling future that the freeze is on again.
    freeze_waiter: aiotasks.Future
    if freeze_checker is not None:
        freeze_waiter = aiotasks.create_task(
            freeze_checker.wait_for(True),
            name=f"freeze-waiter for {resource}")
    else:
        freeze_waiter = asyncio.Future()  # a dummy just to have it

    # Go for the streaming with the prepared freezing/unfreezing setup.
    try:
        yield freeze_waiter
    finally:
        with contextlib.suppress(asyncio.CancelledError):
            freeze_waiter.cancel()
            await freeze_waiter


async def continuous_watch(
        *,
        settings: configuration.OperatorSettings,
        resource: resources.Resource,
        namespace: Optional[str],
        freeze_waiter: aiotasks.Future,
) -> AsyncIterator[bodies.RawEvent]:

    # First, list the resources regularly, and get the list's resource version.
    # Simulate the events with type "None" event - used in detection of causes.
    items, resource_version = await fetching.list_objs_rv(resource=resource, namespace=namespace)
    for item in items:
        yield {'type': None, 'object': item}

    # Repeat through disconnects of the watch as long as the resource version is valid (no errors).
    # The individual watching API calls are disconnected by timeout even if the stream is fine.
    while not freeze_waiter.done():

        # Then, watch the resources starting from the list's resource version.
        stream = watch_objs(
            settings=settings,
            resource=resource, namespace=namespace,
            timeout=settings.watching.server_timeout,
            since=resource_version,
            freeze_waiter=freeze_waiter,
        )
        async for raw_input in stream:
            raw_type = raw_input['type']
            raw_object = raw_input['object']

            # "410 Gone" is for the "resource version too old" error, we must restart watching.
            # The resource versions are lost by k8s after few minutes (5, as per the official doc).
            # The error occurs when there is nothing happening for few minutes. This is normal.
            if raw_type == 'ERROR' and cast(bodies.RawError, raw_object)['code'] == 410:
                where = f'in {namespace!r}' if namespace is not None else 'cluster-wide'
                logger.debug(f"Restarting the watch-stream for {resource} {where}.")
                return  # out of the regular stream, to the infinite stream.

            # Other watch errors should be fatal for the operator.
            if raw_type == 'ERROR':
                raise WatchingError(f"Error in the watch-stream: {raw_object}")

            # Ensure that the event is something we understand and can handle.
            if raw_type not in ['ADDED', 'MODIFIED', 'DELETED']:
                logger.warning("Ignoring an unsupported event type: %r", raw_input)
                continue

            # Keep the latest seen resource version for continuation of the stream on disconnects.
            body = cast(bodies.RawBody, raw_object)
            resource_version = body.get('metadata', {}).get('resourceVersion', resource_version)

            # Yield normal events to the consumer. Errors are already filtered out.
            yield cast(bodies.RawEvent, raw_input)


@auth.reauthenticated_stream
async def watch_objs(
        *,
        settings: configuration.OperatorSettings,
        resource: resources.Resource,
        namespace: Optional[str],
        timeout: Optional[float] = None,
        since: Optional[str] = None,
        context: Optional[auth.APIContext] = None,  # injected by the decorator
        freeze_waiter: aiotasks.Future,
) -> AsyncIterator[bodies.RawInput]:
    """
    Watch objects of a specific resource type.

    The cluster-scoped call is used in two cases:

    * The resource itself is cluster-scoped, and namespacing makes not sense.
    * The operator serves all namespaces for the namespaced custom resource.

    Otherwise, the namespace-scoped call is used:

    * The resource is namespace-scoped AND operator is namespaced-restricted.
    """
    if context is None:
        raise RuntimeError("API instance is not injected by the decorator.")

    is_namespaced = await discovery.is_namespaced(resource=resource, context=context)
    namespace = namespace if is_namespaced else None

    params: Dict[str, str] = {}
    params['watch'] = 'true'
    if since is not None:
        params['resourceVersion'] = since
    if timeout is not None:
        params['timeoutSeconds'] = str(timeout)

    # Stream the parsed events from the response until it is closed server-side,
    # or until it is closed client-side by the freeze-waiting future's callbacks.
    try:
        response = await context.session.get(
            url=resource.get_url(server=context.server, namespace=namespace, params=params),
            timeout=aiohttp.ClientTimeout(
                total=settings.watching.client_timeout,
                sock_connect=settings.watching.connect_timeout,
            ),
        )
        await errors.check_response(response)

        response_close_callback = lambda _: response.close()
        freeze_waiter.add_done_callback(response_close_callback)
        try:
            async with response:
                async for line in _iter_jsonlines(response.content):
                    raw_input = cast(bodies.RawInput, json.loads(line.decode("utf-8")))
                    yield raw_input
        finally:
            freeze_waiter.remove_done_callback(response_close_callback)

    except (aiohttp.ClientConnectionError, aiohttp.ClientPayloadError, asyncio.TimeoutError):
        pass


async def _iter_jsonlines(
        content: aiohttp.StreamReader,
        chunk_size: int = 1024 * 1024,
) -> AsyncIterator[bytes]:
    """
    Iterate line by line over the response's content.

    Usage::

        async for line in _iter_lines(response.content):
            pass

    This is an equivalent of::

        async for line in response.content:
            pass

    Except that the aiohttp's line iteration fails if the accumulated buffer
    length is above 2**17 bytes, i.e. 128 KB (`aiohttp.streams.DEFAULT_LIMIT`
    for the buffer's low-watermark, multiplied by 2 for the high-watermark).
    Kubernetes secrets and other fields can be much longer, up to MBs in length.

    The chunk size of 1MB is an empirical guess for keeping the memory footprint
    reasonably low on huge amount of small lines (limited to 1 MB in total),
    while ensuring the near-instant reads of the huge lines (can be a problem
    with a small chunk size due to too many iterations).

    .. seealso::
        https://github.com/zalando-incubator/kopf/issues/275
    """

    # Minimize the memory footprint by keeping at most 2 copies of a yielded line in memory
    # (in the buffer and as a yielded value), and at most 1 copy of other lines (in the buffer).
    buffer = b''
    async for data in content.iter_chunked(chunk_size):
        buffer += data
        del data

        start = 0
        index = buffer.find(b'\n', start)
        while index >= 0:
            line = buffer[start:index]
            if line:
                yield line
            del line
            start = index + 1
            index = buffer.find(b'\n', start)

        if start > 0:
            buffer = buffer[start:]

    if buffer:
        yield buffer
