"""
WebSocket-based streaming kernel interaction APIs.

NOTE: For nginx-based setups, we need to gather all websocket-based API handlers
      under this "/stream/"-prefixed app.
"""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import secrets
import uuid
import weakref
from collections import defaultdict
from collections.abc import AsyncIterator, Callable, Iterable, Mapping, MutableMapping
from datetime import timedelta
from typing import (
    TYPE_CHECKING,
    Any,
)
from urllib.parse import urlparse

import aiohttp
import aiohttp_cors
import aiotools
import attrs
import trafaret as t
import zmq
import zmq.asyncio
from aiohttp import web
from aiotools import adefer, apartial
from etcd_client import GRPCStatusCode, GRPCStatusError

from ai.backend.common import validators as tx
from ai.backend.common.events.event_types.kernel.broadcast import KernelTerminatingBroadcastEvent
from ai.backend.common.exception import BackendAIError
from ai.backend.common.json import dump_json, load_json
from ai.backend.common.types import AccessKey, AgentId, KernelId, SessionId
from ai.backend.logging import BraceStyleAdapter
from ai.backend.manager.defs import DEFAULT_ROLE
from ai.backend.manager.errors.api import InvalidAPIParameters
from ai.backend.manager.errors.kernel import (
    InvalidStreamMode,
    SessionNotFound,
    TooManySessionsMatched,
)
from ai.backend.manager.errors.resource import AppNotFound, NoCurrentTaskContext
from ai.backend.manager.errors.service import AppServiceStartFailed
from ai.backend.manager.idle import AppStreamingStatus
from ai.backend.manager.models.kernel import KernelRow
from ai.backend.manager.models.session import KernelLoadingStrategy, SessionRow

from .auth import auth_required
from .manager import READ_ALLOWED, server_status_required
from .types import CORSOptions, WebMiddleware
from .utils import call_non_bursty, check_api_params
from .wsproxy import TCPProxy

if TYPE_CHECKING:
    from .context import RootContext

log = BraceStyleAdapter(logging.getLogger(__spec__.name))


