from __future__ import annotations
import asyncio
import socket
import ssl
from dataclasses import dataclass
from typing import Any

from . import constants as C
from . import wire
from .client import send_client_hello

SSL_VERSION_TO_CONST = {
    C.TLS_1_0: ssl.TLSVersion.TLSv1,
    C.TLS_1_1: ssl.TLSVersion.TLSv1_1,
    C.TLS_1_2: ssl.TLSVersion.TLSv1_2,
    C.TLS_1_3: ssl.TLSVersion.TLSv1_3,
}


@dataclass
class VersionProbe:
    version: int
    supported: bool
    negotiated_cipher: str | None = None
    error: str | None = None


def _make_permissive_context(min_v: ssl.TLSVersion, max_v: ssl.TLSVersion, ciphers: str | None = None) -> ssl.SSLContext:
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    try:
        ctx.minimum_version = min_v
        ctx.maximum_version = max_v
    except (ValueError, OSError):
        pass
    if ciphers is not None:
        try:
            ctx.set_ciphers(ciphers)
        except ssl.SSLError:
            pass
    # Allow legacy server connections (e.g. missing renegotiation_info for TLS 1.0 servers).
    try:
        ctx.options &= ~ssl.OP_NO_SSLv3
    except Exception:
        pass
    try:
        # SECLEVEL=0 allows legacy/weak ciphers for enumeration.
        cur_ciphers = ciphers if ciphers else "ALL:@SECLEVEL=0"
        ctx.set_ciphers(cur_ciphers)
    except ssl.SSLError:
        pass
    return ctx


async def probe_tls_version(host: str, port: int, sni: str | None, version: int, timeout: float = 8.0) -> VersionProbe:
    if version not in SSL_VERSION_TO_CONST:
        return VersionProbe(version=version, supported=False, error="unsupported version constant")
    tls_ver = SSL_VERSION_TO_CONST[version]
    ctx = _make_permissive_context(tls_ver, tls_ver)
    loop = asyncio.get_running_loop()

    def _do() -> tuple[bool, str | None, str | None]:
        try:
            with socket.create_connection((host, port), timeout=timeout) as sock:
                with ctx.wrap_socket(sock, server_hostname=sni) as ssock:
                    try:
                        cipher = ssock.cipher()
                        return True, cipher[0] if cipher else None, None
                    except Exception as e:
                        return True, None, str(e)
        except ssl.SSLError as e:
            return False, None, f"{e.__class__.__name__}: {e}"
        except (OSError, socket.timeout, asyncio.TimeoutError) as e:
            return False, None, f"{e.__class__.__name__}: {e}"
        except Exception as e:
            return False, None, f"{e.__class__.__name__}: {e}"

    ok, cipher, err = await loop.run_in_executor(None, _do)
    return VersionProbe(version=version, supported=ok, negotiated_cipher=cipher, error=err)


@dataclass
class CipherProbe:
    version: int
    cipher_id: int
    cipher_name: str
    supported: bool


