# -*- coding: utf-8 -*-
import socket
import struct
import threading
import pytest

from pymrf4.socks5_proxy import ThreadingTCPServer, Socks5Handler


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def _free_port():
    """Return an available TCP port on localhost."""
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def socks5_handshake(sock, *, auth=False, username=None, password=None):
    """Perform the SOCKS5 greeting + optional user/pass auth.

    Returns the server-chosen auth method byte.
    """
    if auth:
        sock.sendall(b"\x05\x01\x02")          # VER, 1 method, USER/PASS
    else:
        sock.sendall(b"\x05\x01\x00")          # VER, 1 method, NO AUTH

    resp = sock.recv(2)
    assert resp[0] == 0x05

    if resp[1] == 0x02 and username is not None:
        # sub-negotiation: VER=0x01, ulen, uname, plen, passwd
        uname = username.encode() if isinstance(username, str) else username
        passwd = password.encode() if isinstance(password, str) else password
        sock.sendall(
            b"\x01"
            + bytes([len(uname)]) + uname
            + bytes([len(passwd)]) + passwd
        )
        auth_resp = sock.recv(2)
        return auth_resp  # (ver, status) – 0x00 = success

    return resp


def socks5_connect(sock, addr, port, *, atyp="ipv4"):
    """Send a CONNECT request and return the reply bytes."""
    if atyp == "ipv4":
        req = struct.pack("!BBB B", 5, 1, 0, 1)
        req += socket.inet_aton(addr) + struct.pack("!H", port)
    elif atyp == "domain":
        encoded = addr.encode() if isinstance(addr, str) else addr
        req = struct.pack("!BBB B B", 5, 1, 0, 3, len(encoded))
        req += encoded + struct.pack("!H", port)
    else:
        raise ValueError(f"unsupported atyp: {atyp}")

    sock.sendall(req)
    resp = sock.recv(64)
    return resp


def socks5_udp_associate(sock):
    """Send a UDP ASSOCIATE request.

    Returns (reply_bytes, relay_addr, relay_port).
    The caller is expected to have already done the greeting/auth handshake.
    """
    # DST.ADDR = 0.0.0.0:0  (ask the server to pick)
    req = struct.pack("!BBB B", 5, 3, 0, 1)
    req += socket.inet_aton("0.0.0.0") + struct.pack("!H", 0)
    sock.sendall(req)

    resp = sock.recv(64)
    # parse reply: VER REP RSV ATYP BND.ADDR BND.PORT
    assert resp[0] == 5 and resp[1] == 0  # success
    atyp = resp[3]
    if atyp == 1:  # IPv4
        relay_addr = socket.inet_ntoa(resp[4:8])
        relay_port = struct.unpack("!H", resp[8:10])[0]
    else:
        raise RuntimeError("unexpected atyp in UDP ASSOCIATE reply")

    return resp, relay_addr, relay_port


def build_socks5_udp_packet(addr, port, payload, *, atyp="ipv4"):
    """Build a SOCKS5 UDP request header + payload."""
    # RSV (2 bytes) | FRAG (1 byte) | ATYP (1 byte) | DST.ADDR | DST.PORT | DATA
    header = struct.pack("!HBB", 0, 0, 1 if atyp == "ipv4" else 3)
    if atyp == "ipv4":
        header += socket.inet_aton(addr) + struct.pack("!H", port)
    elif atyp == "domain":
        encoded = addr.encode() if isinstance(addr, str) else addr
        header += bytes([len(encoded)]) + encoded + struct.pack("!H", port)
    return header + payload


def parse_socks5_udp_packet(data):
    """Parse a SOCKS5 UDP reply: returns (addr, port, payload)."""
    rsv, frag, atyp = struct.unpack("!HBB", data[:4])
    if atyp == 1:  # IPv4
        addr = socket.inet_ntoa(data[4:8])
        port = struct.unpack("!H", data[8:10])[0]
        payload = data[10:]
    elif atyp == 3:  # domain
        alen = data[4]
        addr = data[5:5 + alen].decode()
        port = struct.unpack("!H", data[5 + alen:5 + alen + 2])[0]
        payload = data[5 + alen + 2:]
    else:
        raise ValueError(f"unsupported atyp {atyp}")
    return addr, port, payload


def build_dns_query(domain, qtype=1, txn_id=0x1234):
    """Build a minimal DNS query packet (single A-record question)."""
    # Header: ID, flags=0x0100 (standard query), QDCOUNT=1
    header = struct.pack("!HHHHHH", txn_id, 0x0100, 1, 0, 0, 0)
    # Question section
    question = b""
    for label in domain.split("."):
        question += bytes([len(label)]) + label.encode()
    question += b"\x00"  # root label
    question += struct.pack("!HH", qtype, 1)  # QTYPE=A, QCLASS=IN
    return header + question


# ---------------------------------------------------------------------------
# fixtures
# ---------------------------------------------------------------------------

@pytest.fixture()
def socks5_server():
    """Start a SOCKS5 proxy on a random port; yield (host, port)."""
    port = _free_port()
    server = ThreadingTCPServer(("127.0.0.1", port), Socks5Handler)
    t = threading.Thread(target=server.serve_forever, daemon=True)
    t.start()
    yield "127.0.0.1", port
    server.shutdown()


@pytest.fixture()
def tcp_echo_server():
    """Simple TCP echo server; yield (host, port)."""
    port = _free_port()
    srv_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv_sock.bind(("127.0.0.1", port))
    srv_sock.listen(5)

    def _serve():
        while True:
            try:
                conn, _ = srv_sock.accept()
                data = conn.recv(4096)
                if data:
                    conn.sendall(data)
                conn.close()
            except OSError:
                break

    t = threading.Thread(target=_serve, daemon=True)
    t.start()
    yield "127.0.0.1", port
    srv_sock.close()


@pytest.fixture()
def udp_echo_server():
    """Simple UDP echo server; yield (host, port)."""
    port = _free_port()
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(("127.0.0.1", port))

    def _serve():
        while True:
            try:
                data, addr = sock.recvfrom(4096)
                if data:
                    sock.sendto(data, addr)
            except OSError:
                break

    t = threading.Thread(target=_serve, daemon=True)
    t.start()
    yield "127.0.0.1", port
    sock.close()


@pytest.fixture()
def mock_dns_server():
    """UDP DNS server that answers every A-record query with 127.0.0.1.

    Yield (host, port).
    """
    port = _free_port()
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(("127.0.0.1", port))

    def _serve():
        while True:
            try:
                data, addr = sock.recvfrom(512)
                if len(data) < 12:
                    continue
                # Build a minimal response: copy ID, set QR bit, copy question
                txn_id = data[:2]
                flags = struct.pack("!H", 0x8180)  # QR=1, AA=1, RD=1, RA=1
                counts = struct.pack("!HHHH", 1, 1, 0, 0)  # QD=1, AN=1
                # question section – copy from query (after 12-byte header)
                question = data[12:]
                # answer: NAME=pointer to offset 12, TYPE=A, CLASS=IN, TTL=60, RDLEN=4, RDATA=127.0.0.1
                answer = b"\xc0\x0c"  # pointer to question name
                answer += struct.pack("!HHIH", 1, 1, 60, 4)
                answer += socket.inet_aton("127.0.0.1")
                reply = txn_id + flags + counts + question + answer
                sock.sendto(reply, addr)
            except OSError:
                break

    t = threading.Thread(target=_serve, daemon=True)
    t.start()
    yield "127.0.0.1", port
    sock.close()
