import socket
import ssl
import urllib.request
import urllib.error
from dataclasses import dataclass, field

CONNECT_TIMEOUT = 5
READ_TIMEOUT = 8


@dataclass
class ProbeResult:
    probe_type: str
    success: bool
    banner: str = ""
    body: str = ""
    headers: dict = field(default_factory=dict)
    error: str = ""


def _tcp_banner(host: str, port: int) -> ProbeResult:
    try:
        sock = socket.create_connection((host, port), timeout=CONNECT_TIMEOUT)
        sock.settimeout(READ_TIMEOUT)
        try:
            data = sock.recv(4096)
            return ProbeResult("tcp_banner", True, banner=data.decode("utf-8", errors="replace"))
        except socket.timeout:
            return ProbeResult("tcp_banner", True, banner="")
        finally:
            sock.close()
    except Exception as e:
        return ProbeResult("tcp_banner", False, error=str(e))


def _http_probe(host: str, port: int, probe: dict, use_ssl: bool) -> ProbeResult:
    scheme = "https" if use_ssl else "http"
    path = probe.get("path", "/")
    method = probe.get("method", "GET").upper()
    extra_headers = probe.get("headers", {})
    body_data = probe.get("body", "")

    url = f"{scheme}://{host}:{port}{path}"
    data = body_data.encode() if body_data else None

    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("User-Agent", "Mozilla/5.0")
    req.add_header("Connection", "close")
    for k, v in extra_headers.items():
        req.add_header(k, v)

    try:
        handler = urllib.request.HTTPSHandler(context=ctx) if use_ssl else urllib.request.HTTPHandler()
        opener = urllib.request.build_opener(handler)
        opener.addheaders = []
        with opener.open(req, timeout=READ_TIMEOUT) as resp:
            raw = resp.read(65536)
            body = raw.decode("utf-8", errors="replace")
            headers = dict(resp.headers)
            return ProbeResult(
                "https" if use_ssl else "http",
                True,
                body=body,
                headers=headers,
            )
    except urllib.error.HTTPError as e:
        try:
            body = e.read(16384).decode("utf-8", errors="replace")
        except Exception:
            body = ""
        return ProbeResult("https" if use_ssl else "http", True, body=body, headers=dict(e.headers))
    except Exception as e:
        return ProbeResult("https" if use_ssl else "http", False, error=str(e))


def _rtsp_probe(host: str, port: int) -> ProbeResult:
    try:
        sock = socket.create_connection((host, port), timeout=CONNECT_TIMEOUT)
        sock.settimeout(READ_TIMEOUT)
        request = f"OPTIONS * RTSP/1.0\r\nCSeq: 1\r\nUser-Agent: geo-scout\r\n\r\n"
        sock.sendall(request.encode())
        try:
            data = sock.recv(4096)
            return ProbeResult("rtsp", True, banner=data.decode("utf-8", errors="replace"))
        except socket.timeout:
            return ProbeResult("rtsp", False, error="timeout")
        finally:
            sock.close()
    except Exception as e:
        return ProbeResult("rtsp", False, error=str(e))


def _udp_probe(host: str, port: int) -> ProbeResult:
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(CONNECT_TIMEOUT)
        sock.sendto(b"", (host, port))
        try:
            data, _ = sock.recvfrom(4096)
            return ProbeResult("udp", True, banner=data.decode("utf-8", errors="replace"))
        except socket.timeout:
            return ProbeResult("udp", False, error="timeout")
        finally:
            sock.close()
    except Exception as e:
        return ProbeResult("udp", False, error=str(e))


def run_probe(host: str, port: int, probe: dict) -> ProbeResult:
    ptype = probe.get("type", "tcp_banner")
    if ptype == "tcp_banner":
        return _tcp_banner(host, port)
    elif ptype == "http":
        return _http_probe(host, port, probe, use_ssl=False)
    elif ptype == "https":
        return _http_probe(host, port, probe, use_ssl=True)
    elif ptype == "rtsp":
        return _rtsp_probe(host, port)
    elif ptype == "udp":
        return _udp_probe(host, port)
    return ProbeResult(ptype, False, error=f"unknown probe type: {ptype}")