"""Testing tools."""

import binascii
import io
import mimetypes
import os
import random
from collections import deque
from contextlib import asynccontextmanager, suppress
from functools import partial
from http.cookies import SimpleCookie
from json import loads
from pathlib import Path
from typing import (Any, AsyncGenerator, Awaitable, Callable, Coroutine, Deque, Dict, Mapping,
                    Optional, Tuple, Union, cast)
from urllib.parse import urlencode

from multidict import MultiDict
from yarl import URL

from asgi_tools import BASE_ENCODING, DEFAULT_CHARSET, ASGIConnectionClosed
from asgi_tools._compat import FIRST_COMPLETED, aio_cancel, aio_sleep, aio_spawn, aio_wait
from asgi_tools.response import Response, ResponseJSON, ResponseWebSocket, parse_websocket_msg
from asgi_tools.types import TJSON, TASGIApp, TASGIMessage, TASGIReceive, TASGIScope, TASGISend
from asgi_tools.utils import CIMultiDict, parse_headers


class TestResponse(Response):
    """Response for test client."""

    def __init__(
        self,
    ):
        super().__init__(b"")

    async def __call__(self, _: TASGIScope, receive: TASGIReceive, __: TASGISend):
        self._receive = receive
        msg = await self._receive()
        assert msg.get("type") == "http.response.start", "Invalid Response"
        self.status_code = int(msg.get("status", 502))
        self.headers = cast(MultiDict, parse_headers(msg.get("headers", [])))
        self.content_type = self.headers.get("content-type")
        for cookie in self.headers.getall("set-cookie", []):
            self.cookies.load(cookie)

    async def stream(self) -> AsyncGenerator[bytes, None]:
        """Stream the response."""
        more_body = True
        while more_body:
            msg = await self._receive()
            if msg.get("type") == "http.response.body":
                chunk = msg.get("body")
                if chunk:
                    yield chunk
                more_body = msg.get("more_body", False)

    async def body(self) -> bytes:
        """Load response body."""
        body_ = b""
        async for chunk in self.stream():
            body_ += chunk
        return body_

    async def text(self) -> str:
        body = await self.body()
        return body.decode(DEFAULT_CHARSET)

    async def json(self) -> TJSON:
        text = await self.text()
        return loads(text)


class TestWebSocketResponse(ResponseWebSocket):
    """Support websockets in tests."""

    def connect(self) -> Coroutine[TASGIMessage, Any, Any]:
        return self.send({"type": "websocket.connect"})

    async def disconnect(self):
        await self.send({"type": "websocket.disconnect", "code": 1005})
        self.state = self.STATES.DISCONNECTED

    def send(self, msg, type="websocket.receive"):  # noqa
        """Send a message to a client."""
        return super().send(msg, type=type)

    async def receive(self, raw=False):
        """Receive messages from a client."""
        if self.partner_state == self.STATES.DISCONNECTED:
            raise ASGIConnectionClosed

        msg = await self._receive()
        if not msg["type"].startswith("websocket."):
            raise ValueError(f"Invalid websocket message: {msg!r}")

        if msg["type"] == "websocket.accept":
            self.partner_state = self.STATES.CONNECTED
            return await self.receive(raw=raw)

        if msg["type"] == "websocket.close":
            self.partner_state = self.STATES.DISCONNECTED
            raise ASGIConnectionClosed("Connection has been closed.")

        return msg if raw else parse_websocket_msg(msg, charset=DEFAULT_CHARSET)


