import abc
import asyncio
import concurrent.futures
import logging
import os
import platform
import socket
import sys
from collections.abc import (
    AsyncGenerator,
    Awaitable,
    Callable,
    Generator,
    Iterable,
    Mapping,
)
from contextlib import contextmanager, suppress
from functools import partial, wraps
from inspect import isasyncgenfunction, iscoroutinefunction
from types import ModuleType
from typing import Any, NamedTuple
from unittest.mock import MagicMock

import aiomisc
import pytest
from aiomisc_log import LOG_LEVEL, basic_config

from .compat import sock_set_reuseport
from .context_vars import set_current_loop
from .utils import TimeoutType, bind_socket


log = logging.getLogger("aiomisc_pytest")


USE_UVLOOP = os.getenv("AIOMISC_USE_UVLOOP", "1").lower() in (
    "1",
    "yes",
    "true",
)

uvloop_module: ModuleType | None

try:
    if USE_UVLOOP:
        import uvloop as _uvloop

        uvloop_module = _uvloop
except ImportError:
    uvloop_module = None


log = logging.getLogger(__name__)
ProxyProcessorType = Callable[[bytes], Awaitable[bytes]]


class Delay:
    __slots__ = "__timeout", "future", "lock"

    def __init__(self) -> None:
        self.__timeout: int | float = 0
        self.future: asyncio.Future | None = None
        self.lock = asyncio.Lock()

    @property
    def timeout(self) -> int | float:
        return self.__timeout

    @timeout.setter
    def timeout(self, value: TimeoutType) -> None:
        assert isinstance(value, (int, float))
        assert value >= 0

        self.__timeout = value

        if self.future and not self.future.done():
            self.future.set_result(True)

    async def wait(self) -> None:
        if self.__timeout == 0:
            return

        async with self.lock:
            try:
                self.future = delayed_future(self.__timeout)
                await self.future
            finally:
                self.future = None


def delayed_future(timeout: int | float, result: bool = True) -> asyncio.Future:

    loop = asyncio.get_event_loop()

    def resolve(f: asyncio.Future) -> None:
        nonlocal result  # noqa

        if f.done():
            return
        f.set_result(result)

    future = loop.create_future()
    handle = loop.call_later(timeout, resolve, future)
    future.add_done_callback(lambda _: handle.cancel())

    return future


