from __future__ import annotations

import asyncio
import contextlib
import contextvars
import json
import logging
import signal
import sys
from collections.abc import Callable, Generator
from contextlib import AbstractContextManager

import httpx
import pytest

from tests.protocols.test_http import SIMPLE_GET_REQUEST
from tests.utils import run_server
from uvicorn._types import ASGIApplication, ASGIReceiveCallable, ASGISendCallable, Scope
from uvicorn.config import Config
from uvicorn.protocols.http.flow_control import HIGH_WATER_LIMIT
from uvicorn.protocols.http.h11_impl import H11Protocol
from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol
from uvicorn.server import Server

pytestmark = pytest.mark.anyio


# asyncio does NOT allow raising in signal handlers, so to detect
# raised signals raised a mutable `witness` receives the signal
@contextlib.contextmanager
def capture_signal_sync(sig: signal.Signals) -> Generator[list[int], None, None]:
    """Replace `sig` handling with a normal exception via `signal"""
    witness: list[int] = []
    original_handler = signal.signal(sig, lambda signum, frame: witness.append(signum))
    yield witness
    signal.signal(sig, original_handler)


@contextlib.contextmanager
def capture_signal_async(sig: signal.Signals) -> Generator[list[int], None, None]:  # pragma: py-win32
    """Replace `sig` handling with a normal exception via `asyncio"""
    witness: list[int] = []
    original_handler = signal.getsignal(sig)
    asyncio.get_running_loop().add_signal_handler(sig, witness.append, sig)
    yield witness
    signal.signal(sig, original_handler)


async def dummy_app(scope, receive, send):  # pragma: py-win32
    pass


async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
    assert scope["type"] == "http"
    await send({"type": "http.response.start", "status": 200, "headers": []})
    await send({"type": "http.response.body", "body": b"", "more_body": False})


if sys.platform == "win32":  # pragma: py-not-win32
    signals = [signal.SIGBREAK]
    signal_captures = [capture_signal_sync]
else:  # pragma: py-win32
    signals = [signal.SIGTERM, signal.SIGINT]
    signal_captures = [capture_signal_sync, capture_signal_async]


@pytest.mark.parametrize("exception_signal", signals)
@pytest.mark.parametrize("capture_signal", signal_captures)
async def test_server_interrupt(
    exception_signal: signal.Signals,
    capture_signal: Callable[[signal.Signals], AbstractContextManager[None]],
    unused_tcp_port: int,
):  # pragma: py-win32
    """Test interrupting a Server that is run explicitly inside asyncio"""

    async def interrupt_running(srv: Server):
        while not srv.started:
            await asyncio.sleep(0.01)
        signal.raise_signal(exception_signal)

    server = Server(Config(app=dummy_app, loop="asyncio", port=unused_tcp_port))
    asyncio.create_task(interrupt_running(server))
    with capture_signal(exception_signal) as witness:
        await server.serve()
    assert witness
    # set by the server's graceful exit handler
    assert server.should_exit


async def test_shutdown_on_early_exit_during_startup(unused_tcp_port: int):
    """Test that lifespan.shutdown is called even when should_exit is set during startup."""
    startup_complete = False
    shutdown_complete = False

    async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
        nonlocal startup_complete, shutdown_complete
        if scope["type"] == "lifespan":
            while True:
                message = await receive()
                if message["type"] == "lifespan.startup":
                    await asyncio.sleep(0.5)
                    await send({"type": "lifespan.startup.complete"})
                    startup_complete = True
                elif message["type"] == "lifespan.shutdown":
                    await send({"type": "lifespan.shutdown.complete"})
                    shutdown_complete = True
                    return

    config = Config(app=app, lifespan="on", port=unused_tcp_port)
    server = Server(config=config)

    # Simulate a reload signal arriving during startup:
    # set should_exit before the 0.5s startup sleep finishes.
    async def set_exit():
        await asyncio.sleep(0.2)
        server.should_exit = True

    asyncio.create_task(set_exit())
    await server.serve()

    assert startup_complete
    assert shutdown_complete, "lifespan.shutdown was not called despite startup completing"


async def test_request_than_limit_max_requests_warn_log(
    unused_tcp_port: int, http_protocol_cls: type[H11Protocol | HttpToolsProtocol], caplog: pytest.LogCaptureFixture
):
    caplog.set_level(logging.INFO, logger="uvicorn.error")
    config = Config(app=app, limit_max_requests=1, port=unused_tcp_port, http=http_protocol_cls)
    async with run_server(config):
        async with httpx.AsyncClient() as client:
            tasks = [client.get(f"http://127.0.0.1:{unused_tcp_port}") for _ in range(2)]
            responses = await asyncio.gather(*tasks)
            assert len(responses) == 2
    assert "Maximum request limit of 1 exceeded. Terminating process." in caplog.text


