from __future__ import annotations
import asyncio

from ..protocol import constants as C
from ..protocol import wire


async def _read_records(reader: asyncio.StreamReader, max_records: int = 8, timeout: float = 6.0) -> bytes:
    """Read up to `max_records` complete TLS records from the stream.

    A simple record-layer reader that guarantees we pull entire records
    instead of relying on whatever arrives in the first TCP chunk. Records
    that straddle packet boundaries (common for ServerHello+Certificate
    responses in TLS 1.2) are reassembled.
    """
    buf = bytearray()
    for _ in range(max_records):
        try:
            header = await asyncio.wait_for(reader.readexactly(5), timeout=timeout)
        except (asyncio.IncompleteReadError, asyncio.TimeoutError, ConnectionError):
            break
        except Exception:
            break
        rec_len = (header[3] << 8) | header[4]
        body = b""
        if rec_len:
            try:
                body = await asyncio.wait_for(reader.readexactly(rec_len), timeout=timeout)
            except (asyncio.IncompleteReadError, asyncio.TimeoutError, ConnectionError):
                buf.extend(header)
                break
            except Exception:
                buf.extend(header)
                break
        buf.extend(header)
        buf.extend(body)
        ct = header[0]
        # A fatal alert or end of useful data — stop.
        if ct == C.CT_ALERT and len(body) >= 2 and body[0] == 2:
            break
    return bytes(buf)


def _iter_records(data: bytes):
    i = 0
    while i + 5 <= len(data):
        ct = data[i]
        rec_len = (data[i + 3] << 8) | data[i + 4]
        if i + 5 + rec_len > len(data):
            break
        yield ct, data[i + 5:i + 5 + rec_len]
        i += 5 + rec_len
        if rec_len == 0:
            break


async def probe_secure_renegotiation(host: str, port: int, sni: str | None, timeout: float = 8.0) -> tuple[bool, str | None]:
    """Check whether the server acknowledges RFC 5746 (Secure Renegotiation).

    Behavior:
      - Send a ClientHello that offers only TLS 1.2 (no supported_versions extension,
        so TLS 1.3 servers will downgrade to 1.2 for this probe) and includes
        `renegotiation_info` (empty) per RFC 5746.
      - A server that supports RFC 5746 MUST echo the `renegotiation_info`
        extension in its ServerHello.
      - If the server *only* supports TLS 1.3 (e.g., modern fronts that refuse
        TLS 1.2), RFC 5746 does not apply — TLS 1.3 removed renegotiation
        entirely and is treated as "supported by way of elimination".

    Returns (supported, info_string).
    """
    try:
        reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=timeout)
    except Exception as e:
        return False, f"connect: {e}"
    try:
        exts = (
            (wire.ext_server_name(sni) if sni and not wire._is_ip_literal(sni) else b"")
            + wire.ext_ec_point_formats()
            + wire.ext_supported_groups([0x001d, 0x0017, 0x0018, 0x0019])
            + wire.ext_signature_algorithms([
                0x0403, 0x0804, 0x0401, 0x0503, 0x0805, 0x0501,
                0x0603, 0x0806, 0x0601, 0x0807, 0x0808,
            ])
            + wire.ext_renegotiation_info_empty()
            + wire.ext_extended_master_secret()
        )
        ch = wire.build_client_hello(
            record_version=C.TLS_1_0,
            client_hello_version=C.TLS_1_2,
            hostname=sni,
            cipher_suites=[
                0xc02f, 0xc030, 0xc02b, 0xc02c,
                0xcca8, 0xcca9,
                0xc013, 0xc014, 0x009c, 0x009d,
                0x002f, 0x0035, 0x000a,
            ],
            extensions=exts,
        )
        writer.write(ch)
        try:
            await asyncio.wait_for(writer.drain(), timeout=timeout)
        except Exception:
            pass
        data = await _read_records(reader, max_records=3, timeout=timeout)
        if not data:
            return False, "no response"
        parsed = wire.parse_server_response(data)
        if parsed is None or parsed.handshake_type != C.HS_SERVER_HELLO:
            # If the server returned only an alert on a TLS 1.2 ClientHello we can't
            # observe renegotiation_info; report unknown-leaning-negative.
            if parsed and parsed.alert is not None:
                return False, f"server alert (level={parsed.alert[0]}, desc={parsed.alert[1]})"
            return False, "no server hello"
        # TLS 1.3 removed renegotiation entirely (RFC 8446 §4.1.2). A server that
        # negotiates TLS 1.3 here is safe-by-default.
        neg_ver = parsed.negotiated_version or parsed.server_version
        if neg_ver == C.TLS_1_3:
            return True, "TLS 1.3 (renegotiation obsolete)"
        if parsed.extensions and C.EXT_RENEGOTIATION_INFO in parsed.extensions:
            return True, None
        return False, "no renegotiation_info extension echoed"
    finally:
        try:
            writer.close()
            await asyncio.wait_for(writer.wait_closed(), timeout=1.0)
        except Exception:
            pass


