from __future__ import annotations

import asyncio
from typing import Any

from packetstation.proxy import TrustedProxies, TrustedProxyMiddleware, forwarded

PROXY = "172.18.0.5"


def trusted(addresses: set[str] | None = None) -> TrustedProxies:
    t = TrustedProxies("npm_app_1", resolver=lambda _h: set(addresses or {PROXY}))
    t.refresh()
    return t


def H(**kw: str) -> list[tuple[bytes, bytes]]:
    return [(k.replace("_", "-").encode(), v.encode()) for k, v in kw.items()]


def test_trusted_peer_uses_first_xff_and_proto() -> None:
    ip, scheme = forwarded(PROXY, H(x_forwarded_for="10.1.2.3, 172.18.0.9", x_forwarded_proto="https"), trusted())
    assert ip == "10.1.2.3" and scheme == "https"


def test_untrusted_peer_ignores_headers() -> None:
    ip, scheme = forwarded("10.9.9.9", H(x_forwarded_for="1.2.3.4", x_forwarded_proto="https"), trusted())
    assert ip == "10.9.9.9" and scheme is None


def test_empty_host_trusts_nothing() -> None:
    t = TrustedProxies("")
    t.refresh()
    assert forwarded(PROXY, H(x_forwarded_for="1.2.3.4"), t) == (PROXY, None)


def test_garbage_xff_and_proto_ignored() -> None:
    ip, scheme = forwarded(PROXY, H(x_forwarded_for="not-an-ip", x_forwarded_proto="gopher"), trusted())
    assert ip == PROXY and scheme is None


def test_resolution_failure_keeps_last_good_set() -> None:
    calls = {"n": 0}

    def resolver(_h: str) -> set[str]:
        calls["n"] += 1
        if calls["n"] > 1:
            raise OSError("no such host")
        return {PROXY}

    t = TrustedProxies("npm_app_1", resolver=resolver)
    t.refresh()
    t.refresh()
    assert t.addresses == frozenset({PROXY})


def test_re_resolve_picks_up_new_address() -> None:
    addrs = [{PROXY}, {"172.18.0.77"}]
    t = TrustedProxies("npm_app_1", resolver=lambda _h: addrs.pop(0))
    t.refresh()
    t.refresh()
    assert t.trusts("172.18.0.77") and not t.trusts(PROXY)


def _run_mw(peer: str, headers: list[tuple[bytes, bytes]], t: TrustedProxies) -> dict[str, Any]:
    seen: dict[str, Any] = {}

    async def app(scope: Any, receive: Any, send: Any) -> None:
        seen.update(scope)

    scope = {"type": "http", "client": (peer, 5555), "scheme": "http", "headers": headers, "path": "/"}
    asyncio.run(TrustedProxyMiddleware(app, t)(scope, None, None))  # type: ignore[arg-type]
    return seen


def test_middleware_rewrites_for_trusted_peer() -> None:
    s = _run_mw(PROXY, H(x_forwarded_for="10.1.2.3", x_forwarded_proto="https"), trusted())
    assert s["client"] == ("10.1.2.3", 5555) and s["scheme"] == "https"


def test_middleware_leaves_untrusted_peer() -> None:
    s = _run_mw("10.5.5.5", H(x_forwarded_for="10.1.2.3", x_forwarded_proto="https"), trusted())
    assert s["client"] == ("10.5.5.5", 5555) and s["scheme"] == "http"