"""
When using UDP in network programming its common to write
code so that it has simple loops to retry sending packets if
no response is received. The problem is: if any packets are lost
the time spent waiting keeps accumulating to the cost of a round-trip.

The problem is made worse when you consider slow-downs in DNS and
the often unreliable nature of relying on community servers for
STUN code. If one is not careful they quickly end up with an
algorithm that is prohibitively slow in the best case and quite
unreliable in the worst case.

The original algorithm for doing STUN tests for RFC 3489 used
incremental tests and was very brittle. However -- due to the
nature of how these tests work -- it became clear to me that it was
possible to paralyze the tests into two main phases if a STUN
server's primary and secondary IP were known beforehand.

Phase 1 tests for [open NAT and full cone] while phase 2 tests for
[restrict ip and restrict port] behaviors. The algorithm here is designed to run across multiple public servers where it creates races among
all servers to test a routes NAT. The result is you end up with
the fastest possible determination of NAT behaviors while also
building in safe-guards against packet-loss, inconsistent results, misconfigurations, and slow network conditions.
"""

import asyncio
from ...utility.utils import (
    async_wrap_errors,
    ip_f,
    log,
    fstr,
)
from ...errors import ErrorCantLoadNATInfo
from ...net.net_defs import IP4, IP6, UDP
from ...protocol.stun.stun_defs import (
    RFC3489,
    RFC5389,
    STUN_CHANGE_NONE,
    STUN_CHANGE_PORT,
    STUN_CHANGE_BOTH,
)
from ...protocol.stun.stun_client import STUNClient, get_stun_clients
from ...protocol.stun.stun_utils import stun_reply_to_ret_dic
from ...servers import get_infra
from ...net.pipe.pipe import Pipe
from .nat_defs import (
    BLOCKED_NAT,
    FULL_CONE,
    NA_DELTA,
    OPEN_INTERNET,
    RANDOM_DELTA,
    RESTRICT_NAT,
    RESTRICT_PORT_NAT,
    SYMMETRIC_NAT,
)
from .nat_utils import delta_info, delta_test, nat_info


# Constants for a NAT test.
NAT_TEST_NO = 5
# Per-sub-test timeout. Original value 0.5s was too aggressive for mobile/
# congested links: when q1 or q3 didn't return inside that window, the
# classifier silently fell through to its SYMMETRIC default -- producing
# spurious SYMMETRIC results for NATs that are actually port-restricted.
# 1.5s gives slow STUN servers room to answer while still leaving headroom
# under the outer 4s nic_load_nat budget (worst case 2 * 1.5s = 3s).
NAT_TEST_TIMEOUT = 1.5

# STUN payload to send, Send IP, send port, reply IP, reply port
# Shows order of RFC 3489 NAT enumeration sets.
NAT_TEST_SCHEMA = [
    # Detects: open NAT.
    # dest: primary, primary. reply: primary, primary
    [STUN_CHANGE_NONE, 0, 0, 0, 0],
    # Detects: full cone NAT.
    # Change both reply IP and port.
    # dest: primary, primary. reply: secondary, secondary
    [STUN_CHANGE_BOTH, 0, 0, 3, 3],
    # Detects: non-symmetric NAT.
    # dest secondary, primary. reply: secondary, primary
    [STUN_CHANGE_NONE, 2, 2, 2, 0],
    # Detects: between restrict and port restrict.
    # Change only the reply port.
    # dest: secondary, primary. reply: secondary, secondary.
    [STUN_CHANGE_PORT, 2, 2, 3, 3],
]


async def nat_test_exec(
    dest_addr,
    reply_addr,
    payload,
    mode,
    pipe,
    q,
    test_coro,
):
    stun_client = STUNClient(pipe.route.af, dest_addr, pipe.route.interface, mode=mode)
    if payload == STUN_CHANGE_NONE:
        reply = await stun_client.get_stun_reply(pipe)
    if payload == STUN_CHANGE_PORT:
        reply = await stun_client.get_change_port_reply(reply_addr, pipe)
    if payload == STUN_CHANGE_BOTH:
        reply = await stun_client.get_change_tup_reply(reply_addr, pipe)

    ret = await stun_reply_to_ret_dic(reply)

    # Valid reply.
    if ret is not None and not isinstance(ret, tuple):
        q.append(ret)
        return await test_coro(ret, pipe)

    return None