async def test_limit_max_requests_jitter(
    unused_tcp_port: int, http_protocol_cls: type[H11Protocol | HttpToolsProtocol], caplog: pytest.LogCaptureFixture
):
    caplog.set_level(logging.INFO, logger="uvicorn.error")
    config = Config(
        app=app, limit_max_requests=1, limit_max_requests_jitter=2, port=unused_tcp_port, http=http_protocol_cls
    )
    async with run_server(config) as server:
        limit = server.limit_max_requests
        assert limit is not None
        assert 1 <= limit <= 3
        async with httpx.AsyncClient() as client:
            tasks = [client.get(f"http://127.0.0.1:{unused_tcp_port}") for _ in range(limit + 1)]
            await asyncio.gather(*tasks)
    assert f"Maximum request limit of {limit} exceeded. Terminating process." in caplog.text


@contextlib.asynccontextmanager
async def _raw_server(
    *,
    app: ASGIApplication,
    port: int,
    http_protocol_cls: type[H11Protocol | HttpToolsProtocol],
    reset_contextvars: bool = False,
):
    config = Config(app=app, port=port, loop="asyncio", http=http_protocol_cls, reset_contextvars=reset_contextvars)
    server = Server(config=config)
    task = asyncio.create_task(server.serve())

    while not server.started:
        await asyncio.sleep(0.01)

    reader, writer = await asyncio.open_connection("127.0.0.1", port)

    async def extract_json_body(request: bytes):
        writer.write(request)
        await writer.drain()

        status, *headers = (await reader.readuntil(b"\r\n\r\n")).split(b"\r\n")[:-2]
        assert status == b"HTTP/1.1 200 OK"

        content_length = next(int(h.split(b":", 1)[1]) for h in headers if h.lower().startswith(b"content-length:"))
        return json.loads(await reader.readexactly(content_length))

    try:
        yield extract_json_body
    finally:
        writer.close()
        await writer.wait_closed()
        server.should_exit = True
        await task


async def test_contextvars_preserved_by_default(
    http_protocol_cls: type[H11Protocol | HttpToolsProtocol], unused_tcp_port: int
):
    """By default, context set outside the ASGI task is visible inside it."""
    ctx: contextvars.ContextVar[str] = contextvars.ContextVar("ctx")
    ctx.set("outer-value")

    async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
        assert scope["type"] == "http"
        while True:
            message = await receive()
            assert message["type"] == "http.request"
            if not message["more_body"]:
                break
        body = json.dumps({"ctx": ctx.get("MISSING")}).encode("utf-8")
        headers = [(b"content-type", b"application/json"), (b"content-length", str(len(body)).encode("utf-8"))]
        await send({"type": "http.response.start", "status": 200, "headers": headers})
        await send({"type": "http.response.body", "body": body})

    async with _raw_server(app=app, http_protocol_cls=http_protocol_cls, port=unused_tcp_port) as extract_json_body:
        assert await extract_json_body(SIMPLE_GET_REQUEST) == {"ctx": "outer-value"}


async def test_reset_contextvars_asyncio(
    http_protocol_cls: type[H11Protocol | HttpToolsProtocol], unused_tcp_port: int
):
    """With reset_contextvars=True, each ASGI run starts with a fresh context.

    Non-regression test for https://github.com/encode/uvicorn/issues/2167.
    """
    default_contextvars = {c.name for c in contextvars.copy_context().keys()}
    ctx: contextvars.ContextVar[str] = contextvars.ContextVar("ctx")

    async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
        assert scope["type"] == "http"

        # initial context should be empty
        initial_context = {
            n: v for c, v in contextvars.copy_context().items() if (n := c.name) not in default_contextvars
        }
        # set any contextvar before the body is read
        ctx.set(scope["path"])

        while True:
            message = await receive()
            assert message["type"] == "http.request"
            if not message["more_body"]:
                break

        body = json.dumps(initial_context).encode("utf-8")
        headers = [(b"content-type", b"application/json"), (b"content-length", str(len(body)).encode("utf-8"))]
        await send({"type": "http.response.start", "status": 200, "headers": headers})
        await send({"type": "http.response.body", "body": body})

    # body larger than HIGH_WATER_LIMIT forces a reading pause on the main thread
    # and a resumption inside the ASGI task, which is where the original pollution showed up.
    large_body = b"a" * (HIGH_WATER_LIMIT + 1)
    large_request = b"\r\n".join(
        [
            b"POST /large-body HTTP/1.1",
            b"Host: example.org",
            b"Content-Type: application/octet-stream",
            f"Content-Length: {len(large_body)}".encode(),
            b"",
            large_body,
        ]
    )

    async with _raw_server(
        app=app, http_protocol_cls=http_protocol_cls, port=unused_tcp_port, reset_contextvars=True
    ) as extract_json_body:
        assert await extract_json_body(large_request) == {}
        assert await extract_json_body(SIMPLE_GET_REQUEST) == {}