class TCPProxy:
    DEFAULT_TIMEOUT = 30

    __slots__ = (
        "proxy_port",
        "clients",
        "server",
        "proxy_host",
        "target_port",
        "target_host",
        "listen_host",
        "read_delay",
        "write_delay",
        "read_processor",
        "write_processor",
        "buffered",
    )

    def __init__(
        self,
        target_host: str,
        target_port: int,
        listen_host: str = "127.0.0.1",
        buffered: bool = True,
    ):
        self.target_port = target_port
        self.target_host = target_host
        self.proxy_host = listen_host
        self.proxy_port = unused_port()
        self.read_delay: TimeoutType = 0
        self.write_delay: TimeoutType = 0
        self.buffered = buffered

        self.clients: set[TCPProxyClient] = set()
        self.server: asyncio.AbstractServer | None = None

        self.read_processor: ProxyProcessorType | None = None
        self.write_processor: ProxyProcessorType | None = None

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}[{id(self):x}]: tcp://{self.proxy_host}:{self.proxy_port} => tcp://{self.target_host}:{self.target_port}>"

    async def start(
        self, timeout: TimeoutType | None = None
    ) -> asyncio.AbstractServer:
        log.debug("Starting %r", self)
        server = await asyncio.wait_for(
            asyncio.start_server(
                self._handle_client, host=self.proxy_host, port=self.proxy_port
            ),
            timeout=timeout,
        )
        self.server = server
        return server

    ClientType = tuple[asyncio.StreamReader, asyncio.StreamWriter]

    async def create_client(self) -> ClientType:
        log.debug("Creating client for %r", self)
        return await asyncio.open_connection(self.proxy_host, self.proxy_port)

    async def __aenter__(self) -> "TCPProxy":
        if self.server is None:
            await self.start(timeout=self.DEFAULT_TIMEOUT)
        return self

    async def __aexit__(
        self, exc_type: type[Exception], exc_val: Exception, exc_tb: Any
    ) -> None:
        await self.close(timeout=self.DEFAULT_TIMEOUT)

    async def close(self, timeout: TimeoutType | None = None) -> None:
        async def close() -> None:
            await self.disconnect_all()

            if self.server is None:
                return

            self.server.close()
            await self.server.wait_closed()

        await asyncio.wait_for(close(), timeout=timeout)

    def set_delay(
        self, read_delay: TimeoutType, write_delay: TimeoutType = 0
    ) -> None:
        log.debug("Setting delay [R/W %f %f]", read_delay, write_delay)

        for client in self.clients:
            log.debug(
                "Applying delays [R/W: %f %f] for %r",
                read_delay,
                write_delay,
                client,
            )
            client.read_delay.timeout = read_delay
            client.write_delay.timeout = write_delay

        self.read_delay = read_delay
        self.write_delay = write_delay

    def set_content_processors(
        self, read: ProxyProcessorType | None, write: ProxyProcessorType | None
    ) -> None:
        log.debug(
            "Setting content processors for %r: read=%r write=%r",
            self,
            read,
            write,
        )

        for client in self.clients:
            log.debug(
                "Applying context processors for %r: read=%r write=%r",
                client,
                read,
                write,
            )

            client.read_processor = read
            client.write_processor = write

        self.read_processor = read
        self.write_processor = write

    def disconnect_all(self) -> asyncio.Future:
        log.debug("Disconnecting %s clients of %r", len(self.clients), self)
        return asyncio.ensure_future(
            asyncio.gather(
                *[client.close() for client in self.clients],
                return_exceptions=True,
            )
        )

    async def _handle_client(
        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        client = TCPProxyClient(reader, writer, buffered=self.buffered)
        self.clients.add(client)

        client.read_delay.timeout = self.read_delay
        client.write_delay.timeout = self.write_delay
        client.read_processor = self.read_processor
        client.write_processor = self.write_processor

        await client.connect(self.target_host, self.target_port)

        client.closing.add_done_callback(lambda _: self.clients.remove(client))

    @contextmanager
    def slowdown(
        self, read_delay: TimeoutType = 0, write_delay: TimeoutType = 0
    ) -> Generator[None, None, None]:
        old_read_delay = self.read_delay
        old_write_delay = self.write_delay

        self.set_delay(read_delay, write_delay)

        try:
            yield
        finally:
            self.set_delay(old_read_delay, old_write_delay)


class TCPProxyClient:
    __slots__ = (
        "client_reader",
        "client_writer",
        "server_reader",
        "server_writer",
        "chunk_size",
        "tasks",
        "loop",
        "read_delay",
        "write_delay",
        "closing",
        "__processors",
        "__client_repr",
        "__server_repr",
        "buffered",
    )

    @staticmethod
    async def _blank_processor(body: bytes) -> bytes:
        return body

    def __init__(
        self,
        client_reader: asyncio.StreamReader,
        client_writer: asyncio.StreamWriter,
        chunk_size: int = 64 * 1024,
        buffered: bool = False,
    ):

        self.loop = asyncio.get_event_loop()
        self.client_reader: asyncio.StreamReader = client_reader
        self.client_writer: asyncio.StreamWriter = client_writer

        self.server_reader: asyncio.StreamReader | None = None
        self.server_writer: asyncio.StreamWriter | None = None

        self.tasks: Iterable[asyncio.Task] = ()
        self.chunk_size = chunk_size  # type: int
        self.read_delay = Delay()
        self.write_delay = Delay()

        self.closing: asyncio.Future = self.loop.create_future()
        self.__processors: dict[str, ProxyProcessorType] = {
            "read": self._blank_processor,
            "write": self._blank_processor,
        }

        self.buffered = bool(buffered)

        self.__client_repr = ""
        self.__server_repr = ""

    @property
    def read_processor(self) -> ProxyProcessorType:
        return self.__processors["read"]

    @read_processor.setter
    def read_processor(self, value: ProxyProcessorType | None) -> None:
        if value is None:
            self.__processors["read"] = self._blank_processor
            return
        self.__processors["read"] = aiomisc.awaitable(value)

    @property
    def write_processor(self) -> ProxyProcessorType:
        return self.__processors["write"]

    @write_processor.setter
    def write_processor(self, value: ProxyProcessorType | None) -> None:
        if value is None:
            self.__processors["write"] = self._blank_processor
            return
        self.__processors["write"] = aiomisc.awaitable(value)

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}[{id(self):x}]: {self.__client_repr} => {self.__server_repr}>"

    @staticmethod
    async def _close_writer(writer: asyncio.StreamWriter) -> None:
        writer.close()
        await writer.wait_closed()

    async def pipe(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        processor: str,
        delay: Delay,
    ) -> None:
        try:
            while not reader.at_eof():
                chunk = await reader.read(self.chunk_size)

                if not chunk:
                    break

                if delay.timeout > 0:
                    log.debug(
                        "%r sleeping %.3f seconds on %s",
                        self,
                        delay.timeout,
                        processor,
                    )

                    await delay.wait()

                writer.write(await self.__processors[processor](chunk))

                if not self.buffered:
                    await writer.drain()
        finally:
            await self._close_writer(writer)

    async def connect(self, target_host: str, target_port: int) -> None:
        log.debug("Establishing connection for %r", self)

        self.server_reader, self.server_writer = await asyncio.open_connection(
            host=target_host, port=target_port
        )

        self.__client_repr = ":".join(
            map(str, self.client_writer.get_extra_info("peername")[:2])
        )
        self.__server_repr = ":".join(
            map(str, self.server_writer.get_extra_info("peername")[:2])
        )

        self.tasks = (
            self.loop.create_task(
                self.pipe(
                    self.server_reader,
                    self.client_writer,
                    "write",
                    self.write_delay,
                )
            ),
            self.loop.create_task(
                self.pipe(
                    self.client_reader,
                    self.server_writer,
                    "read",
                    self.read_delay,
                )
            ),
        )

    async def close(self) -> None:
        log.debug("Closing %r", self)
        if self.closing.done():
            return

        await aiomisc.cancel_tasks(self.tasks)
        self.loop.call_soon(self.closing.set_result, True)
        await self.closing


