"""Scanner engine — faithful port of VlessScanner.kt (MIT, Unknown37841).

Original: app/src/main/java/com/spider/railwayipscanner/VlessScanner.kt (219 lines)

Protocol (unchanged from the Android original):
  1. TCP connect to IP:port (default 443), timeout 2400 ms
  2. TLS handshake with custom SNI = domain, ALPN http/1.1, trust-all certs
  3. WebSocket upgrade GET <path>, expect HTTP 101, headers read strictly
     until \\r\\n\\r\\n
  4. Binary VLESS header + HTTP GET /generate_204 payload, sent as a masked
     WebSocket binary frame
  5. Success only if the response contains "HTTP/1.1 204" or "204 No Content"

Verification (test_ip_real_delay): 3 consecutive probes, 60 ms sleep between,
average delay; a drop on the 2nd/3rd probe is reported as "Unstable".

Stdlib only: socket, ssl, threading / concurrent.futures, time, uuid.
"""
from __future__ import annotations

import socket
import ssl
import time
import uuid as uuid_module
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Callable, Optional

DEFAULT_TIMEOUT_MS = 2400
WS_MASK = bytes([0x12, 0x34, 0x56, 0x78])
TARGET_HOST = b"www.google.com"
TARGET_PORT = 80


@dataclass
class ScanResult:
    ip: str
    is_success: bool
    delay_ms: int = 0
    error: Optional[str] = None


# ---------------------------------------------------------------------------
# Packet / frame construction (separated so it can be unit-tested offline)
# ---------------------------------------------------------------------------

def build_vless_header(uuid_string: str, target_host: bytes = TARGET_HOST,
                       target_port: int = TARGET_PORT) -> bytes:
    """Binary VLESS header, byte-for-byte as in VlessScanner.kt:

    [0] version 0 | [1:17] UUID (16 bytes) | addon length 0 | command 1 (TCP)
    | port (2 bytes, big-endian) | address type 2 (domain) | host length | host
    """
    uuid_bytes = uuid_module.UUID(uuid_string).bytes  # 16 bytes, MSB-first
    header = bytearray()
    header.append(0)                       # Version 0
    header.extend(uuid_bytes)              # UUID auth
    header.append(0)                       # Proto addon length
    header.append(1)                       # Command: 1 = TCP
    header.append((target_port >> 8) & 0xFF)
    header.append(target_port & 0xFF)      # Port 80 in the original
    header.append(2)                       # Address Type: Domain
    header.append(len(target_host) & 0xFF)
    header.extend(target_host)
    return bytes(header)


def build_http_payload() -> bytes:
    crlf = "\r\n"
    return (
        "GET /generate_204 HTTP/1.1" + crlf +
        "Host: www.google.com" + crlf +
        "User-Agent: v2rayN" + crlf +
        "Connection: close" + crlf + crlf
    ).encode("utf-8")


def build_full_vless_data(uuid_string: str) -> bytes:
    return build_vless_header(uuid_string) + build_http_payload()


def build_ws_frame(data: bytes, mask: bytes = WS_MASK) -> bytes:
    """Masked WebSocket client binary frame, as in VlessScanner.kt.

    First byte 0x82 (FIN + binary), second byte 0x80|len (<126) or
    0x80|126 + 2-byte length, then the 4-byte mask, then masked payload.
    """
    masked = bytes(b ^ mask[i % 4] for i, b in enumerate(data))
    if len(data) < 126:
        frame_header = bytes([0x82, 0x80 | len(data)]) + mask
    else:
        frame_header = bytes([
            0x82,
            0x80 | 126,
            (len(data) >> 8) & 0xFF,
            len(data) & 0xFF,
        ]) + mask
    return frame_header + masked


def build_ws_upgrade_request(domain: str, path: str) -> bytes:
    clean_path = path if path.startswith("/") else "/" + path
    crlf = "\r\n"
    ws_req = (
        "GET " + clean_path + " HTTP/1.1" + crlf +
        "Host: " + domain + crlf +
        "Upgrade: websocket" + crlf +
        "Connection: Upgrade" + crlf +
        "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==" + crlf +
        "Sec-WebSocket-Version: 13" + crlf +
        "User-Agent: Mozilla/5.0 (Android; SpiderPanel)" + crlf + crlf
    )
    return ws_req.encode("utf-8")


def _read_until_double_crlf(recv_byte: Callable[[], Optional[int]]) -> Optional[bytes]:
    """Read one byte at a time until \\r\\n\\r\\n — mirrors readUntilDoubleCRLF()."""
    buffer = bytearray()
    pattern = bytes([13, 10, 13, 10])
    matched = 0
    max_bytes = 4096
    total_read = 0
    while total_read < max_bytes:
        b = recv_byte()
        if b is None or b == -1:
            break
        total_read += 1
        buffer.append(b)
        if b == pattern[matched]:
            matched += 1
            if matched == len(pattern):
                return bytes(buffer)
        else:
            matched = 1 if b == pattern[0] else 0
    return bytes(buffer) if len(buffer) > 0 else None


# ---------------------------------------------------------------------------
# Probing
# ---------------------------------------------------------------------------