@server_status_required(READ_ALLOWED)
@auth_required
@adefer
async def stream_pty(
    defer: Callable[[Callable[[], None]], None], request: web.Request
) -> web.StreamResponse:
    root_ctx: RootContext = request.app["_root.context"]
    app_ctx: PrivateContext = request.app["stream.context"]
    database_ptask_group: aiotools.PersistentTaskGroup = request.app["database_ptask_group"]
    session_name = request.match_info["session_name"]
    access_key = request["keypair"]["access_key"]
    api_version = request["api_version"]
    try:
        async with root_ctx.db.begin_readonly_session() as db_sess:
            session = await asyncio.shield(
                database_ptask_group.create_task(
                    SessionRow.get_session(
                        db_sess,
                        session_name,
                        access_key,
                        kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
                    )
                ),
            )
    except SessionNotFound:
        raise
    log.info("STREAM_PTY(ak:{0}, s:{1})", access_key, session_name)
    compute_session: KernelRow = session.main_kernel
    stream_key = compute_session.id

    await asyncio.shield(
        database_ptask_group.create_task(
            root_ctx.registry.increment_session_usage(session),
        )
    )
    ws = web.WebSocketResponse(max_msg_size=root_ctx.config_provider.config.manager.max_wsmsg_size)
    await ws.prepare(request)

    myself = asyncio.current_task()
    if myself is None:
        raise NoCurrentTaskContext("No current task context")
    app_ctx.stream_pty_handlers[stream_key].add(myself)
    defer(lambda: app_ctx.stream_pty_handlers[stream_key].discard(myself))

    async def connect_streams(
        compute_session: KernelRow,
    ) -> tuple[zmq.asyncio.Socket, zmq.asyncio.Socket]:
        # TODO: refactor as custom row/table method
        if compute_session.kernel_host is None:
            hostname = urlparse(compute_session.agent_addr).hostname
            kernel_host = hostname.decode() if isinstance(hostname, bytes) else hostname
        else:
            kernel_host = compute_session.kernel_host
        stdin_addr = f"tcp://{kernel_host}:{compute_session.repl_in_port}"
        log.debug("stream_pty({0}): stdin: {1}", stream_key, stdin_addr)
        stdin_sock = app_ctx.zctx.socket(zmq.PUB)
        stdin_sock.connect(stdin_addr)
        stdin_sock.setsockopt(zmq.LINGER, 100)
        stdout_addr = f"tcp://{kernel_host}:{compute_session.repl_out_port}"
        log.debug("stream_pty({0}): stdout: {1}", stream_key, stdout_addr)
        stdout_sock = app_ctx.zctx.socket(zmq.SUB)
        stdout_sock.connect(stdout_addr)
        stdout_sock.setsockopt(zmq.LINGER, 100)
        stdout_sock.subscribe(b"")
        return stdin_sock, stdout_sock

    # Wrap sockets in a list so that below coroutines can share reference changes.
    socks = list(await connect_streams(compute_session))
    app_ctx.stream_stdin_socks[stream_key].add(socks[0])
    defer(lambda: app_ctx.stream_stdin_socks[stream_key].discard(socks[0]))
    stream_sync = asyncio.Event()

    async def stream_stdin() -> None:
        nonlocal socks
        try:
            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    data = load_json(msg.data)
                    if data["type"] == "stdin":
                        raw_data = base64.b64decode(data["chars"].encode("ascii"))
                        try:
                            await socks[0].send_multipart([raw_data])
                        except (RuntimeError, zmq.error.ZMQError):
                            # when socks[0] is closed, re-initiate the connection.
                            app_ctx.stream_stdin_socks[stream_key].discard(socks[0])
                            socks[1].close()
                            stdin_sock, stdout_sock = await connect_streams(compute_session)
                            socks[0] = stdin_sock
                            socks[1] = stdout_sock
                            app_ctx.stream_stdin_socks[stream_key].add(socks[0])
                            await socks[0].send_multipart([raw_data])
                            log.debug("stream_stdin({0}): zmq stream reset", stream_key)
                            stream_sync.set()
                            continue
                    else:
                        await asyncio.shield(
                            database_ptask_group.create_task(
                                root_ctx.registry.increment_session_usage(session),
                            ),
                        )
                        run_id = secrets.token_hex(8)
                        if data["type"] == "resize":
                            code = f"%resize {data['rows']} {data['cols']}"
                            await root_ctx.registry.execute(
                                session,
                                api_version,
                                run_id,
                                "query",
                                code,
                                {},
                                flush_timeout=None,
                            )
                        elif data["type"] == "ping":
                            await root_ctx.registry.execute(
                                session,
                                api_version,
                                run_id,
                                "query",
                                "%ping",
                                {},
                                flush_timeout=None,
                            )
                        elif data["type"] == "restart":
                            # Close existing zmq sockets and let stream
                            # handlers get a new one with changed stdin/stdout
                            # ports.
                            log.debug("stream_stdin: restart requested")
                            if not socks[0].closed:
                                await asyncio.shield(
                                    database_ptask_group.create_task(
                                        root_ctx.registry.restart_session(
                                            session,
                                        ),
                                    ),
                                )
                                socks[0].close()
                            else:
                                log.warning(
                                    "stream_stdin({0}): "
                                    "duplicate kernel restart request; "
                                    "ignoring it.",
                                    stream_key,
                                )
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    log.warning(
                        "stream_stdin({0}): connection closed ({1})", stream_key, ws.exception()
                    )
        except asyncio.CancelledError:
            # Agent or kernel is terminated.
            raise
        except Exception:
            await root_ctx.error_monitor.capture_exception(
                context={"user": request["user"]["uuid"]}
            )
            log.exception("stream_stdin({0}): unexpected error", stream_key)
        finally:
            log.debug("stream_stdin({0}): terminated", stream_key)
            if not socks[0].closed:
                socks[0].close()

    async def stream_stdout() -> None:
        nonlocal socks
        log.debug("stream_stdout({0}): started", stream_key)
        data: list[bytes] = []
        try:
            while True:
                try:
                    data = await socks[1].recv_multipart()
                except (asyncio.CancelledError, zmq.error.ZMQError):
                    if socks[0] not in app_ctx.stream_stdin_socks:
                        # we are terminating
                        return
                    # connection is closed, so wait until stream_stdin() recovers it.
                    await stream_sync.wait()
                    stream_sync.clear()
                    log.debug("stream_stdout({0}): zmq stream reset", stream_key)
                    continue
                if ws.closed:
                    break
                await ws.send_bytes(
                    dump_json({
                        "type": "out",
                        "data": base64.b64encode(data[0]).decode("ascii"),
                    })
                )
        except asyncio.CancelledError:
            pass
        except Exception:
            await root_ctx.error_monitor.capture_exception(
                context={"user": request["user"]["uuid"]}
            )
            log.exception("stream_stdout({0}): unexpected error", stream_key)
        finally:
            log.debug("stream_stdout({0}): terminated", stream_key)
            socks[1].close()

    # According to aiohttp docs, reading ws must be done inside this task.
    # We execute the stdout handler as another task.
    stdout_task = asyncio.create_task(stream_stdout())
    try:
        await stream_stdin()
    except Exception:
        await root_ctx.error_monitor.capture_exception(context={"user": request["user"]["uuid"]})
        log.exception("stream_pty({0}): unexpected error", stream_key)
    finally:
        stdout_task.cancel()
        await stdout_task
    return ws