def unused_port(*args: Any) -> int:
    with socket.socket(*args) as sock:
        sock.bind(("", 0))
        return sock.getsockname()[1]


@contextmanager
def mock_get_event_loop() -> Generator[Any, MagicMock, None]:
    original_get_event_loop = asyncio.get_event_loop
    getter_mock = MagicMock(asyncio.get_event_loop)
    getter_mock.side_effect = original_get_event_loop

    try:
        asyncio.get_event_loop = getter_mock
        yield getter_mock
    finally:
        asyncio.get_event_loop = original_get_event_loop


@pytest.fixture(scope="session")
def tcp_proxy() -> type[TCPProxy]:
    return TCPProxy


def isasyncgenerator(func: Callable[..., Any]) -> bool | None:
    if isasyncgenfunction(func):
        return True
    if iscoroutinefunction(func):
        return False
    return None


def pytest_fixture_setup(fixturedef, request):  # type: ignore
    func = fixturedef.func

    is_async_gen = isasyncgenerator(func)

    if is_async_gen is None:
        return

    strip_request = False
    if "request" not in fixturedef.argnames:
        fixturedef.argnames += ("request",)
        strip_request = True

    if "event_loop" not in request.fixturenames:
        raise Exception("`event_loop` fixture required")

    # noinspection PyProtectedMember
    loop_fixturedef = request._get_active_fixturedef("event_loop")

    def wrapper(*args, **kwargs):  # type: ignore
        request = kwargs.pop("request") if strip_request else kwargs["request"]

        event_loop = request.getfixturevalue("event_loop")

        if not is_async_gen:
            return event_loop.run_until_complete(func(*args, **kwargs))

        gen = func(*args, **kwargs)

        def finalizer():  # type: ignore
            try:
                return event_loop.run_until_complete(gen.__anext__())
            except StopAsyncIteration:  # NOQA
                pass

        loop_fixturedef.addfinalizer(
            partial(fixturedef.finish, request=request)
        )

        request.addfinalizer(finalizer)
        return event_loop.run_until_complete(gen.__anext__())

    fixturedef.func = wrapper


@pytest.fixture(scope="session")
def localhost() -> str:
    params = ((socket.AF_INET, "127.0.0.1"), (socket.AF_INET6, "::1"))
    for family, addr in params:
        with socket.socket(family, socket.SOCK_STREAM) as sock:
            try:
                sock.bind((addr, 0))
            except Exception:
                pass
            else:
                return addr
    raise RuntimeError("localhost unavailable")


@pytest.fixture(scope="session")
def loop_debug(pytestconfig: pytest.Config) -> bool:
    return pytestconfig.getoption("--aiomisc-debug")


@pytest.fixture(scope="session")
def aiomisc_test_timeout(pytestconfig: pytest.Config) -> int | float | None:
    return pytestconfig.getoption("--aiomisc-test-timeout")


