"""Synthetic HTTP/1.1 pool exhaustion with real loopback TCP origin and proxy.

The intentionally small fixture handles one GET per connection. It cannot
forward to any destination except its own ephemeral loopback origin.
"""

import asyncio
from datetime import datetime, timezone
from importlib.metadata import version
import json
import platform
import sys
from time import monotonic
from urllib.parse import urlsplit

import httpx
from async_proxy_check import make_client

BODY = b'{"ok":true,"source":"loopback-origin"}\n'
FIXTURE_DEADLINE = 8.0


def require(condition: bool, message: str) -> None:
    if not condition:
        raise AssertionError(message)


class LoopbackFixture:
    """Owned sockets/tasks are closed and awaited, including exceptional exits."""

    def __init__(self):
        self.events = []
        self.requests = {"proxy": [], "origin": []}
        self.errors = []
        self.tasks = set()
        self.writers = set()
        self.servers = []
        self.forced_cancellations = 0
        self.ready = {path: asyncio.Event() for path in ("/hold/1", "/hold/2")}
        self.closed = {path: asyncio.Event() for path in self.ready}
        self.started = monotonic()

    def log(self, side, event, **fields):
        self.events.append({"sequence": len(self.events) + 1,
                            "seconds": round(monotonic() - self.started, 6),
                            "side": side, "event": event, **fields})

    async def __aenter__(self):
        try:
            origin = await asyncio.start_server(self.accept("origin"), "127.0.0.1", 0)
            self.servers.append(origin)
            self.origin_port = origin.sockets[0].getsockname()[1]
            proxy = await asyncio.start_server(self.accept("proxy"), "127.0.0.1", 0)
            self.servers.append(proxy)
            self.proxy_port = proxy.sockets[0].getsockname()[1]
            self.origin_url = f"http://127.0.0.1:{self.origin_port}"
            self.proxy_url = f"http://127.0.0.1:{self.proxy_port}"
            return self
        except BaseException:
            await self.__aexit__(None, None, None)
            raise

    def accept(self, side):
        def completed(task):
            self.tasks.discard(task)
            if not task.cancelled() and task.exception() is not None:
                self.errors.append({"side": side, "error": type(task.exception()).__name__})

        def connected(reader, writer):
            self.writers.add(writer)
            task = asyncio.create_task(self.serve(side, reader, writer))
            self.tasks.add(task)
            task.add_done_callback(completed)
        return connected

    async def close_writer(self, writer):
        writer.close()
        try:
            await asyncio.wait_for(writer.wait_closed(), 1.0)
        except (ConnectionError, BrokenPipeError):
            pass
        finally:
            self.writers.discard(writer)

    async def serve(self, side, reader, writer):
        try:
            async with asyncio.timeout(FIXTURE_DEADLINE):
                await getattr(self, side)(reader, writer)
        except (ConnectionError, asyncio.IncompleteReadError):
            # Early client closure is expected when abandoning a streaming body.
            self.log(side, "peer_closed")
        except Exception as error:
            self.errors.append({"side": side, "error": type(error).__name__})
        finally:
            await self.close_writer(writer)

    async def origin(self, reader, writer):
        header = await reader.readuntil(b"\r\n\r\n")
        method, path, protocol = header.split(b"\r\n", 1)[0].decode("ascii").split()
        require(method == "GET" and protocol == "HTTP/1.1", "fixture accepts HTTP/1.1 GET only")
        self.requests["origin"].append(path)
        self.log("origin", "request", path=path)
        if path in self.ready:
            # The body remains incomplete until the client closes the stream.
            writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 20\r\nConnection: close\r\n\r\nheld")
            await writer.drain()
            self.log("origin", "body_held", path=path)
            self.ready[path].set()
            require(await reader.read(1) == b"", "unexpected bytes on held request")
            self.log("origin", "held_peer_closed", path=path)
            self.closed[path].set()
            return
        body = BODY if path != "/large" else b"x" * 65_537
        status = b"503 Service Unavailable" if path == "/status" else b"200 OK"
        writer.write(b"HTTP/1.1 " + status + b"\r\nContent-Length: "
                     + str(len(body)).encode() + b"\r\nContent-Type: application/json"
                     + b"\r\nConnection: close\r\n\r\n" + body)
        await writer.drain()

    async def proxy(self, reader, writer):
        header = await reader.readuntil(b"\r\n\r\n")
        method, target, protocol = header.split(b"\r\n", 1)[0].decode("ascii").split()
        url = urlsplit(target)
        require(method == "GET" and protocol == "HTTP/1.1", "fixture accepts GET only")
        require(url.scheme == "http" and url.hostname == "127.0.0.1"
                and url.port == self.origin_port and not url.username and not url.password,
                "fixture refuses destinations outside its loopback origin")
        require(not url.query and not url.fragment, "fixture paths have no query or fragment")
        path = url.path or "/"
        self.requests["proxy"].append(path)
        self.log("proxy", "request", path=path)
        upstream_reader, upstream_writer = await asyncio.open_connection("127.0.0.1", self.origin_port)
        self.writers.add(upstream_writer)
        pumps = []
        try:
            upstream_writer.write(f"GET {path} HTTP/1.1\r\nHost: 127.0.0.1:{self.origin_port}\r\nConnection: close\r\n\r\n".encode("ascii"))
            await upstream_writer.drain()

            async def copy(source, destination):
                while data := await source.read(8192):
                    destination.write(data)
                    await destination.drain()

            pumps = [asyncio.create_task(copy(upstream_reader, writer)),
                     asyncio.create_task(copy(reader, upstream_writer))]
            done, _ = await asyncio.wait(pumps, return_when=asyncio.FIRST_COMPLETED)
            for task in done:
                task.result()
        finally:
            for task in pumps:
                task.cancel()
            await asyncio.gather(*pumps, return_exceptions=True)
            await self.close_writer(upstream_writer)

    async def __aexit__(self, *_):
        for server in self.servers:
            server.close()
        for server in self.servers:
            await server.wait_closed()
        # Closing all peers lets normal handlers finish; cancellation is a fallback.
        for writer in tuple(self.writers):
            writer.close()
        if self.tasks:
            _, pending = await asyncio.wait(tuple(self.tasks), timeout=1.0)
            self.forced_cancellations += len(pending)
            for task in pending:
                task.cancel()
            await asyncio.gather(*pending, return_exceptions=True)
        self.log("fixture", "closed", handlers=len(self.tasks), writers=len(self.writers))

    def assert_clean(self):
        require(not self.tasks and not self.writers, "fixture leaked handlers or writers")
        require(not any(server.is_serving() for server in self.servers), "listener still serving")
        require(self.forced_cancellations == 0, "fixture required forced handler cancellation")
        require(not self.errors, f"fixture failed: {self.errors}")