@server_status_required(READ_ALLOWED)
@auth_required
@adefer
async def stream_execute(
    defer: Callable[[Callable[[], None]], None], request: web.Request
) -> web.StreamResponse:
    """
    WebSocket-version of gateway.kernel.execute().
    """
    root_ctx: RootContext = request.app["_root.context"]
    app_ctx: PrivateContext = request.app["stream.context"]
    database_ptask_group: aiotools.PersistentTaskGroup = request.app["database_ptask_group"]
    rpc_ptask_group: aiotools.PersistentTaskGroup = request.app["rpc_ptask_group"]

    config = root_ctx.config_provider.config
    registry = root_ctx.registry
    session_name = request.match_info["session_name"]
    access_key = request["keypair"]["access_key"]
    api_version = request["api_version"]
    log.info("STREAM_EXECUTE(ak:{0}, s:{1})", access_key, session_name)
    try:
        async with root_ctx.db.begin_readonly_session() as db_sess:
            session: SessionRow = await asyncio.shield(
                database_ptask_group.create_task(
                    SessionRow.get_session(
                        db_sess,
                        session_name,
                        access_key,
                        kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
                    ),
                ),
            )
    except SessionNotFound:
        raise
    stream_key = session.main_kernel.id

    await asyncio.shield(
        database_ptask_group.create_task(
            registry.increment_session_usage(session),
        )
    )
    ws = web.WebSocketResponse(max_msg_size=config.manager.max_wsmsg_size)
    await ws.prepare(request)

    myself = asyncio.current_task()
    if myself is None:
        raise NoCurrentTaskContext("No current task context")
    app_ctx.stream_execute_handlers[stream_key].add(myself)
    defer(lambda: app_ctx.stream_execute_handlers[stream_key].discard(myself))

    # This websocket connection itself is a "run".
    run_id = secrets.token_hex(8)

    try:
        if ws.closed:
            log.debug("STREAM_EXECUTE: client disconnected (cancelled)")
            return ws
        params = await ws.receive_json()
        if not params.get("mode"):
            raise InvalidStreamMode("mode is missing or empty!")
        mode = params["mode"]
        if mode not in {"query", "batch"}:
            raise InvalidStreamMode("mode has an invalid value.")
        code = params.get("code", "")
        opts = params.get("options", None) or {}

        while True:
            # TODO: rewrite agent and kernel-runner for unbuffered streaming.
            raw_result = await registry.execute(
                session, api_version, run_id, mode, code, opts, flush_timeout=0.2
            )
            if ws.closed:
                log.debug("STREAM_EXECUTE: client disconnected (interrupted)")  # type: ignore[unreachable]
                await asyncio.shield(
                    rpc_ptask_group.create_task(
                        registry.interrupt_session(session),
                    )
                )
                break
            await ws.send_json({
                "status": raw_result["status"],
                "console": raw_result.get("console"),
                "exitCode": raw_result.get("exitCode"),
                "options": raw_result.get("options"),
                "files": raw_result.get("files"),
            })
            if raw_result["status"] == "waiting-input":
                mode = "input"
                code = await ws.receive_str()
            elif raw_result["status"] == "finished":
                break
            else:
                # repeat until we get finished
                mode = "continue"
                code = ""
                opts.clear()
    except (json.decoder.JSONDecodeError, AssertionError) as e:
        log.warning("STREAM_EXECUTE: invalid/missing parameters: {0!r}", e)
        if not ws.closed:
            await ws.send_json({
                "status": "error",
                "msg": f"Invalid API parameters: {e!r}",
            })
    except BackendAIError as e:
        log.exception("STREAM_EXECUTE: exception")
        if not ws.closed:
            await ws.send_json({
                "status": "error",
                "msg": f"BackendError: {e!r}",
            })
        raise
    except asyncio.CancelledError:
        if not ws.closed:
            await ws.send_json({
                "status": "server-restarting",
                "msg": (
                    "The API server is going to restart for maintenance. "
                    "Please connect again with the same run ID."
                ),
            })
        raise
    finally:
        return ws