@pytest.fixture(autouse=True)
def aiomisc_func_wrap() -> Callable:
    return aiomisc.awaitable


def pytest_pycollect_makeitem(collector, name, obj):  # type: ignore
    if collector.funcnamefilter(name) and iscoroutinefunction(obj):
        return list(collector._genfunctions(name, obj))
    return None


@pytest.mark.tryfirst
def pytest_pyfunc_call(pyfuncitem):  # type: ignore
    if not iscoroutinefunction(pyfuncitem.function):
        return None

    event_loop = pyfuncitem.funcargs.get("event_loop", None)
    func_wraper = pyfuncitem.funcargs.get(
        "aiomisc_func_wrap", aiomisc.awaitable
    )

    aiomisc_test_timeout = pyfuncitem.funcargs.get("aiomisc_test_timeout", None)

    kwargs = {
        arg: pyfuncitem.funcargs[arg]
        for arg in pyfuncitem._fixtureinfo.argnames
    }

    @wraps(pyfuncitem.obj)
    async def func() -> Any:
        return await asyncio.wait_for(
            func_wraper(pyfuncitem.obj)(**kwargs), timeout=aiomisc_test_timeout
        )

    event_loop.run_until_complete(func())

    return True


@pytest.fixture(scope="session")
def thread_pool_size(request: pytest.FixtureRequest) -> int:
    return request.config.getoption("--aiomisc-pool-size")


@pytest.fixture
def services() -> Iterable[aiomisc.Service]:
    return []


@pytest.fixture
def default_context() -> Mapping[str, Any]:
    return {}


loop_autouse = os.getenv("AIOMISC_LOOP_AUTOUSE", "1") == "1"


@pytest.fixture(name="thread_pool_executor")
def _thread_pool_executor() -> type[concurrent.futures.ThreadPoolExecutor]:
    from aiomisc.thread_pool import ThreadPoolExecutor

    return ThreadPoolExecutor


@pytest.fixture(name="entrypoint_kwargs")
def _entrypoint_kwargs() -> dict:
    return {"log_config": False}


def _create_event_loop() -> asyncio.AbstractEventLoop:
    """Create a new event loop, using uvloop if available and enabled."""
    if USE_UVLOOP and uvloop_module is not None:
        return uvloop_module.new_event_loop()
    return asyncio.new_event_loop()


@pytest.fixture(autouse=loop_autouse)
def event_loop(
    request: pytest.FixtureRequest,
    caplog: pytest.LogCaptureFixture,
    thread_pool_size: int,
    loop_debug: bool,
    thread_pool_executor: Callable[..., concurrent.futures.ThreadPoolExecutor],
) -> Generator[Any, asyncio.AbstractEventLoop, None]:
    basic_config(log_format="plain", stream=caplog.handler.stream)

    get_marker = request.node.get_closest_marker
    forbid_loop_getter_marker = get_marker("forbid_get_event_loop")
    catch_unhandled_marker = get_marker("catch_loop_exceptions")

    exceptions: list[dict[str, Any]] = []

    def catch_exceptions(_: Any, catched: dict[str, Any]) -> None:
        nonlocal exceptions  # noqa
        exceptions.append(catched)

    loop = _create_event_loop()
    loop.set_debug(loop_debug)

    asyncio.set_event_loop(loop)
    set_current_loop(loop)

    pool = thread_pool_executor(thread_pool_size)
    loop.set_default_executor(pool)

    if catch_unhandled_marker:
        loop.set_exception_handler(catch_exceptions)

    if LOG_LEVEL:
        LOG_LEVEL.set(logging.getLogger().getEffectiveLevel())

    try:
        with mock_get_event_loop() as event_loop_getter_mock:
            if forbid_loop_getter_marker:
                event_loop_getter_mock.side_effect = partial(
                    pytest.fail, "get_event_loop is forbidden"
                )
            yield loop
    finally:
        if exceptions:
            logging.error(
                "Unhandled exceptions found:\n\n\t%s",
                "\n\t".join(
                    ("Message: {m}\n\tFuture: {f}\n\tException: {e}").format(
                        m=e["message"],
                        f=repr(e.get("future")),
                        e=repr(e.get("exception")),
                    )
                    for e in exceptions
                ),
            )
            pytest.fail("Unhandled exceptions found. See logs.")

        basic_config(log_format="plain", stream=sys.stderr)

        if not loop.is_closed():
            with suppress(Exception):
                loop.run_until_complete(loop.shutdown_asyncgens())
            with suppress(Exception):
                loop.close()

        asyncio.set_event_loop(None)


