"""
Contains registries for coordinating active WS connections (aka "channels", see
`./channels/README.md`).

NOTE: in AWS we had to turn on what amounts to sticky sessions for frontend apps,
so that an individual frontend app instance is guaranteed to always connect to
the same backend api instance. This is because the two registries here are
tied together via a `client_id` string.

The tale-of-the-tape is this:
  - frontend app requires two different channels (WS connections) to the backend api
    - one dedicated to streaming VNC's RFB protocol
    - the other dedicated to messaging (JSON)
  - both of these channels are stateful and need to coordinate with one another

Additionally, this module manages:
  - CDP input channels for interactive browser control
  - Stream reference counts that defer browser state cleanup while active
    CDP streams hold references to a workflow run's browser
"""

from __future__ import annotations

import typing as t

import structlog

if t.TYPE_CHECKING:
    from skyvern.forge.sdk.routes.streaming.cdp_input import CdpInputChannel
    from skyvern.forge.sdk.routes.streaming.channels.message import MessageChannel
    from skyvern.forge.sdk.routes.streaming.channels.vnc import VncChannel

LOG = structlog.get_logger()


# a registry for VNC channels, keyed by `client_id`
vnc_channels: dict[str, VncChannel] = {}


def add_vnc_channel(vnc_channel: VncChannel) -> None:
    vnc_channels[vnc_channel.client_id] = vnc_channel


def get_vnc_channel(client_id: str) -> VncChannel | None:
    return vnc_channels.get(client_id)


def del_vnc_channel(client_id: str, *, expected: VncChannel | None = None) -> None:
    candidate = vnc_channels.get(client_id)

    if candidate is None:
        return

    # Prevent stale channel shutdown from deleting a newer channel that reused
    # the same client_id during route transitions/reconnects.
    if expected is not None and candidate is not expected:
        return

    del vnc_channels[client_id]


# a registry for message channels, keyed by `client_id`
message_channels: dict[str, MessageChannel] = {}


def add_message_channel(message_channel: MessageChannel) -> None:
    message_channels[message_channel.client_id] = message_channel


def get_message_channel(client_id: str) -> MessageChannel | None:
    candidate = message_channels.get(client_id)

    if candidate is None:
        return None

    if candidate.is_open:
        return candidate

    LOG.info(
        "MessageChannel: message channel is not open; deleting it",
        client_id=candidate.client_id,
    )
    del_message_channel(candidate.client_id, expected=candidate)
    return None


def del_message_channel(client_id: str, *, expected: MessageChannel | None = None) -> None:
    candidate = message_channels.get(client_id)

    if candidate is None:
        return

    # Prevent stale channel shutdown from deleting a newer channel that reused
    # the same client_id during route transitions/reconnects.
    if expected is not None and candidate is not expected:
        return

    del message_channels[client_id]


# Stream reference counts per workflow_run_id.
_stream_refcounts: dict[str, int] = {}
_deferred_close_params: dict[str, bool] = {}


def stream_ref_inc(workflow_run_id: str) -> None:
    _stream_refcounts[workflow_run_id] = _stream_refcounts.get(workflow_run_id, 0) + 1


async def stream_ref_dec(workflow_run_id: str) -> None:
    count = _stream_refcounts.get(workflow_run_id, 0) - 1
    if count <= 0:
        _stream_refcounts.pop(workflow_run_id, None)
        close_on_completion = _deferred_close_params.pop(workflow_run_id, None)
        if close_on_completion is not None:
            from skyvern.forge import app

            browser_state = app.BROWSER_MANAGER.pages.get(workflow_run_id)
            if browser_state is not None:
                try:
                    await browser_state.close(close_browser_on_completion=close_on_completion)
                except Exception:
                    LOG.warning(
                        "stream_ref_dec: error closing deferred browser state",
                        workflow_run_id=workflow_run_id,
                        exc_info=True,
                    )
            app.BROWSER_MANAGER.evict_page(workflow_run_id)
    else:
        _stream_refcounts[workflow_run_id] = count


def stream_ref_active(workflow_run_id: str) -> bool:
    return _stream_refcounts.get(workflow_run_id, 0) > 0


def set_deferred_close_params(workflow_run_id: str, close_browser_on_completion: bool) -> None:
    _deferred_close_params[workflow_run_id] = close_browser_on_completion


# a registry for CDP input channels, keyed by `client_id`
cdp_input_channels: dict[str, CdpInputChannel] = {}


def add_cdp_input_channel(channel: CdpInputChannel) -> None:
    cdp_input_channels[channel.client_id] = channel


def del_cdp_input_channel(client_id: str) -> None:
    cdp_input_channels.pop(client_id, None)