async def probe_cipher(host: str, port: int, sni: str | None, version: int, cipher_id: int, cipher_name: str) -> CipherProbe:
    """Enumerate whether a given cipher suite is accepted under a given TLS version.

    For TLS 1.2 and below we use the ssl module with set_ciphers().
    For TLS 1.3 we use set_ciphersuites() equivalents via the OpenSSL names.
    Falls back to False on any handshake error.
    """
    if version == C.TLS_1_3:
        # TLS 1.3 only recognizes cipher suites by their IANA name in OpenSSL.
        name_map = {
            0x1301: "TLS_AES_128_GCM_SHA256",
            0x1302: "TLS_AES_256_GCM_SHA384",
            0x1303: "TLS_CHACHA20_POLY1305_SHA256",
            0x1304: "TLS_AES_128_CCM_SHA256",
            0x1305: "TLS_AES_128_CCM_8_SHA256",
        }
        openssl_name = name_map.get(cipher_id)
        if not openssl_name:
            return CipherProbe(version, cipher_id, cipher_name, False)
        tls_ver = SSL_VERSION_TO_CONST[version]
        ctx = _make_permissive_context(tls_ver, tls_ver)
        try:
            ctx.set_ciphersuites(openssl_name)
        except Exception:
            return CipherProbe(version, cipher_id, cipher_name, False)
        loop = asyncio.get_running_loop()

        def _do() -> bool:
            try:
                with socket.create_connection((host, port), timeout=6.0) as sock:
                    with ctx.wrap_socket(sock, server_hostname=sni) as ssock:
                        actual = ssock.cipher()
                        return bool(actual and actual[0] == openssl_name)
            except Exception:
                return False
        ok = await loop.run_in_executor(None, _do)
        return CipherProbe(version, cipher_id, cipher_name, ok)

    # For < TLS 1.3: send a ClientHello offering only this suite.
    exts = b""
    parts = []
    if sni:
        try:
            parts.append(wire.ext_server_name(sni))
        except Exception:
            pass
    parts.append(wire.ext_ec_point_formats())
    parts.append(wire.ext_supported_groups([0x001d, 0x0017, 0x0018]))
    parts.append(wire.ext_signature_algorithms([
        0x0403, 0x0804, 0x0401, 0x0503, 0x0805, 0x0501,
        0x0603, 0x0806, 0x0601,
    ]))
    parts.append(wire.ext_renegotiation_info_empty())
    exts = b"".join(parts)

    res = await send_client_hello(
        host, port,
        record_version=C.TLS_1_0,
        client_hello_version=version,
        cipher_suites=[cipher_id],
        extensions=exts,
        sni=sni,
    )
    if not res.connected:
        return CipherProbe(version, cipher_id, cipher_name, False)
    sh = res.server_hello
    if sh is None or sh.alert is not None:
        return CipherProbe(version, cipher_id, cipher_name, False)
    ok = sh.cipher_suite == cipher_id and sh.server_version == version
    return CipherProbe(version, cipher_id, cipher_name, ok)


async def _fetch_chain_via_wire(host: str, port: int, sni: str | None, timeout: float = 8.0) -> list[bytes]:
    """Pull the full TLS 1.2 Certificate handshake message ourselves and
    return the DER-encoded chain. Used as fallback on Python < 3.13 where
    ``ssl.SSLSocket.get_unverified_chain`` isn't available.
    """
    parts = []
    if sni:
        try:
            parts.append(wire.ext_server_name(sni))
        except Exception:
            pass
    parts.append(wire.ext_ec_point_formats())
    parts.append(wire.ext_supported_groups([0x001d, 0x0017, 0x0018, 0x0019]))
    parts.append(wire.ext_signature_algorithms([
        0x0403, 0x0804, 0x0401, 0x0503, 0x0805, 0x0501,
        0x0603, 0x0806, 0x0601, 0x0807, 0x0808,
    ]))
    parts.append(wire.ext_renegotiation_info_empty())
    exts = b"".join(parts)

    try:
        reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=timeout)
    except Exception:
        return []
    try:
        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
        # Read enough records to cover ServerHello + Certificate; TLS 1.2
        # Certificate messages are often fragmented across several records.
        buf = bytearray()
        deadline = timeout
        for _ in range(20):
            try:
                header = await asyncio.wait_for(reader.readexactly(5), timeout=deadline)
            except Exception:
                break
            rec_len = (header[3] << 8) | header[4]
            if rec_len == 0 or rec_len > 1 << 14:
                break
            try:
                body = await asyncio.wait_for(reader.readexactly(rec_len), timeout=deadline)
            except Exception:
                break
            buf.extend(header); buf.extend(body)
            # Stop once we've seen a Certificate handshake in the accumulated buffer.
            if _has_certificate_message(bytes(buf)):
                break
        return _extract_certificates_from_records(bytes(buf))
    finally:
        try:
            writer.close()
            await asyncio.wait_for(writer.wait_closed(), timeout=1.0)
        except Exception:
            pass


def _has_certificate_message(records: bytes) -> bool:
    i = 0
    hs_buf = bytearray()
    while i + 5 <= len(records):
        ct = records[i]
        rlen = (records[i + 3] << 8) | records[i + 4]
        if i + 5 + rlen > len(records):
            break
        if ct == C.CT_HANDSHAKE:
            hs_buf.extend(records[i + 5:i + 5 + rlen])
        i += 5 + rlen
    j = 0
    while j + 4 <= len(hs_buf):
        hs_type = hs_buf[j]
        hs_len = (hs_buf[j + 1] << 16) | (hs_buf[j + 2] << 8) | hs_buf[j + 3]
        if j + 4 + hs_len > len(hs_buf):
            return False
        if hs_type == C.HS_CERTIFICATE:
            return True
        j += 4 + hs_len
    return False