@pytest.fixture(autouse=loop_autouse)
async def entrypoint(
    services: Iterable[aiomisc.Service],
    loop_debug: bool,
    default_context: Mapping[str, Any],
    entrypoint_kwargs: Mapping[str, Any],
    thread_pool_size: int,
    thread_pool_executor: Callable[..., concurrent.futures.ThreadPoolExecutor],
    event_loop: asyncio.AbstractEventLoop,
) -> AsyncGenerator[aiomisc.entrypoint, None]:
    from aiomisc.context import get_context
    from aiomisc.entrypoint import entrypoint

    async with entrypoint(
        *services, loop=event_loop, **entrypoint_kwargs
    ) as ep:
        ctx = get_context()
        for key, value in default_context.items():
            ctx[key] = value
        yield ep


class PortSocket(NamedTuple):
    port: int
    socket: socket.socket


@pytest.fixture
def aiomisc_socket_factory(
    request: pytest.FixtureRequest, localhost: str
) -> Callable[..., PortSocket]:
    """Returns a"""

    def factory(*args: Any, **kwargs: Any) -> PortSocket:
        sock = bind_socket(*args, address=localhost, port=0, **kwargs)
        port = sock.getsockname()[1]

        # Close socket after teardown
        request.addfinalizer(sock.close)

        return PortSocket(port=port, socket=sock)

    return factory


class SocketWrapper(abc.ABC):
    address: str
    port: int

    def __init__(self, *args: Any):
        self._socket_args: tuple[Any, ...] = args
        self.address: str = ""
        self.port: int = 0

    @abc.abstractmethod
    def prepare(self, address: str) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    def close(self) -> None:
        pass


class SocketWrapperUnix(SocketWrapper):
    socket: socket.socket
    fd: int

    def __init__(self, *args: Any):
        super().__init__(*args)
        self.socket = socket.socket(*args)
        self.fd = -1

    def close(self) -> None:
        self.socket.close()
        if self.fd > 0:
            os.close(self.fd)

    def prepare(self, address: str) -> None:
        self.socket.bind((address, 0))
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock_set_reuseport(self.socket, True)
        self.address, self.port = self.socket.getsockname()[:2]
        # detach socket object and save file descriptor
        self.fd = self.socket.detach()


class SocketWrapperWindows(SocketWrapper):
    def close(self) -> None:
        pass

    def prepare(self, address: str) -> None:
        self.address = address
        self.port = unused_port(*self._socket_args)


socket_wrapper: type[SocketWrapper]

if platform.system() == "Windows":
    socket_wrapper = SocketWrapperWindows
else:
    socket_wrapper = SocketWrapperUnix


@pytest.fixture
def aiomisc_unused_port_factory(
    request: pytest.FixtureRequest, localhost: str
) -> Callable[[], int]:
    def port_factory(*args: Any) -> int:
        wrapper = socket_wrapper(*args)
        wrapper.prepare("::" if ":" in localhost else "0.0.0.0")
        request.addfinalizer(wrapper.close)
        return wrapper.port

    return port_factory


@pytest.fixture
def aiomisc_unused_port(aiomisc_unused_port_factory: Callable[..., int]) -> int:
    return aiomisc_unused_port_factory()


__all__ = (
    # Classes
    "Delay",
    "PortSocket",
    "ProxyProcessorType",
    "SocketWrapper",
    "SocketWrapperUnix",
    "SocketWrapperWindows",
    "TCPProxy",
    "TCPProxyClient",
    # Functions
    "delayed_future",
    "isasyncgenerator",
    "mock_get_event_loop",
    "unused_port",
    # Pytest hooks (not plugin hooks)
    "pytest_fixture_setup",
    "pytest_pycollect_makeitem",
    "pytest_pyfunc_call",
    # Fixtures
    "aiomisc_func_wrap",
    "aiomisc_socket_factory",
    "aiomisc_test_timeout",
    "aiomisc_unused_port",
    "aiomisc_unused_port_factory",
    "default_context",
    "entrypoint",
    "event_loop",
    "localhost",
    "loop_autouse",
    "loop_debug",
    "services",
    "socket_wrapper",
    "tcp_proxy",
    "thread_pool_size",
    # Internal helpers
    "_create_event_loop",
    "_entrypoint_kwargs",
    "_thread_pool_executor",
)