async def run_diagnostic():
    baseline_tasks = set(asyncio.all_tasks())
    observed_at = datetime.now(timezone.utc).isoformat()
    fixture = LoopbackFixture()
    async with fixture:
        async with make_client(fixture.proxy_url) as client:
            held = []
            try:
                for path in ("/hold/1", "/hold/2"):
                    fixture.log("client", "attempt", path=path)
                    response = await client.send(client.build_request("GET", fixture.origin_url + path), stream=True)
                    held.append(response)
                    response.raise_for_status()
                    await asyncio.wait_for(fixture.ready[path].wait(), 1.0)
                    require(not response.is_closed, "held response closed unexpectedly")
                    fixture.log("client", "stream_open", path=path)
                before = {side: list(paths) for side, paths in fixture.requests.items()}
                fixture.log("client", "attempt", path="/blocked")
                start = monotonic()
                try:
                    await client.get(fixture.origin_url + "/blocked")
                except httpx.PoolTimeout:
                    elapsed = monotonic() - start
                    fixture.log("client", "PoolTimeout", path="/blocked")
                else:
                    raise AssertionError("third request did not raise PoolTimeout")
                require(before == fixture.requests, "blocked request reached a server")
                await held[0].aclose()
                fixture.log("client", "explicit_aclose", path="/hold/1")
                await asyncio.wait_for(fixture.closed["/hold/1"].wait(), 1.0)
                fixture.log("client", "attempt", path="/ok")
                recovered = await client.get(fixture.origin_url + "/ok")
                recovered.raise_for_status()
                require(recovered.content == BODY, "recovery body failed validation")
                require(recovered.http_version == "HTTP/1.1", "unexpected HTTP version")
                fixture.log("client", "validated_body", path="/ok", status=recovered.status_code)
                require(not held[1].is_closed, "second stream should still be occupied")
                expected = ["/hold/1", "/hold/2", "/ok"]
                require(all(paths == expected for paths in fixture.requests.values()), "unexpected server sequence")
            finally:
                for response in held:
                    await response.aclose()
    fixture.assert_clean()
    require(client.is_closed, "HTTPX client still open")
    require(not set(asyncio.all_tasks()) - baseline_tasks, "diagnostic left new asyncio tasks")
    return {"outcome": "pass", "observation": "synthetic-loopback-http1",
            "observed_at_utc": observed_at,
            "runtime": {"python": platform.python_version(), "implementation": platform.python_implementation(),
                        "system": platform.system(), "machine": platform.machine(),
                        "packages": {name: version(name) for name in ("httpx", "httpcore", "anyio", "h11", "certifi", "idna", "typing_extensions")}},
            "config": {"max_connections": 2, "held_streams": 2, "pool_timeout_seconds": 0.25,
                       "trust_env": False, "http2": False},
            "pool_wait_observed_seconds": round(elapsed, 6),
            "blocked_request_reached_proxy": False, "blocked_request_reached_origin": False,
            "recovery": {"status": recovered.status_code, "body": recovered.json(), "body_validated": True},
            "requests": fixture.requests,
            "cleanup": {"handlers": len(fixture.tasks), "writers": len(fixture.writers), "errors": fixture.errors,
                        "forced_cancellations": fixture.forced_cancellations, "client_closed": client.is_closed},
            "events": fixture.events}


async def bounded_diagnostic():
    async with asyncio.timeout(12.0):
        return await run_diagnostic()


def main():
    try:
        result = asyncio.run(bounded_diagnostic())
    except Exception as error:
        print(json.dumps({"outcome": "fail", "error": type(error).__name__}))
        return 1
    print(json.dumps(result, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())