from __future__ import annotations

import asyncio
import base64
from collections.abc import Awaitable, Callable

import pytest

from app.fleet.proxy import ProxyEndpoint, parse_proxy_url, probe_proxy

ServerHandler = Callable[[asyncio.StreamReader, asyncio.StreamWriter], Awaitable[None]]


@pytest.mark.parametrize(
    ("value", "expected"),
    [
        (
            "socks5://proxy-user:proxy-password@proxy.example.com:10808",
            ProxyEndpoint(
                protocol="socks5",
                host="proxy.example.com",
                port=10808,
                username="proxy-user",
                password="proxy-password",
            ),
        ),
        (
            "http://proxy-user:proxy-password@proxy.example.com:10808",
            ProxyEndpoint(
                protocol="http",
                host="proxy.example.com",
                port=10808,
                username="proxy-user",
                password="proxy-password",
            ),
        ),
        (
            "socks5://proxy.example.com:10808@proxy-user:proxy-password",
            ProxyEndpoint(
                protocol="socks5",
                host="proxy.example.com",
                port=10808,
                username="proxy-user",
                password="proxy-password",
            ),
        ),
        (
            "socks5://proxy.example.com:1080@user:1234",
            ProxyEndpoint(
                protocol="socks5",
                host="proxy.example.com",
                port=1080,
                username="user",
                password="1234",
            ),
        ),
        (
            "socks5://1.2.3.4:1234@host.example.com:5678",
            ProxyEndpoint(
                protocol="socks5",
                host="1.2.3.4",
                port=1234,
                username="host.example.com",
                password="5678",
            ),
        ),
        (
            "SOCKS5://user%40team:p%3Aa%2Fss@PROXY.EXAMPLE.COM:1080",
            ProxyEndpoint(
                protocol="socks5",
                host="proxy.example.com",
                port=1080,
                username="user@team",
                password="p:a/ss",
            ),
        ),
        (
            "http://proxy.example.com:8080",
            ProxyEndpoint(
                protocol="http",
                host="proxy.example.com",
                port=8080,
                username=None,
                password=None,
            ),
        ),
    ],
)
def test_parse_proxy_url_supported_formats(value: str, expected: ProxyEndpoint) -> None:
    assert parse_proxy_url(value) == expected


@pytest.mark.parametrize(
    "value",
    [
        "",
        "ftp://user:password@proxy.example.com:1080",
        "socks5://user:password@proxy.example.com",
        "socks5://user:password@proxy.example.com:0",
        "socks5://user:password@proxy.example.com:65536",
        "socks5://user:password@proxy.example.com:not-a-port",
        "socks5://user:password@proxy.example.com:1080/path",
        "socks5://user:password@proxy.example.com:1080?mode=test",
        "socks5://user:password@proxy.example.com:1080#fragment",
        "socks5://user:password@proxy.example.com:1080\n",
        "socks5://user:pass%0Dword@proxy.example.com:1080",
        "socks5://user:pass%ZZword@proxy.example.com:1080",
        "socks5://user@proxy.example.com:1080",
        "socks5://user:@proxy.example.com:1080",
        "socks5://user:password@proxy.example.com:1080@extra",
    ],
)
def test_parse_proxy_url_rejects_invalid_values(value: str) -> None:
    with pytest.raises(ValueError):
        parse_proxy_url(value)


def test_parse_errors_do_not_echo_credentials() -> None:
    username = "sensitive-user"
    password = "sensitive-password"
    with pytest.raises(ValueError) as captured:
        parse_proxy_url(f"socks5://{username}:{password}@proxy.example.com:70000")
    message = str(captured.value)
    assert username not in message
    assert password not in message


@pytest.mark.asyncio
async def test_probe_socks5_performs_username_password_auth_and_connect() -> None:
    observed: asyncio.Future[tuple[str, str, bytes]] = asyncio.get_running_loop().create_future()

    async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            greeting = await reader.readexactly(3)
            writer.write(b"\x05\x02")
            await writer.drain()
            auth_version, username_length = await reader.readexactly(2)
            username = (await reader.readexactly(username_length)).decode()
            password_length = (await reader.readexactly(1))[0]
            password = (await reader.readexactly(password_length)).decode()
            writer.write(b"\x01\x00")
            await writer.drain()
            connect_request = await reader.readexactly(10)
            writer.write(b"\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00")
            await writer.drain()
            observed.set_result(
                (username, password, greeting + bytes((auth_version,)) + connect_request)
            )
        except Exception as exc:  # pragma: no cover - surfaced through the future
            if not observed.done():
                observed.set_exception(exc)
        finally:
            writer.close()
            await writer.wait_closed()

    server, port = await _start_server(handler)
    try:
        await probe_proxy(
            "127.0.0.1",
            port,
            "socks5",
            username="proxy-user",
            password="proxy-password",
        )
        username, password, protocol_bytes = await observed
    finally:
        await _close_server(server)

    assert username == "proxy-user"
    assert password == "proxy-password"
    assert protocol_bytes[:3] == b"\x05\x01\x02"
    assert protocol_bytes[3] == 1
    assert protocol_bytes[4:] == b"\x05\x01\x00\x01\x01\x01\x01\x01\x01\xbb"