@server_status_required(READ_ALLOWED)
@auth_required
@check_api_params(
    t.Dict({
        tx.AliasedKey(["app", "service"]): t.String,
        # The port argument is only required to use secondary ports
        # when the target app listens multiple TCP ports.
        # Otherwise it should be omitted or set to the same value of
        # the actual port number used by the app.
        tx.AliasedKey(["port"], default=None): t.Null | t.Int[1024:65535],
        tx.AliasedKey(["envs"], default=None): t.Null | t.String,  # stringified JSON
        # e.g., '{"PASSWORD": "12345"}'
        tx.AliasedKey(["arguments"], default=None): t.Null | t.String,  # stringified JSON
        # e.g., '{"-P": "12345"}'
        # The value can be one of:
        # None, str, List[str]
    })
)
@adefer
async def stream_proxy(
    defer: Callable[[Callable[[], None]], None], request: web.Request, params: Mapping[str, Any]
) -> web.StreamResponse:
    root_ctx: RootContext = request.app["_root.context"]
    app_ctx: PrivateContext = request.app["stream.context"]
    database_ptask_group: aiotools.PersistentTaskGroup = request.app["database_ptask_group"]
    rpc_ptask_group: aiotools.PersistentTaskGroup = request.app["rpc_ptask_group"]
    session_name: str = request.match_info["session_name"]
    access_key: AccessKey = request["keypair"]["access_key"]
    service: str = params["app"]
    myself = asyncio.current_task()
    if myself is None:
        raise NoCurrentTaskContext("No current task context")
    try:
        async with root_ctx.db.begin_readonly_session() as db_sess:
            session = await asyncio.shield(
                database_ptask_group.create_task(
                    SessionRow.get_session(
                        db_sess,
                        session_name,
                        access_key,
                        kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
                    ),
                )
            )
    except (SessionNotFound, TooManySessionsMatched):
        raise
    kernel: KernelRow = session.main_kernel
    kernel_id = kernel.id
    session_id = SessionId(session.id)
    stream_key = kernel_id
    stream_id = uuid.uuid4().hex
    app_ctx.stream_proxy_handlers[stream_key].add(myself)
    defer(lambda: app_ctx.stream_proxy_handlers[stream_key].discard(myself))
    kernel_host: str
    if kernel.kernel_host is None:
        hostname = urlparse(kernel.agent_addr).hostname
        if hostname is None:
            raise InvalidAPIParameters(
                f"Cannot determine kernel host from agent address: {kernel.agent_addr}"
            )
        kernel_host = hostname.decode() if isinstance(hostname, bytes) else hostname
    else:
        kernel_host = kernel.kernel_host
    service_ports: list[dict[str, Any]] = kernel.service_ports or []
    sport: dict[str, Any] = {}
    host_port: int = 0
    dest: tuple[str, int] = ("", 0)
    for sport in service_ports:
        if sport["name"] == service:
            if params["port"]:
                # using one of the primary/secondary ports of the app
                try:
                    hport_idx = sport["container_ports"].index(params["port"])
                except ValueError as e:
                    raise InvalidAPIParameters(
                        f"Service {service} does not open the port number {params['port']}."
                    ) from e
                host_port = sport["host_ports"][hport_idx]
            else:  # using the default (primary) port of the app
                if "host_ports" not in sport:
                    host_port = sport["host_port"]  # legacy kernels
                else:
                    host_port = sport["host_ports"][0]
            dest = (kernel_host, host_port)
            break
    else:
        raise AppNotFound(f"{session_name}:{service}")

    log.info(
        "STREAM_WSPROXY (ak:{}, s:{}): tunneling {}:{} to {}",
        access_key,
        session_name,
        service,
        sport["protocol"],
        "{}:{}".format(*dest),
    )
    if sport["protocol"] == "tcp":
        proxy_cls = TCPProxy
    elif sport["protocol"] == "pty":
        raise NotImplementedError
    elif sport["protocol"] in ("http", "preopen", "vnc", "rdp"):
        proxy_cls = TCPProxy
    else:
        raise InvalidAPIParameters(f"Unsupported service protocol: {sport['protocol']}")

    valkey_live = root_ctx.valkey_live
    conn_tracker_key = f"session.{kernel_id}.active_app_connections"

    async def update_connection_tracker() -> None:
        """Update connection tracker with current timestamp."""
        await valkey_live.update_app_connection_tracker(str(kernel_id), service, stream_id)

    async def refresh_cb(_kernel_id_str: str, _data: bytes) -> None:
        await asyncio.shield(
            rpc_ptask_group.create_task(
                call_non_bursty(
                    conn_tracker_key,
                    update_connection_tracker,
                    max_bursts=128,
                    max_idle=5000,
                ),
            )
        )

    down_cb = apartial(refresh_cb, str(kernel_id))
    up_cb = apartial(refresh_cb, str(kernel_id))
    ping_cb = apartial(refresh_cb, str(kernel_id))

    async def add_conn_track() -> None:
        async with app_ctx.conn_tracker_lock:
            app_ctx.active_session_ids[kernel_id] += 1
            await valkey_live.update_connection_tracker(str(kernel_id), service, stream_id)
            await root_ctx.idle_checker_host.update_app_streaming_status(
                session_id,
                AppStreamingStatus.HAS_ACTIVE_CONNECTIONS,
            )

    async def clear_conn_track() -> None:
        async with app_ctx.conn_tracker_lock:
            app_ctx.active_session_ids[kernel_id] -= 1
            if app_ctx.active_session_ids[kernel_id] <= 0:
                del app_ctx.active_session_ids[kernel_id]
            await valkey_live.remove_connection_tracker(str(kernel_id), service, stream_id)
            remaining_count = await valkey_live.count_active_connections(str(kernel_id))
            if remaining_count == 0:
                await root_ctx.idle_checker_host.update_app_streaming_status(
                    session_id,
                    AppStreamingStatus.NO_ACTIVE_CONNECTIONS,
                )

    try:
        await asyncio.shield(
            database_ptask_group.create_task(
                add_conn_track(),
            )
        )
        await asyncio.shield(
            database_ptask_group.create_task(
                root_ctx.registry.increment_session_usage(session),
            )
        )

        opts: MutableMapping[str, None | str | list[str]] = {}
        if params["arguments"] is not None:
            opts["arguments"] = load_json(params["arguments"])
        if params["envs"] is not None:
            opts["envs"] = load_json(params["envs"])

        result = await asyncio.shield(
            rpc_ptask_group.create_task(
                root_ctx.registry.start_service(session, service, opts),
            ),
        )
        if result["status"] == "failed":
            raise AppServiceStartFailed(
                "Failed to launch the app service", extra_data=result["error"]
            )

        # TODO: weakref to proxies for graceful shutdown?
        ws = web.WebSocketResponse(
            autoping=False,
            max_msg_size=root_ctx.config_provider.config.manager.max_wsmsg_size,
        )
        await ws.prepare(request)
        proxy = proxy_cls(
            ws,
            dest[0],
            dest[1],
            downstream_callback=down_cb,
            upstream_callback=up_cb,
            ping_callback=ping_cb,
        )
        return await proxy.proxy()
    except asyncio.CancelledError:
        log.debug("stream_proxy({}, {}) cancelled", stream_key, service)
        raise
    finally:
        await asyncio.shield(database_ptask_group.create_task(clear_conn_track()))


