from __future__ import annotations
import asyncio
import ipaddress
import socket
from dataclasses import dataclass

from . import constants as C
from . import wire

CONNECT_TIMEOUT = 6.0
IO_TIMEOUT = 8.0


@dataclass
class ProbeResult:
    connected: bool
    server_hello: wire.ParsedServerHello | None
    alert: tuple[int, int] | None
    raw: bytes
    error: str | None = None


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


async def _open(host: str, port: int, timeout: float = CONNECT_TIMEOUT) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
    return await asyncio.wait_for(asyncio.open_connection(host, port), timeout=timeout)


async def send_client_hello(
    host: str,
    port: int,
    *,
    record_version: int = C.TLS_1_0,
    client_hello_version: int = C.TLS_1_2,
    cipher_suites: list[int],
    extensions: bytes = b"",
    sni: str | None = None,
    timeout: float = IO_TIMEOUT,
) -> ProbeResult:
    try:
        reader, writer = await _open(host, port)
    except Exception as e:
        return ProbeResult(False, None, None, b"", str(e))
    try:
        ch = wire.build_client_hello(
            record_version=record_version,
            client_hello_version=client_hello_version,
            hostname=sni,
            cipher_suites=cipher_suites,
            extensions=extensions,
        )
        writer.write(ch)
        try:
            await asyncio.wait_for(writer.drain(), timeout=timeout)
        except Exception:
            pass
        data = b""
        try:
            data = await asyncio.wait_for(reader.read(4096), timeout=timeout)
        except Exception:
            pass
        parsed = wire.parse_server_response(data) if data else None
        alert = parsed.alert if parsed else None
        return ProbeResult(True, parsed, alert, data)
    except Exception as e:
        return ProbeResult(True, None, None, b"", str(e))
    finally:
        try:
            writer.close()
            await asyncio.wait_for(writer.wait_closed(), timeout=1.5)
        except Exception:
            pass


async def send_ssl2_client_hello(host: str, port: int) -> ProbeResult:
    try:
        reader, writer = await _open(host, port)
    except Exception as e:
        return ProbeResult(False, None, None, b"", str(e))
    try:
        writer.write(wire.build_ssl2_client_hello())
        try:
            await asyncio.wait_for(writer.drain(), timeout=IO_TIMEOUT)
        except Exception:
            pass
        try:
            data = await asyncio.wait_for(reader.read(4096), timeout=IO_TIMEOUT)
        except Exception:
            data = b""
        # In SSLv2 a SERVER-HELLO response has msg_type 4 in the body.
        is_sslv2 = False
        if len(data) >= 3:
            # 2-byte length header with MSB set is SSLv2, then body[0] = msg type
            if data[0] & 0x80:
                length = ((data[0] & 0x7f) << 8) | data[1]
                if length >= 1 and len(data) >= 3 and data[2] == 4:
                    is_sslv2 = True
        alert = None if is_sslv2 else (0, 0)
        parsed = None
        return ProbeResult(True, parsed, alert, data)
    finally:
        try:
            writer.close()
            await asyncio.wait_for(writer.wait_closed(), timeout=1.5)
        except Exception:
            pass


async def tcp_reachable(host: str, port: int, timeout: float = CONNECT_TIMEOUT) -> bool:
    try:
        _, w = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=timeout)
        try:
            w.close()
            try:
                await asyncio.wait_for(w.wait_closed(), timeout=1.0)
            except Exception:
                pass
        except Exception:
            pass
        return True
    except Exception:
        return False


async def resolve_host(host: str) -> str | None:
    if is_ip_literal(host):
        return host.strip("[]")
    loop = asyncio.get_running_loop()
    try:
        infos = await loop.getaddrinfo(host, None, type=socket.SOCK_STREAM)
    except Exception:
        return None
    for family, _, _, _, sockaddr in infos:
        if family == socket.AF_INET:
            return sockaddr[0]
    for family, _, _, _, sockaddr in infos:
        if family == socket.AF_INET6:
            return sockaddr[0]
    return None