def _extract_certificates_from_records(records: bytes) -> list[bytes]:
    """Parse TLS 1.2 Certificate handshake message and return the DER chain."""
    # 1) Reassemble handshake layer from all handshake-typed records.
    i = 0
    hs_buf = bytearray()
    while i + 5 <= len(records):
        ct = records[i]
        rlen = (records[i + 3] << 8) | records[i + 4]
        if i + 5 + rlen > len(records):
            break
        if ct == C.CT_HANDSHAKE:
            hs_buf.extend(records[i + 5:i + 5 + rlen])
        i += 5 + rlen
    # 2) Walk handshake messages, find Certificate.
    j = 0
    while j + 4 <= len(hs_buf):
        hs_type = hs_buf[j]
        hs_len = (hs_buf[j + 1] << 16) | (hs_buf[j + 2] << 8) | hs_buf[j + 3]
        if j + 4 + hs_len > len(hs_buf):
            break
        if hs_type == C.HS_CERTIFICATE and hs_len >= 3:
            body = bytes(hs_buf[j + 4:j + 4 + hs_len])
            # TLS 1.2 Certificate: uint24 total_len || [ uint24 cert_len || cert ] *
            total = (body[0] << 16) | (body[1] << 8) | body[2]
            k = 3
            end = 3 + total
            chain: list[bytes] = []
            while k + 3 <= end:
                clen = (body[k] << 16) | (body[k + 1] << 8) | body[k + 2]
                k += 3
                if k + clen > end:
                    break
                chain.append(body[k:k + clen])
                k += clen
            return chain
        j += 4 + hs_len
    return []


async def get_peer_certificate_chain(host: str, port: int, sni: str | None, timeout: float = 8.0) -> list[bytes]:
    """Return DER-encoded peer certificate chain.

    Tries the ssl module's Python-3.13+ chain APIs first, then falls back to
    a raw TLS 1.2 wire probe that parses the Certificate handshake message
    directly. Returns leaf first, then intermediates. Empty on failure.
    """
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    try:
        ctx.set_ciphers("ALL:@SECLEVEL=0")
    except ssl.SSLError:
        pass
    loop = asyncio.get_running_loop()

    def _do() -> list[bytes]:
        try:
            with socket.create_connection((host, port), timeout=timeout) as sock:
                with ctx.wrap_socket(sock, server_hostname=sni) as ssock:
                    leaf = ssock.getpeercert(binary_form=True)
                    chain: list[bytes] = [leaf] if leaf else []
                    # Intermediate chain via internal API (CPython 3.13+) if available.
                    # On older Pythons these attributes are absent — callers will
                    # fall through to the wire-level fetcher below.
                    try:
                        verified = ssock.get_verified_chain()  # type: ignore[attr-defined]
                        if verified:
                            chain = list(verified)
                    except (AttributeError, ssl.SSLError):
                        try:
                            unverified = ssock.get_unverified_chain()  # type: ignore[attr-defined]
                            if unverified:
                                chain = list(unverified)
                        except (AttributeError, ssl.SSLError):
                            pass
                    return chain
        except Exception:
            return []

    chain = await loop.run_in_executor(None, _do)
    # If the ssl module only gave us a single cert (common on Python < 3.13),
    # fetch the full chain via our own TLS wire probe.
    if len(chain) <= 1:
        try:
            wire_chain = await _fetch_chain_via_wire(host, port, sni, timeout=timeout)
        except Exception:
            wire_chain = []
        if len(wire_chain) > len(chain):
            chain = wire_chain
    return chain


async def alpn_negotiate(host: str, port: int, sni: str | None, alpn_list: list[str], timeout: float = 6.0) -> str | None:
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    try:
        ctx.set_alpn_protocols(alpn_list)
    except NotImplementedError:
        return None
    try:
        ctx.set_ciphers("ALL:@SECLEVEL=0")
    except ssl.SSLError:
        pass
    loop = asyncio.get_running_loop()

    def _do() -> str | None:
        try:
            with socket.create_connection((host, port), timeout=timeout) as sock:
                with ctx.wrap_socket(sock, server_hostname=sni) as ssock:
                    return ssock.selected_alpn_protocol()
        except Exception:
            return None
    return await loop.run_in_executor(None, _do)


# Note: OCSP stapling detection requires probing the `status_request` TLS
# extension and reading the CertificateStatus handshake message. Python's
# stdlib ssl module does not expose stapled OCSP to Python; the engine checks
# live OCSP/CRL instead (see certs/revocation.py).