def _make_ssl_context() -> ssl.SSLContext:
    # Trust-all, exactly like the original's X509TrustManager that accepts all.
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    try:
        ctx.set_alpn_protocols(["http/1.1"])
    except Exception:
        pass
    return ctx


_SSL_CONTEXT = None


def _ssl_context() -> ssl.SSLContext:
    global _SSL_CONTEXT
    if _SSL_CONTEXT is None:
        _SSL_CONTEXT = _make_ssl_context()
    return _SSL_CONTEXT


def run_single_probe(ip: str, domain: str, uuid_string: str,
                     path: str = "/", port: int = 443) -> ScanResult:
    t0 = time.time()
    raw_socket = None
    ssl_socket = None
    try:
        timeout_s = DEFAULT_TIMEOUT_MS / 1000.0

        # 1. Direct TCP Connect
        raw_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        raw_socket.settimeout(timeout_s)
        raw_socket.connect((ip, port))

        # 2. TLS Handshake with Domain SNI & ALPN http/1.1
        ssl_socket = _ssl_context().wrap_socket(
            raw_socket, server_hostname=domain
        )
        ssl_socket.settimeout(timeout_s)
        raw_socket = None  # ownership transferred to ssl_socket

        out = ssl_socket
        inp = ssl_socket

        # 3. Send WebSocket Upgrade Request
        out.sendall(build_ws_upgrade_request(domain, path))

        def _recv_byte():
            data = inp.recv(1)
            if not data:
                return -1
            return data[0]

        header_bytes = _read_until_double_crlf(_recv_byte)
        if header_bytes is None:
            return ScanResult(ip, False, 0, "No WS upgrade response")
        header_str = header_bytes.decode("utf-8", errors="replace")
        first_line = header_str.splitlines()[0] if header_str.splitlines() else ""
        if "101" not in first_line:
            return ScanResult(ip, False, 0, f"WS Upgrade failed: {first_line}")

        # 4+5. Binary VLESS frame + HTTP 204 probe payload
        full_data = build_full_vless_data(uuid_string)
        out.sendall(build_ws_frame(full_data))

        # 5. Read End-to-End Google Response (single read, 1024 bytes, as original)
        resp_buf = inp.recv(1024)
        delay_ms = int((time.time() - t0) * 1000)

        if resp_buf:
            resp_str = resp_buf.decode("iso-8859-1")
            if "HTTP/1.1 204" in resp_str or "204 No Content" in resp_str:
                return ScanResult(ip, True, delay_ms)
        return ScanResult(ip, False, delay_ms, "No valid Google 204 response")

    except Exception as e:
        return ScanResult(ip, False, 0, str(e) or "Connection error")
    finally:
        try:
            if ssl_socket is not None:
                ssl_socket.close()
        except Exception:
            pass
        try:
            if raw_socket is not None:
                raw_socket.close()
        except Exception:
            pass


def test_ip_real_delay(ip: str, domain: str, uuid_string: str,
                       path: str = "/", port: int = 443,
                       double_check: bool = True) -> ScanResult:
    """Faithful port of testIpRealDelay(): 3 probes when double_check is on."""
    first_attempt = run_single_probe(ip, domain, uuid_string, path, port)
    if not first_attempt.is_success:
        return first_attempt

    if double_check:
        time.sleep(0.060)
        second_attempt = run_single_probe(ip, domain, uuid_string, path, port)
        if not second_attempt.is_success:
            return ScanResult(ip, False, 0, "Unstable connection (packet drop on 2nd probe)")
        time.sleep(0.060)
        third_attempt = run_single_probe(ip, domain, uuid_string, path, port)
        if not third_attempt.is_success:
            return ScanResult(ip, False, 0, "Unstable connection (packet drop on 3rd probe)")
        avg_delay = (first_attempt.delay_ms + second_attempt.delay_ms + third_attempt.delay_ms) // 3
        return ScanResult(ip, True, avg_delay)

    return first_attempt


def scan_subnet(subnet_prefix: str, domain: str, uuid_string: str,
                path: str = "/", port: int = 443, max_workers: int = 12,
                on_progress: Optional[Callable[[int, int, ScanResult], None]] = None,
                should_stop: Optional[Callable[[], bool]] = None) -> list[ScanResult]:
    """Scan .1-.254 of a subnet prefix in parallel, sorted by delay.

    Android MainActivity uses a fixed pool of 12 threads and launches one
    job per IP; results are kept sorted by delayMs as they arrive.
    """
    prefix = subnet_prefix if subnet_prefix.endswith(".") else subnet_prefix + "."
    all_ips = [f"{prefix}{i}" for i in range(1, 255)]
    results: list[ScanResult] = []
    completed = 0

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_ip = {
            executor.submit(test_ip_real_delay, ip, domain, uuid_string, path, port, True): ip
            for ip in all_ips
        }
        for future in as_completed(future_to_ip):
            if should_stop is not None and should_stop():
                for f in future_to_ip:
                    f.cancel()
                break
            completed += 1
            try:
                result = future.result()
            except Exception as e:
                result = ScanResult(future_to_ip[future], False, 0, str(e))
            if result.is_success:
                results.append(result)
                results.sort(key=lambda r: r.delay_ms)
            if on_progress is not None:
                on_progress(completed, 254, result)

    results.sort(key=lambda r: r.delay_ms)
    return results