import asyncio

from app.models.health import HealthSnapshot
from app.models.proxy import ProxyDocument
from app.models.service import ServiceEndpointDocument
from app.services.mitmproxy.health import check_proxy


def test_check_proxy_reports_upstream_and_proxy_healthy():
    async def run():
        upstream = await asyncio.start_server(http_handler, "127.0.0.1", 0)
        proxy = await asyncio.start_server(http_handler, "127.0.0.1", 0)
        try:
            upstream_port = upstream.sockets[0].getsockname()[1]
            proxy_port = proxy.sockets[0].getsockname()[1]
            document = ProxyDocument(
                proxy_port=proxy_port,
                name="health",
                code="",
                service=ServiceEndpointDocument(port=upstream_port),
            )

            health = await check_proxy(document)
        finally:
            upstream.close()
            proxy.close()
            await upstream.wait_closed()
            await proxy.wait_closed()

        assert isinstance(health, HealthSnapshot)
        assert health.overall == "healthy"
        assert [check.target for check in health.checks] == [
            "upstream",
            "proxy",
        ]
        assert all(check.duration_ms is None for check in health.checks)

    asyncio.run(run())


def test_check_proxy_reports_benchmark_latency_when_enabled():
    async def run():
        upstream = await asyncio.start_server(http_handler, "127.0.0.1", 0)
        proxy = await asyncio.start_server(http_handler, "127.0.0.1", 0)
        try:
            upstream_port = upstream.sockets[0].getsockname()[1]
            proxy_port = proxy.sockets[0].getsockname()[1]
            document = ProxyDocument(
                proxy_port=proxy_port,
                name="health",
                code="",
                healthcheck={"benchmark_enabled": True},
                service=ServiceEndpointDocument(port=upstream_port),
            )

            health = await check_proxy(document)
        finally:
            upstream.close()
            proxy.close()
            await upstream.wait_closed()
            await proxy.wait_closed()

        assert isinstance(health, HealthSnapshot)
        assert all(
            check.duration_ms is not None and check.duration_ms >= 0
            for check in health.checks
        )

    asyncio.run(run())


async def http_handler(reader, writer):
    await reader.readline()
    writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")
    await writer.drain()
    writer.close()
    await writer.wait_closed()