async def nat_test_workers(
    pipe,
    q,
    test_index,
    test_coro,
    servers,
    test_no,
):
    # Make list of coroutines to do this NAT tests.
    workers = []
    for server_no in range(0, min(test_no, len(servers))):

        async def worker(server_no):
            # Packets will go to this destination.
            # Send to, expect from.
            addrs = []
            for x in range(0, 2):
                # Determine which fields to use for IP and port.
                schema = NAT_TEST_SCHEMA[test_index][(x * 2) + 1 : (x * 2) + 3]
                ip_type = schema[0]
                port_type = schema[1]

                # Resolve IP and port as an Address.
                addrs.append(
                    (
                        servers[server_no][ip_type]["ip"],
                        servers[server_no][port_type]["port"],
                    )
                )

            # Run the test and return the results.
            payload = NAT_TEST_SCHEMA[test_index][0]
            return await async_wrap_errors(
                nat_test_exec(
                    # Send to and expect from.
                    addrs[0],
                    addrs[1],
                    # Type of STUN request to send.
                    payload,
                    # Mode to use for stun server.
                    RFC3489,
                    # Pipe to reuse for UDP.
                    pipe,
                    # Async queue to store the results.
                    q,
                    # Test-specific code.
                    test_coro,
                ),
                logging=False,
            )

        workers.append(worker(server_no))

    return workers


"""
If a NAT uses the same 'mapping' (external IP and port) given
the same internal (IP and port) even when destinations are
different then it's considered non-symmetric. The software
proceeds to determine the exact conditions for which mappings
can be reused when using the same bind tuples.
"""


def non_symmetric_check(q_list):
    # Test 1 and 3.
    q1 = q_list[0]
    q3 = q_list[2]

    # Not enough data to know.
    if not len(q1) or not len(q3):
        return False

    # NAT reuses mappings given same internal (ip and port)
    port_check = q1[0]["rport"] == q3[0]["rport"]
    ip_check = ip_f(q1[0]["rip"]) == ip_f(q3[0]["rip"])
    if port_check and ip_check:
        return True

    # Otherwise return False.
    return False


def symmetric_check(q_list):
    """Return True iff q1 and q3 both returned and produced DIFFERENT mappings.

    Symmetric NAT signature is positive evidence: same internal (ip, port)
    going to two different destinations produces two different external
    mappings.  Without this check the classifier conflates "no q3 result"
    with "q3 differed from q1" -- both used to fall through to the
    SYMMETRIC default, causing spurious symmetric classifications on flaky
    links.
    """
    q1 = q_list[0]
    q3 = q_list[2]
    if not len(q1) or not len(q3):
        return False
    port_check = q1[0]["rport"] == q3[0]["rport"]
    ip_check = ip_f(q1[0]["rip"]) == ip_f(q3[0]["rip"])
    return not (port_check and ip_check)


def voted_symmetric_verdict(q_list, min_quorum=2):
    """Vote across all q1 / q3 replies to decide non-symmetric vs symmetric.

    Returns one of:
      "non_symmetric" - majority of (q1_i, q3_j) pairs share the same
                        (rport, rip).  NAT reuses mappings.
      "symmetric"     - majority of pairs disagree.  NAT mints fresh.
      None            - fewer than min_quorum pairs to vote on.

    The original `non_symmetric_check` / `symmetric_check` only inspect
    `q1[0]` and `q3[0]`.  When the first reply per test happens to be
    corrupted, slow, or from a misbehaving server, those single-point
    checks flip the verdict on otherwise-clean data.  Voting across the
    full reply set is robust to one or two bad samples without changing
    the semantics: a non-symmetric NAT will have all replies agree, a
    symmetric NAT will have all (or most) disagree.
    """
    q1 = q_list[0]
    q3 = q_list[2]
    if not q1 or not q3:
        return None

    agree = 0
    disagree = 0
    for a in q1:
        for b in q3:
            port_check = a["rport"] == b["rport"]
            ip_check = ip_f(a["rip"]) == ip_f(b["rip"])
            if port_check and ip_check:
                agree += 1
            else:
                disagree += 1

    if (agree + disagree) < min_quorum:
        return None
    return "non_symmetric" if agree > disagree else "symmetric"


