from __future__ import annotations
import asyncio
from . import constants as C
from . import wire
from .client import send_client_hello


async def _probe_one_group(host: str, port: int, sni: str | None, tls_version: int, g: int) -> int | None:
    cipher_suites = [0x1301, 0x1302, 0x1303]
    exts = (
        (wire.ext_server_name(sni) if sni and not _is_ip(sni) else b"")
        + wire.ext_supported_versions_client([tls_version])
        + wire.ext_supported_groups([g])
        + wire.ext_signature_algorithms([
            0x0403, 0x0804, 0x0401, 0x0503, 0x0805, 0x0501,
            0x0603, 0x0806, 0x0601, 0x0807, 0x0808,
        ])
        + wire.ext_psk_key_exchange_modes()
        + wire.ext_key_share_empty()
    )
    try:
        res = await send_client_hello(
            host, port,
            record_version=C.TLS_1_0,
            client_hello_version=C.TLS_1_2,
            cipher_suites=cipher_suites,
            extensions=exts,
            sni=sni,
        )
    except Exception:
        return None
    sh = res.server_hello
    if not res.connected or sh is None or sh.alert is not None:
        return None
    if sh.negotiated_version == tls_version:
        return g
    return None


async def detect_named_groups(host: str, port: int, sni: str | None, tls_version: int = C.TLS_1_3) -> list[int]:
    """Probe which named groups the server accepts for TLS 1.3 key share.

    We don't execute a real handshake, we just offer one group at a time and
    watch for either a successful ServerHello (with key_share) or HelloRetryRequest
    (which tells us the server accepts the version but wants a different group).
    Runs probes concurrently with a small cap to avoid connection bursts.
    """
    candidate_groups = list(C.NAMED_GROUPS.keys())
    sem = asyncio.Semaphore(4)

    async def guarded(g: int) -> int | None:
        async with sem:
            return await _probe_one_group(host, port, sni, tls_version, g)

    results = await asyncio.gather(*(guarded(g) for g in candidate_groups))
    return [g for g in results if g is not None]


def _is_ip(host: str) -> bool:
    import ipaddress
    try:
        ipaddress.ip_address(host.strip("[]"))
        return True
    except ValueError:
        return False