@server_status_required(READ_ALLOWED)
@auth_required
async def get_stream_apps(request: web.Request) -> web.Response:
    session_name = request.match_info["session_name"]
    access_key = request["keypair"]["access_key"]
    root_ctx: RootContext = request.app["_root.context"]
    async with root_ctx.db.begin_readonly_session() as db_sess:
        compute_session = await SessionRow.get_session(
            db_sess,
            session_name,
            access_key,
            kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
        )
    raw_service_ports = compute_session.main_kernel.service_ports
    if raw_service_ports is None:
        return web.json_response([])
    service_ports: list[dict[str, Any]] = raw_service_ports
    resp = []
    for item in service_ports:
        response_dict: dict[str, Any] = {
            "name": item["name"],
            "protocol": item["protocol"],
            "ports": item["container_ports"],
        }
        if "url_template" in item.keys():
            response_dict["url_template"] = item["url_template"]
        if "allowed_arguments" in item.keys():
            response_dict["allowed_arguments"] = item["allowed_arguments"]
        if "allowed_envs" in item.keys():
            response_dict["allowed_envs"] = item["allowed_envs"]
        resp.append(response_dict)
    return web.json_response(resp)


async def handle_kernel_terminating(
    app: web.Application,
    _source: AgentId,
    event: KernelTerminatingBroadcastEvent,
) -> None:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["stream.context"]
    try:
        kernel = await KernelRow.get_kernel(
            root_ctx.db,
            event.kernel_id,
            allow_stale=True,
        )
    except SessionNotFound:
        return
    if kernel.cluster_role == DEFAULT_ROLE:
        stream_key = kernel.id
        cancelled_tasks = []
        for sock in app_ctx.stream_stdin_socks[stream_key]:
            sock.close()
        for handler in list(app_ctx.stream_pty_handlers.get(stream_key, [])):
            handler.cancel()
            cancelled_tasks.append(handler)
        for handler in list(app_ctx.stream_execute_handlers.get(stream_key, [])):
            handler.cancel()
            cancelled_tasks.append(handler)
        for handler in list(app_ctx.stream_proxy_handlers.get(stream_key, [])):
            handler.cancel()
            cancelled_tasks.append(handler)
        await asyncio.gather(*cancelled_tasks, return_exceptions=True)
        # TODO: reconnect if restarting?