@pytest.mark.asyncio
async def test_probe_socks5_supports_anonymous_method() -> None:
    observed: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()

    async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            greeting = await reader.readexactly(3)
            writer.write(b"\x05\x00")
            await writer.drain()
            connect_request = await reader.readexactly(10)
            writer.write(b"\x05\x00\x00\x03\x03foo\x00\x50")
            await writer.drain()
            observed.set_result(greeting + connect_request)
        finally:
            writer.close()
            await writer.wait_closed()

    server, port = await _start_server(handler)
    try:
        await probe_proxy("127.0.0.1", port, "socks5")
        protocol_bytes = await observed
    finally:
        await _close_server(server)

    assert protocol_bytes[:3] == b"\x05\x01\x00"
    assert protocol_bytes[3:] == b"\x05\x01\x00\x01\x01\x01\x01\x01\x01\xbb"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("method_reply", "auth_reply", "connect_reply"),
    [
        (b"\x05\x00", None, None),
        (b"\x05\x02", b"\x01\x01", None),
        (
            b"\x05\x02",
            b"\x01\x00",
            b"\x05\x05\x00\x01\x00\x00\x00\x00\x00\x00",
        ),
    ],
)
async def test_probe_socks5_rejects_protocol_and_auth_failures_without_secret_leak(
    method_reply: bytes,
    auth_reply: bytes | None,
    connect_reply: bytes | None,
) -> None:
    async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            await reader.readexactly(3)
            writer.write(method_reply)
            await writer.drain()
            if auth_reply is not None:
                _version, username_length = await reader.readexactly(2)
                await reader.readexactly(username_length)
                password_length = (await reader.readexactly(1))[0]
                await reader.readexactly(password_length)
                writer.write(auth_reply)
                await writer.drain()
            if connect_reply is not None:
                await reader.readexactly(10)
                writer.write(connect_reply)
                await writer.drain()
        finally:
            writer.close()
            await writer.wait_closed()

    server, port = await _start_server(handler)
    try:
        with pytest.raises(OSError) as captured:
            await probe_proxy(
                "127.0.0.1",
                port,
                "socks5",
                username="secret-user",
                password="secret-password",
            )
    finally:
        await _close_server(server)

    assert "secret-user" not in str(captured.value)
    assert "secret-password" not in str(captured.value)


@pytest.mark.asyncio
async def test_probe_http_sends_connect_and_basic_authorization() -> None:
    observed: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()

    async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            request = await reader.readuntil(b"\r\n\r\n")
            observed.set_result(request)
            writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
            await writer.drain()
        finally:
            writer.close()
            await writer.wait_closed()

    server, port = await _start_server(handler)
    try:
        await probe_proxy(
            "127.0.0.1",
            port,
            "http",
            username="http-user",
            password="http-password",
        )
        request = (await observed).decode("ascii")
    finally:
        await _close_server(server)

    expected_auth = base64.b64encode(b"http-user:http-password").decode("ascii")
    assert request.startswith("CONNECT 1.1.1.1:443 HTTP/1.1\r\n")
    assert "Host: 1.1.1.1:443\r\n" in request
    assert f"Proxy-Authorization: Basic {expected_auth}\r\n" in request


@pytest.mark.asyncio
@pytest.mark.parametrize("status", [407, 500])
async def test_probe_http_rejects_non_2xx_without_secret_leak(status: int) -> None:
    async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            await reader.readuntil(b"\r\n\r\n")
            writer.write(f"HTTP/1.1 {status} Failure\r\n\r\n".encode("ascii"))
            await writer.drain()
        finally:
            writer.close()
            await writer.wait_closed()

    server, port = await _start_server(handler)
    try:
        with pytest.raises(OSError) as captured:
            await probe_proxy(
                "127.0.0.1",
                port,
                "http",
                username="secret-http-user",
                password="secret-http-password",
            )
    finally:
        await _close_server(server)

    assert "secret-http-user" not in str(captured.value)
    assert "secret-http-password" not in str(captured.value)


@pytest.mark.asyncio
async def test_probe_timeout_is_sanitized() -> None:
    async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            await reader.read()
        finally:
            writer.close()
            await writer.wait_closed()

    server, port = await _start_server(handler)
    try:
        with pytest.raises(OSError, match="代理探测超时"):
            await probe_proxy("127.0.0.1", port, "socks5", timeout_seconds=0.01)
    finally:
        await _close_server(server)


@pytest.mark.asyncio
async def test_probe_rejects_invalid_configuration() -> None:
    with pytest.raises(ValueError):
        await probe_proxy("127.0.0.1", 1080, "https")
    with pytest.raises(ValueError):
        await probe_proxy("127.0.0.1", 0, "socks5")
    with pytest.raises(ValueError):
        await probe_proxy(
            "127.0.0.1",
            1080,
            "socks5",
            username="only-username",
        )


async def _start_server(handler: ServerHandler) -> tuple[asyncio.AbstractServer, int]:
    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    socket_address = server.sockets[0].getsockname()
    return server, int(socket_address[1])


async def _close_server(server: asyncio.AbstractServer) -> None:
    server.close()
    await server.wait_closed()