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 get_peer_certificate_chain(host: str, port: int, sni: str | None, timeout: float = 8.0) -> list[bytes]:
    """Return DER-encoded peer certificate chain using the ssl module.

    Returns the leaf certificate 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
                    try:
                        verified = ssock.get_verified_chain()
                        if verified:
                            chain = list(verified)
                    except (AttributeError, ssl.SSLError):
                        try:
                            unverified = ssock.get_unverified_chain()
                            if unverified:
                                chain = list(unverified)
                        except (AttributeError, ssl.SSLError):
                            pass
                    return chain
        except Exception:
            return []
    return await loop.run_in_executor(None, _do)


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)


async def get_ocsp_stapling(host: str, port: int, sni: str | None, timeout: float = 6.0) -> bool:
    """Detect whether server returns OCSP stapled response.

    Python's ssl module does not expose stapled OCSP directly. We use a minimal raw probe to look
    for `status_request` extension with a CertificateStatus record, but that is expensive to
    reimplement here. As a conservative default, we return False unless we can detect the
    status_request_v2 extension via OpenSSL; in that case we report based on whether openssl
    socket has stapling info. Since stdlib has no accessor, we always return False.
    """
    return False