async def stream_conn_tracker_gc(root_ctx: RootContext, app_ctx: PrivateContext) -> None:
    valkey_live = root_ctx.valkey_live
    try:
        while True:
            try:
                # TODO: Use `no_packet_timeout` from the shared config after resolving type issue.
                no_packet_timeout: timedelta = tx.TimeDuration().check(
                    await root_ctx.etcd.get("config/idle/app-streaming-packet-timeout") or "5m",
                )
            except GRPCStatusError as e:
                err_detail = e.args[0]

                if err_detail["code"] == GRPCStatusCode.Unavailable:
                    log.warning(
                        "stream_conn_tracker_gc(): error while connecting to Etcd server,"
                        " retrying..."
                    )
                    continue
                raise e
            async with app_ctx.conn_tracker_lock:
                now = await valkey_live.get_server_time()
                for session_id in app_ctx.active_session_ids.keys():
                    prev_remaining_count = await valkey_live.count_active_connections(
                        str(session_id)
                    )
                    removed_count = await valkey_live.remove_stale_connections(
                        str(session_id),
                        now - no_packet_timeout.total_seconds(),
                    )
                    remaining_count = await valkey_live.count_active_connections(str(session_id))
                    log.debug(
                        f"conn_tracker: gc {session_id} "
                        f"removed/remaining = {removed_count}/{remaining_count}",
                    )
                    if prev_remaining_count > 0 and remaining_count == 0:
                        # Note: kernel_id is used as session_id key here for connection tracking
                        # The idle checker operates on session granularity
                        await root_ctx.idle_checker_host.update_app_streaming_status(
                            SessionId(session_id),
                            AppStreamingStatus.NO_ACTIVE_CONNECTIONS,
                        )
            await asyncio.sleep(10)
    except asyncio.CancelledError:
        pass