async def probe_fallback_scsv(host: str, port: int, sni: str | None, timeout: float = 8.0) -> tuple[bool, str | None]:
    """Check whether the server enforces TLS_FALLBACK_SCSV (RFC 7507).

    Strategy:
      - Send a ClientHello with legacy_version=TLS 1.1, cipher list containing
        TLS_FALLBACK_SCSV (0x5600), and NO supported_versions extension.
      - If the server's highest supported version is > TLS 1.1, it MUST respond
        with fatal alert `inappropriate_fallback` (86) per RFC 7507.
      - If the server genuinely has TLS 1.1 as its max (rare today), it will
        complete the handshake normally — in that case SCSV enforcement is
        moot. We treat that as `n/a` rather than `not enforced`.

    Returns (enforced_or_not_applicable, info_string).
    """
    try:
        reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=timeout)
    except Exception as e:
        return False, f"connect: {e}"
    try:
        exts = (
            (wire.ext_server_name(sni) if sni and not wire._is_ip_literal(sni) else b"")
            + wire.ext_ec_point_formats()
            + wire.ext_supported_groups([0x001d, 0x0017, 0x0018])
            + wire.ext_signature_algorithms([
                0x0403, 0x0401, 0x0501, 0x0601,
            ])
            + wire.ext_renegotiation_info_empty()
        )
        ch = wire.build_client_hello(
            record_version=C.TLS_1_0,
            client_hello_version=C.TLS_1_1,
            hostname=sni,
            cipher_suites=[
                0x5600,                      # TLS_FALLBACK_SCSV
                0xc013, 0xc014, 0xc009, 0xc00a,
                0x002f, 0x0035, 0x000a,
            ],
            extensions=exts,
        )
        writer.write(ch)
        try:
            await asyncio.wait_for(writer.drain(), timeout=timeout)
        except Exception:
            pass
        data = await _read_records(reader, max_records=3, timeout=timeout)
        if not data:
            # A fatal alert may cause the server to abort the TCP connection
            # before the alert record is flushed on some stacks. Treat as
            # unknown-leaning-negative, but also probe once more without SCSV
            # to distinguish from "server totally broken".
            return False, "no response (possibly RST)"
        saw_server_hello = False
        negotiated_version: int | None = None
        for ct, body in _iter_records(data):
            if ct == C.CT_ALERT and len(body) >= 2:
                level, desc = body[0], body[1]
                if desc == 86:
                    return True, None
                # Some servers reply with handshake_failure(40) / protocol_version(70)
                # when downgrade is refused; treat as enforcement.
                if desc in (70,) and level == 2:
                    return True, f"alert protocol_version({desc})"
            if ct == C.CT_HANDSHAKE and len(body) >= 4 and body[0] == C.HS_SERVER_HELLO:
                saw_server_hello = True
                if len(body) >= 6:
                    negotiated_version = (body[4] << 8) | body[5]
        if saw_server_hello:
            if negotiated_version == C.TLS_1_1:
                # Server genuinely maxes out at TLS 1.1 → SCSV is not applicable.
                return True, "server max = TLS 1.1 (SCSV not applicable)"
            return False, "server completed handshake without inappropriate_fallback"
        return False, "no inappropriate_fallback alert"
    finally:
        try:
            writer.close()
            await asyncio.wait_for(writer.wait_closed(), timeout=1.0)
        except Exception:
            pass
