"""Pytest Plugin for owsclient."""
import weakref
from typing import Iterator, TypeAlias, cast

import pytest
import respx

from owsclient import ImpersonationM2MTokenManager
from owsclient.base import (
    _ASYNC_OWS_CLIENT_REGISTRY,
    _OWS_CLIENT_REGISTRY,
    AsyncOwsClient,
    OwsClient,
)
from owsclient.impersonation_client import (
    _IMPERSONATION_OWS_CLIENT_REGISTRY,
    ImpersonationOwsClient,
)
from owsclient.protocols import AsyncM2MTokenManager, M2MTokenManager
from owsclient.test.mock import OwsClientMock

AnyOwsClient: TypeAlias = OwsClient | AsyncOwsClient


def pytest_configure(config: pytest.Config) -> None:
    """Configure plugin for pytest."""
    config.addinivalue_line(
        "markers",
        "ows_client(assert_all_called=False, assert_all_mocked=False): "
        "configure the ows_client_mock fixture. ",
    )

    config.addinivalue_line(
        "markers",
        "ows_client_mock_disable_m2m_patch: "
        "disable m2m_token_manager patching in ows_client_mock fixture",
    )


def mock_m2m_token_manager_with(
    m2m_token_manager_mock: M2MTokenManager,
) -> weakref.WeakKeyDictionary[OwsClient, M2MTokenManager | None]:
    """Replace the M2MTokenManager attributes in the registered owclients."""
    original_managers: weakref.WeakKeyDictionary[
        OwsClient, M2MTokenManager | None
    ] = weakref.WeakKeyDictionary()

    for _ows_client in _OWS_CLIENT_REGISTRY:
        original_managers[_ows_client] = _ows_client.m2m_token_manager
        _ows_client.m2m_token_manager = m2m_token_manager_mock

    return original_managers


def mock_async_m2m_token_manager_with(
    m2m_token_manager_mock: AsyncM2MTokenManager,
) -> weakref.WeakKeyDictionary[AsyncOwsClient, AsyncM2MTokenManager | None]:
    """Replace the AsyncM2MTokenManager attributes in the registered async owclients."""
    original_managers: weakref.WeakKeyDictionary[
        AsyncOwsClient, AsyncM2MTokenManager | None
    ] = weakref.WeakKeyDictionary()

    for _async_ows_client in _ASYNC_OWS_CLIENT_REGISTRY:
        original_managers[_async_ows_client] = _async_ows_client.m2m_token_manager
        _async_ows_client.m2m_token_manager = m2m_token_manager_mock

    return original_managers


def mock_impersonation_m2m_token_manager_with(
    m2m_token_manager_mock: ImpersonationM2MTokenManager,
) -> weakref.WeakKeyDictionary[ImpersonationOwsClient, ImpersonationM2MTokenManager]:
    """Replace the ImpersonationM2MTokenManager attributes in the registered async owclients.

    NOTE: `m2m_token_manager` cannot be null in an `ImpersonationOwsClient` object, so the
    `WeakKeyDictionary` does not support null values.
    """
    original_impersonation_managers: weakref.WeakKeyDictionary[
        ImpersonationOwsClient, ImpersonationM2MTokenManager
    ] = weakref.WeakKeyDictionary()

    for _impersonation_ows_client in _IMPERSONATION_OWS_CLIENT_REGISTRY:
        original_impersonation_managers[
            _impersonation_ows_client
        ] = _impersonation_ows_client.m2m_token_manager
        _impersonation_ows_client.m2m_token_manager = m2m_token_manager_mock

    return original_impersonation_managers


@pytest.fixture
def ows_client_mock(request: pytest.FixtureRequest) -> Iterator[OwsClientMock]:
    """Reusable test fixture for interacting with owsclient."""
    ows_client_marker = request.node.get_closest_marker("ows_client")
    disable_patch_marker = request.node.get_closest_marker(
        "ows_client_mock_disable_m2m_patch"
    )
    ows_client_mock = OwsClientMock()
    mock_router = (
        ows_client_mock.mock
        if ows_client_marker is None
        else cast(respx.MockRouter, ows_client_mock.mock(**ows_client_marker.kwargs))
    )

    # Store original managers with weak references to clients
    original_managers, original_async_managers, original_impersonation_managers = (
        None,
        None,
        None,
    )

    if disable_patch_marker is None:
        original_managers = mock_m2m_token_manager_with(
            ows_client_mock.mock_m2m_token_manager
        )

        original_async_managers = mock_async_m2m_token_manager_with(
            ows_client_mock.mock_async_m2m_token_manager
        )

        original_impersonation_managers = mock_impersonation_m2m_token_manager_with(
            ows_client_mock.mock_impersonation_m2m_token_manager
        )

    with mock_router:
        yield ows_client_mock

    if disable_patch_marker is None:
        if original_managers:
            for _ows_client, original_manager in original_managers.items():
                _ows_client.m2m_token_manager = original_manager

        if original_async_managers:
            for (
                _async_ows_client,
                original_async_manager,
            ) in original_async_managers.items():
                _async_ows_client.m2m_token_manager = original_async_manager

        if original_impersonation_managers:
            for (
                _impersonation_ows_client,
                original_impersonation_manager,
            ) in original_impersonation_managers.items():
                _impersonation_ows_client.m2m_token_manager = (
                    original_impersonation_manager
                )