@attrs.define(slots=True, auto_attribs=True, init=False)
class PrivateContext:
    stream_pty_handlers: defaultdict[KernelId, weakref.WeakSet[asyncio.Task[Any]]]
    stream_execute_handlers: defaultdict[KernelId, weakref.WeakSet[asyncio.Task[Any]]]
    stream_proxy_handlers: defaultdict[KernelId, weakref.WeakSet[asyncio.Task[Any]]]
    stream_stdin_socks: defaultdict[KernelId, weakref.WeakSet[zmq.asyncio.Socket]]
    zctx: zmq.asyncio.Context
    conn_tracker_lock: asyncio.Lock
    conn_tracker_gc_task: asyncio.Task[Any]
    active_session_ids: defaultdict[KernelId, int]


async def stream_app_ctx(app: web.Application) -> AsyncIterator[None]:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["stream.context"]

    app_ctx.stream_pty_handlers = defaultdict(weakref.WeakSet)
    app_ctx.stream_execute_handlers = defaultdict(weakref.WeakSet)
    app_ctx.stream_proxy_handlers = defaultdict(weakref.WeakSet)
    app_ctx.stream_stdin_socks = defaultdict(weakref.WeakSet)
    app_ctx.zctx = zmq.asyncio.Context()
    app_ctx.conn_tracker_lock = asyncio.Lock()
    app_ctx.active_session_ids = defaultdict(int)  # multiset[int]
    app_ctx.conn_tracker_gc_task = asyncio.create_task(stream_conn_tracker_gc(root_ctx, app_ctx))

    root_ctx.event_dispatcher.subscribe(
        KernelTerminatingBroadcastEvent, app, handle_kernel_terminating
    )

    yield

    # The shutdown handler below is called before this cleanup.
    app_ctx.zctx.term()


async def stream_shutdown(app: web.Application) -> None:
    database_ptask_group: aiotools.PersistentTaskGroup = app["database_ptask_group"]
    rpc_ptask_group: aiotools.PersistentTaskGroup = app["rpc_ptask_group"]
    await database_ptask_group.shutdown()
    await rpc_ptask_group.shutdown()
    cancelled_tasks: list[asyncio.Task[Any]] = []
    app_ctx: PrivateContext = app["stream.context"]
    app_ctx.conn_tracker_gc_task.cancel()
    cancelled_tasks.append(app_ctx.conn_tracker_gc_task)
    for per_kernel_handlers in app_ctx.stream_pty_handlers.values():
        for handler in list(per_kernel_handlers):
            if not handler.done():
                handler.cancel()
                cancelled_tasks.append(handler)
    for per_kernel_handlers in app_ctx.stream_execute_handlers.values():
        for handler in list(per_kernel_handlers):
            if not handler.done():
                handler.cancel()
                cancelled_tasks.append(handler)
    for per_kernel_handlers in app_ctx.stream_proxy_handlers.values():
        for handler in list(per_kernel_handlers):
            if not handler.done():
                handler.cancel()
                cancelled_tasks.append(handler)
    await asyncio.gather(*cancelled_tasks, return_exceptions=True)


def create_app(
    default_cors_options: CORSOptions,
) -> tuple[web.Application, Iterable[WebMiddleware]]:
    app = web.Application()
    app.cleanup_ctx.append(stream_app_ctx)
    app.on_shutdown.append(stream_shutdown)
    app["prefix"] = "stream"
    app["api_versions"] = (2, 3, 4)
    app["stream.context"] = PrivateContext()
    app["database_ptask_group"] = aiotools.PersistentTaskGroup()
    app["rpc_ptask_group"] = aiotools.PersistentTaskGroup()
    cors = aiohttp_cors.setup(app, defaults=default_cors_options)
    add_route = app.router.add_route
    cors.add(add_route("GET", r"/session/{session_name}/pty", stream_pty))
    cors.add(add_route("GET", r"/session/{session_name}/execute", stream_execute))
    cors.add(add_route("GET", r"/session/{session_name}/apps", get_stream_apps))
    # internally both tcp/http proxies use websockets as API/agent-level transports,
    # and thus they have the same implementation here.
    cors.add(add_route("GET", r"/session/{session_name}/httpproxy", stream_proxy))
    cors.add(add_route("GET", r"/session/{session_name}/tcpproxy", stream_proxy))
    return app, []
