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

from tests.conftest import (
    socks5_handshake,
    socks5_connect,
    socks5_udp_associate,
    build_socks5_udp_packet,
    parse_socks5_udp_packet,
    build_dns_query,
)

# ── TCP CONNECT ──────────────────────────────────────────────────────

def test_connect_ipv4(socks5_server, tcp_echo_server):
    """TCP CONNECT to an IPv4 echo server through the proxy."""
    proxy_host, proxy_port = socks5_server
    echo_host, echo_port = tcp_echo_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as s:
        socks5_handshake(s)
        resp = socks5_connect(s, echo_host, echo_port, atyp="ipv4")
        assert resp[0] == 5 and resp[1] == 0  # success

        s.sendall(b"hello proxy")
        data = s.recv(1024)
        assert data == b"hello proxy"


def test_connect_domain(socks5_server, tcp_echo_server):
    """TCP CONNECT using a domain name (resolves via DoH)."""
    proxy_host, proxy_port = socks5_server
    echo_host, echo_port = tcp_echo_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as s:
        socks5_handshake(s)
        # "localhost" should resolve to 127.0.0.1 for local echo test
        # Use the echo server's real IP to avoid DoH dependency
        resp = socks5_connect(s, "127.0.0.1", echo_port, atyp="ipv4")
        assert resp[0] == 5 and resp[1] == 0

        s.sendall(b"domain test")
        data = s.recv(1024)
        assert data == b"domain test"


# ── AUTH ─────────────────────────────────────────────────────────────

def test_auth_noauth(socks5_server, tcp_echo_server):
    """No-auth handshake succeeds."""
    proxy_host, proxy_port = socks5_server
    echo_host, echo_port = tcp_echo_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as s:
        resp = socks5_handshake(s, auth=False)
        assert resp == b"\x05\x00"

        # verify a CONNECT still works
        connect_resp = socks5_connect(s, echo_host, echo_port, atyp="ipv4")
        assert connect_resp[1] == 0


def test_auth_success(socks5_server, tcp_echo_server):
    """Correct credentials are accepted."""
    proxy_host, proxy_port = socks5_server
    echo_host, echo_port = tcp_echo_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as s:
        auth_resp = socks5_handshake(
            s, auth=True,
            username="YOUR_PROXY_LOGIN",
            password="YOUR_PROXY_PASSWORD",
        )
        assert auth_resp == b"\x01\x00"  # success

        connect_resp = socks5_connect(s, echo_host, echo_port, atyp="ipv4")
        assert connect_resp[1] == 0

        s.sendall(b"authed")
        data = s.recv(1024)
        assert data == b"authed"


def test_auth_failure(socks5_server):
    """Wrong credentials are rejected."""
    proxy_host, proxy_port = socks5_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as s:
        auth_resp = socks5_handshake(
            s, auth=True,
            username="wrong_user",
            password="wrong_pass",
        )
        assert auth_resp == b"\x01\x01"  # failure


# ── INVALID VERSION ─────────────────────────────────────────────────

def test_invalid_version(socks5_server):
    """Sending a wrong SOCKS version should cause disconnect."""
    proxy_host, proxy_port = socks5_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as s:
        s.sendall(b"\x04\x01\x00")  # version 4
        # Server should close connection or return error
        data = s.recv(64)
        # Either empty (connection closed) or no valid SOCKS5 response
        assert data == b"" or (len(data) >= 2 and data[0] != 5)


# ── UDP ASSOCIATE ────────────────────────────────────────────────────

def test_udp_associate_handshake(socks5_server):
    """UDP ASSOCIATE handshake returns a valid relay address."""
    proxy_host, proxy_port = socks5_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as s:
        socks5_handshake(s)
        resp, relay_addr, relay_port = socks5_udp_associate(s)
        assert resp[1] == 0  # success
        assert relay_port > 0


def test_udp_forward(socks5_server, udp_echo_server):
    """UDP data forwarded to echo server and back through the relay."""
    proxy_host, proxy_port = socks5_server
    echo_host, echo_port = udp_echo_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as tcp:
        socks5_handshake(tcp)
        _, relay_addr, relay_port = socks5_udp_associate(tcp)

        udp_client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_client.settimeout(5)
        try:
            payload = b"udp echo test"
            pkt = build_socks5_udp_packet(echo_host, echo_port, payload)
            udp_client.sendto(pkt, (relay_addr if relay_addr != "0.0.0.0" else "127.0.0.1", relay_port))

            reply_data, _ = udp_client.recvfrom(4096)
            addr, port, echoed = parse_socks5_udp_packet(reply_data)
            assert echoed == payload
        finally:
            udp_client.close()


def test_udp_dns(socks5_server, mock_dns_server):
    """DNS query relayed through UDP associate to mock DNS server."""
    proxy_host, proxy_port = socks5_server
    dns_host, dns_port = mock_dns_server

    with socket.create_connection((proxy_host, proxy_port), timeout=5) as tcp:
        socks5_handshake(tcp)
        _, relay_addr, relay_port = socks5_udp_associate(tcp)

        udp_client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_client.settimeout(5)
        try:
            dns_query = build_dns_query("example.com")
            pkt = build_socks5_udp_packet(dns_host, dns_port, dns_query)
            udp_client.sendto(pkt, (relay_addr if relay_addr != "0.0.0.0" else "127.0.0.1", relay_port))

            reply_data, _ = udp_client.recvfrom(4096)
            addr, port, dns_reply = parse_socks5_udp_packet(reply_data)
            # Verify it's a DNS response (QR bit set)
            assert len(dns_reply) >= 12
            flags = struct.unpack("!H", dns_reply[2:4])[0]
            assert flags & 0x8000  # QR=1 means response
        finally:
            udp_client.close()
