"""
A small tool that connects to a local port and tries to read the first message
(the “banner”) that the service sends back. Many services announce themselves
right after a connection is made, so grabbing this banner can help identify
what’s running on the port.

The probe is designed to be safe: it only works on local addresses, and it uses
the correct IPv4/IPv6 socket family so it behaves consistently on Windows. It
records how long the attempt took, whether the connection succeeded, and any
errors that happened along the way.

The result is returned as a BannerProbeResult, which includes the raw bytes
received, a decoded text version, and simple fields explaining why the probe
succeeded or failed. This makes it easy for the rest of the agent to use the
banner as a clue when classifying services.
"""

import socket
import time
from dataclasses import dataclass
from typing import Optional
from .utility import resolve_local_address


# -----------------------------
# Summary-oriented return type
# -----------------------------
@dataclass
class BannerProbeResult:
    # Pure summary fields
    # status: str            # SUCCESS | FAILURE
    reason: str            # NON_LOCAL | CONNECT_FAILED | READ_FAILED | OK
    success: bool          # True if banner read
    duration_ms: float     # Total time in ms

    # Raw signals
    raw_bytes: bytes
    decoded_text: Optional[str]

    # Embedded low-level diagnostics
    errno: Optional[int]
    error: Optional[str]

    # The resolved address used for the socket
    resolved_address: Optional[str]


# -----------------------------
# Helper: determine socket family
# -----------------------------
def _select_family(address: str):
    """Return (family, sockaddr) for IPv4 or IPv6."""
    try:
        socket.inet_pton(socket.AF_INET, address)
        return socket.AF_INET, (address, None)  # port added later
    except OSError:
        pass

    try:
        socket.inet_pton(socket.AF_INET6, address)
        return socket.AF_INET6, (address, None, 0, 0)
    except OSError:
        pass

    return None, None


# -----------------------------
# Meta-probe implementation
# -----------------------------
def banner_probe(local_address_mode: str, port: int, timeout: float = 2.5) -> BannerProbeResult:
    """
    Banner meta-probe.
    Top-level fields are pure summary.
    Raw socket diagnostics are embedded.
    Windows-safe: uses correct AF_INET/AF_INET6 family.
    """

    start = time.time()

    resolved = resolve_local_address(local_address_mode)
    if not resolved:
        return BannerProbeResult(
            # status="FAILURE",
            reason="NON_LOCAL",
            success=False,
            duration_ms=0.0,
            raw_bytes=b"",
            decoded_text=None,
            errno=None,
            error="NON_LOCAL_ADDRESS",
            resolved_address=None,
        )

    # Determine correct socket family
    family, sockaddr = _select_family(resolved)
    if family is None:
        return BannerProbeResult(
            # status="FAILURE",
            reason="CONNECT_FAILED",
            success=False,
            duration_ms=0.0,
            raw_bytes=b"",
            decoded_text=None,
            errno=None,
            error="INVALID_ADDRESS_FAMILY",
            resolved_address=resolved,
        )

    # Insert port into sockaddr
    if family == socket.AF_INET:
        sockaddr = (resolved, port)
    else:
        sockaddr = (resolved, port, 0, 0)

    # Create socket with correct family (fixes Windows AF_UNSPEC bug)
    s = socket.socket(family, socket.SOCK_STREAM)
    s.settimeout(timeout)

    # --- Connect ---
    try:
        err = s.connect_ex(sockaddr)
        if err != 0:
            return BannerProbeResult(
                # status="FAILURE",
                reason="CONNECT_FAILED",
                success=False,
                duration_ms=(time.time() - start) * 1000,
                raw_bytes=b"",
                decoded_text=None,
                errno=err,
                error=f"CONNECT_EX_FAILED_{err}",
                resolved_address=resolved,
            )
    except Exception as e:
        return BannerProbeResult(
            # status="FAILURE",
            reason="CONNECT_FAILED",
            success=False,
            duration_ms=(time.time() - start) * 1000,
            raw_bytes=b"",
            decoded_text=None,
            errno=-4,
            error=f"UNEXPECTED_EXCEPTION: {e}",
            resolved_address=resolved,
        )

    # --- Read banner ---
    try:
        data = s.recv(4096)
    except Exception as e:
        return BannerProbeResult(
            # status="FAILURE",
            reason="READ_FAILED",
            success=False,
            duration_ms=(time.time() - start) * 1000,
            raw_bytes=b"",
            decoded_text=None,
            errno=-5,
            error=f"READ_EXCEPTION: {e}",
            resolved_address=resolved,
        )
    finally:
        s.close()

    # Decode banner
    try:
        text = data.decode("utf-8", errors="replace") if data else None
    except Exception:
        text = None

    return BannerProbeResult(
        # status="SUCCESS",
        reason="OK",
        success=True,
        duration_ms=(time.time() - start) * 1000,
        raw_bytes=data,
        decoded_text=text,
        errno=0,
        error=None,
        resolved_address=resolved,
    )


# -----------------------------
# Manual test
# -----------------------------
if __name__ == "__main__":
    print("Running banner_probe(ipv4_loopback, 9999)...")
    result = banner_probe("ipv4_loopback", 9999)

    print("\n=== Banner Probe Result ===")
    # print(f"Status        : {result.status}")
    print(f"Reason        : {result.reason}")
    print(f"Success       : {result.success}")
    print(f"Duration (ms) : {result.duration_ms:.2f}")
    print(f"Resolved      : {result.resolved_address}")
    print(f"Raw Bytes     : {result.raw_bytes}")
    print(f"Decoded Text  : {result.decoded_text}")
    print(f"Errno         : {result.errno}")
    print(f"Error         : {result.error}")