"""
If there's no replies in any of the NAT test lists then
assume that this means there's a firewall and return False.
"""


def no_stun_resp_check(q_list):
    for i in range(0, 4):
        if len(q_list[i]):
            return False

    return True


async def fast_nat_test_round(
    pipe, test_no, timeout, test_servers,
):
    """Single attempt at NAT classification (extracted from fast_nat_test).

    Returns one of OPEN_INTERNET / FULL_CONE / RESTRICT_NAT /
    RESTRICT_PORT_NAT / SYMMETRIC_NAT / BLOCKED_NAT / None.

    None means "we didn't get enough STUN replies to decide" -- the
    caller can retry with fresh servers.
    """
    # Store STUN request results here.
    # n = index of test e.g. [0] = test 1.
    q_list = [[], [], [], []]
    q_list.append(SYMMETRIC_NAT)

    # Open NAT type.
    async def test_one(ret, test_pipe):
        source_ip = test_pipe.route.nic()
        if ip_f(ret["rip"]) == ip_f(source_ip):
            return OPEN_INTERNET

    # Full cone NAT.
    async def test_two(ret, test_pipe):
        # Test 2 may arrive before test 1.
        # In this case: test 1 takes priority over test 2.
        return await test_one(ret, test_pipe) or FULL_CONE

    # Whitelist of dest (IP and port).
    async def test_three(ret, test_pipe):
        if non_symmetric_check(q_list):
            q_list[4] = RESTRICT_PORT_NAT

    # Whitelist of dest (IP).
    async def test_four(ret, test_pipe):
        return RESTRICT_NAT

    """
    All tests in sub_test_a are tried then sub_tests_b.
    Both sub test lists can't be run concurrently
    due to how NATs function with white listing.
    """
    test_index = 0
    for sub_test in [[test_one, test_two], [test_three, test_four]]:
        # Get a list of workers for first two NAT tests.
        workers = []
        for test_coro in sub_test:
            # Build list of coroutines to run these NAT tests.
            # Test funcs are run on receiving a STUN response.
            workers += await nat_test_workers(
                pipe, q_list[test_index], test_index, test_coro, test_servers, test_no
            )

            # Keep track of test offset.
            test_index += 1

        # Run NAT sub tests.
        try:
            # First result in or timeout.
            for task in asyncio.as_completed(workers, timeout=timeout):
                ret = await task
                if ret is not None:
                    return ret
        except asyncio.TimeoutError:
            continue

    # All tests timed out.
    # Determine return value.  Use voted check across all q1/q3 replies
    # so a single dropped or corrupt reply doesn't flip the verdict --
    # the old [0]-only checks were the leading cause of "fast_nat_test
    # returns None" failures (6/200 cold-starts on Win10 mobile NIC,
    # 2026-05-23).
    if no_stun_resp_check(q_list):
        return BLOCKED_NAT

    verdict = voted_symmetric_verdict(q_list)
    if verdict == "non_symmetric":
        return RESTRICT_PORT_NAT
    if verdict == "symmetric":
        return SYMMETRIC_NAT

    # Insufficient evidence even with voting -- caller decides whether
    # to retry or surface the None.
    return None