class ASGITestClient:
    """The test client allows you to make requests against an ASGI application.

    Features:

    * cookies
    * multipart/form-data
    * follow redirects
    * request streams
    * response streams
    * websocket support
    * lifespan management

    """

    def __init__(self, app: TASGIApp, base_url: str = "http://localhost"):
        self.app = app
        self.base_url = URL(base_url)
        self.cookies: SimpleCookie = SimpleCookie()
        self.headers: Dict[str, str] = {}

    def __getattr__(self, name: str) -> Callable[..., Awaitable]:
        return partial(self.request, method=name.upper())

    async def request(
        self,
        path: str,
        method: str = "GET",
        query: Union[str, Dict] = "",
        headers: Optional[Dict[str, str]] = None,
        cookies: Optional[Dict[str, str]] = None,
        data: Union[bytes, str, Dict, AsyncGenerator] = b"",
        json: TJSON = None,
        follow_redirect: bool = True,
        timeout: float = 10.0,
    ) -> TestResponse:
        """Make a HTTP requests."""

        headers = headers or dict(self.headers)

        if isinstance(data, str):
            data = Response.process_content(data)

        elif isinstance(data, dict):
            is_multipart = any(isinstance(value, io.IOBase) for value in data.values())
            if is_multipart:
                data, headers["Content-Type"] = encode_multipart(data)

            else:
                headers["Content-Type"] = "application/x-www-form-urlencoded"
                data = urlencode(data).encode(DEFAULT_CHARSET)

        elif json is not None:
            headers["Content-Type"] = "application/json"
            data = ResponseJSON.process_content(json)

        pipe = Pipe()

        if isinstance(data, bytes):
            headers.setdefault("Content-Length", str(len(data)))

        scope = self.build_scope(
            path,
            headers=headers,
            query=query,
            cookies=cookies,
            type="http",
            method=method,
        )

        await aio_wait(
            aio_wait(
                stream_data(data, pipe.send_to_app),
                self.app(scope, pipe.receive_from_app, pipe.send_to_client),
            ),
            raise_timeout(timeout),
            strategy=FIRST_COMPLETED,
        )

        res = TestResponse()
        await res(scope, pipe.receive_from_client, pipe.send_to_app)
        for n, v in res.cookies.items():
            self.cookies[n] = v

        if follow_redirect and res.status_code in {301, 302, 303, 307, 308}:
            return await self.get(res.headers["location"])

        return res

    # TODO: Timeouts for websockets
    @asynccontextmanager
    async def websocket(
        self,
        path: str,
        query: Union[str, Dict, None] = None,
        headers: Optional[Dict] = None,
        cookies: Optional[Dict] = None,
    ):
        """Connect to a websocket."""
        pipe = Pipe()

        ci_headers = CIMultiDict(headers or {})

        scope = self.build_scope(
            path,
            headers=ci_headers,
            query=query,
            cookies=cookies,
            type="websocket",
            subprotocols=str(ci_headers.get("Sec-WebSocket-Protocol", "")).split(","),
        )
        ws = TestWebSocketResponse(scope, pipe.receive_from_client, pipe.send_to_app)
        async with aio_spawn(
            self.app, scope, pipe.receive_from_app, pipe.send_to_client
        ):
            await ws.connect()
            yield ws
            await ws.disconnect()

    def lifespan(self, timeout: float = 3e-2):
        """Manage `Lifespan <https://asgi.readthedocs.io/en/latest/specs/lifespan.html>`_
        protocol."""
        return manage_lifespan(self.app, timeout=timeout)

    def build_scope(
        self,
        path: str,
        headers: Union[Dict, CIMultiDict, None] = None,
        query: Union[str, Dict, None] = None,
        cookies: Optional[Dict] = None,
        **scope,
    ) -> TASGIScope:
        """Prepare a request scope."""
        headers = headers or {}
        headers.setdefault("User-Agent", "ASGI-Tools-Test-Client")
        headers.setdefault("Host", self.base_url.host)

        if cookies:
            for c, v in cookies.items():
                self.cookies[c] = v

        if len(self.cookies):
            headers.setdefault("Cookie", self.cookies.output(header="", sep=";"))

        url = URL(path)
        if query:
            url = url.with_query(query)

        # Setup client
        scope.setdefault("client", ("127.0.0.1", random.randint(1024, 65535)))

        return dict(
            {
                "asgi": {"version": "3.0"},
                "http_version": "1.1",
                "path": url.path,
                "query_string": url.raw_query_string.encode(),
                "raw_path": url.raw_path.encode(),
                "root_path": "",
                "scheme": scope.get("type") == "http" and self.base_url.scheme or "ws",
                "headers": [
                    (key.lower().encode(BASE_ENCODING), str(val).encode(BASE_ENCODING))
                    for key, val in (headers or {}).items()
                ],
                "server": ("127.0.0.1", self.base_url.port),
            },
            **scope,
        )