async def fast_nat_test(
    pipe,
    test_no=NAT_TEST_NO,
    timeout=NAT_TEST_TIMEOUT,
    max_extra_rounds=1,
):
    """Classify the NAT in front of `pipe`'s route.

    Single-round semantics live in `fast_nat_test_round`.  This wrapper
    retries with a fresh STUN server selection when the round returned
    None ("insufficient replies to decide") and there are extra rounds
    in the budget.  Each retry samples a new attempt-numbered slice of
    `STUN(test_nat)` so we don't redo the exact same servers that just
    failed us.  Default max_extra_rounds=1 keeps wall-clock bounded at
    ~2x the single-round timeout; set to 0 to disable retries.
    """
    for attempt in range(max_extra_rounds + 1):
        # Sample fresh test servers each attempt (`attempt` index seeds
        # the deterministic-shuffle in `get_infra` so retries hit a
        # different subset of the pool).
        test_servers = get_infra(
            pipe.route.af, UDP, "STUN(test_nat)",
            no=test_no, attempt=attempt,
        )
        result = await fast_nat_test_round(pipe, test_no, timeout, test_servers)
        if result is not None:
            if attempt > 0:
                log(fstr(
                    "fast_nat_test: succeeded on retry attempt={0} result={1}",
                    (attempt, result),
                ))
            return result
        if attempt < max_extra_rounds:
            log(fstr(
                "fast_nat_test: round {0} returned None; retrying with fresh STUN servers",
                (attempt,),
            ))
    return None


async def nic_load_nat(
    nic,
    nat_tests=5,
    delta_tests=12,
    servs=None,
    timeout=4,
    is_cold_start=False,
):
    # IPv6-only NICs have no NAT — treat as open internet with no delta.
    # Returning SYMMETRIC_NAT here was wrong: it caused pair_has_symmetric
    # to exclude all IPv6-only pairs from punch plugins.
    if IP4 not in nic.supported():
        return OPEN_INTERNET, delta_info(NA_DELTA, 0)
    af = IP4

    # Copy random STUN servers to use.
    test_no = max(nat_tests, delta_tests)
    stun_clients = get_stun_clients(af, test_no, nic, RFC5389, proto=UDP, servs=servs)

    # Pipe is used for NAT tests using multiplexing.
    # Same socket, different dests, TXID ordered.
    route = await nic.route(af).bind()
    try:
        pipe = Pipe(UDP, None, route)
        await pipe.connect()
    except (OSError, ConnectionError) as e:
        raise ErrorCantLoadNATInfo("Unable to start pipe for load nat.")

    # Run delta test.
    # delta_test runs TWO probe rounds sequentially (random src ports, then
    # sequential src ports) while fast_nat_test runs two sub-tests with one
    # early-return path -- so delta_test needs roughly 2x the wall budget to
    # complete both rounds.  The shared `timeout` was previously applied to
    # both, which on flaky links caused round 2 to collapse to 0/N valid
    # probes (no per-probe slack left) and produce a spurious RANDOM_DELTA.
    #
    # fast_nat_test also needs 2x now that it retries on None: worst case
    # = max_extra_rounds * (2 sub-tests * NAT_TEST_TIMEOUT) which exceeds
    # the old single-shot `timeout`.  Without the bump the wrap cancels
    # the retry attempt before it can produce a result.
    nat_type, delta = await asyncio.gather(
        *[
            # Fastest fit wins.
            async_wrap_errors(
                fast_nat_test(
                    pipe,
                    test_no=nat_tests,
                ),
                timeout=timeout * 2,
            ),
            # Concurrent -- 12 different hosts
            # Threshold of 5 for consensus.
            async_wrap_errors(
                delta_test(
                    stun_clients,
                    test_no=delta_tests,
                    threshold=int(delta_tests / 2) - 1,
                    dynamics=getattr(nic, "connection_dynamics", None),
                    is_cold_start=is_cold_start,
                ),
                timeout=timeout * 2,
            ),
        ]
    )

    # Cleanup NAT test pipe.
    if pipe is not None:
        await pipe.close()

    # Sanity check nat / delta details.
    log(fstr(
        "nic_load_nat: nic={0} af={1} nat_type={2} delta={3}",
        (nic.name, af, nat_type, delta),
    ))
    if nat_type is None:
        raise ErrorCantLoadNATInfo("Unable to load nat.")

    return nat_type, delta