def encode_multipart(data: Dict) -> Tuple[bytes, str]:
    body = io.BytesIO()
    boundary = binascii.hexlify(os.urandom(16))
    for name, value in data.items():
        headers = f'Content-Disposition: form-data; name="{ name }"'
        if hasattr(value, "read"):
            filename = getattr(value, "name", None)
            if filename:
                headers = f'{ headers }; filename="{ Path(filename).name }"'
                content_type = (
                    mimetypes.guess_type(filename)[0] or "application/octet-stream"
                )
                headers = f"{ headers }\r\nContent-Type: { content_type }"
            value = value.read()

        body.write(b"--%b\r\n" % boundary)
        body.write(headers.encode("utf-8"))
        body.write(b"\r\n\r\n")
        if isinstance(value, str):
            value = value.encode("utf-8")
        body.write(value)
        body.write(b"\r\n")

    body.write(b"--%b--\r\n" % boundary)
    return body.getvalue(), (b"multipart/form-data; boundary=%s" % boundary).decode()


class Pipe:

    __slots__ = (
        "delay",
        "app_is_closed",
        "client_is_closed",
        "app_queue",
        "client_queue",
    )

    def __init__(self, delay: float = 1e-3):
        self.delay = delay
        self.app_is_closed = False
        self.client_is_closed = False
        self.app_queue: Deque[TASGIMessage] = deque()
        self.client_queue: Deque[TASGIMessage] = deque()

    async def send_to_client(self, msg: TASGIMessage):
        if self.client_is_closed:
            raise RuntimeError(f"Unexpected ASGI message to client '{msg.get('type')}'")

        if msg.get("type") == "websocket.close":
            self.client_is_closed = True

        elif msg.get("type") == "http.response.body":
            self.client_is_closed = not msg.get("more_body", False)

        self.client_queue.append(msg)

    async def send_to_app(self, msg: TASGIMessage):
        if self.app_is_closed:
            raise RuntimeError(f"Unexpected ASGI message to app '{msg.get('type')}'")

        if msg.get("type") == "http.disconnect":
            self.app_is_closed = True

        self.app_queue.append(msg)

    async def receive_from_client(self):
        while not self.client_queue:
            await aio_sleep(self.delay)
        return self.client_queue.popleft()

    async def receive_from_app(self):
        while not self.app_queue:
            await aio_sleep(self.delay)
        return self.app_queue.popleft()


async def raise_timeout(timeout: Union[int, float]):
    await aio_sleep(timeout)
    raise TimeoutError("Timeout occured")


async def stream_data(
    data: Union[bytes, AsyncGenerator[Any, bytes]],
    send: Callable[[TASGIMessage], Awaitable[None]],
):
    """Stream a data to an application."""

    if isinstance(data, bytes):
        return await send({"type": "http.request", "body": data, "more_body": False})

    async for chunk in data:
        await send({"type": "http.request", "body": chunk, "more_body": True})
    await send({"type": "http.request", "body": b"", "more_body": False})


@asynccontextmanager
async def manage_lifespan(app, timeout: float = 3e-2):
    """Manage `Lifespan <https://asgi.readthedocs.io/en/latest/specs/lifespan.html>`_ protocol."""
    pipe = Pipe()

    scope = {"type": "lifespan"}

    async def safe_spawn():
        with suppress(BaseException):
            await app(scope, pipe.receive_from_app, pipe.send_to_client)

    async with aio_spawn(safe_spawn) as task:
        await pipe.send_to_app({"type": "lifespan.startup"})
        msg = await aio_wait(
            pipe.receive_from_client(), aio_sleep(timeout), strategy=FIRST_COMPLETED
        )
        if msg and isinstance(msg, Mapping):
            if msg["type"] == "lifespan.startup.failed":
                await aio_cancel(task)
            else:
                assert msg["type"] == "lifespan.startup.complete"

        yield

        await pipe.send_to_app({"type": "lifespan.shutdown"})
        msg = await aio_wait(
            pipe.receive_from_client(), aio_sleep(timeout), strategy=FIRST_COMPLETED
        )
        if msg and isinstance(msg, Mapping):
            assert msg["type"] == "lifespan.shutdown.complete"


# pylama:ignore=D
