package checker

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

// Sentinel errors returned by the SOCKS5 primitives.
var (
	ErrSocks5BadVersion      = errors.New("socks5: server returned wrong version")
	ErrSocks5RejectedAllAuth = errors.New("socks5: server rejected all offered auth methods (0xFF)")
	ErrAuthRejected          = errors.New("socks5: user/pass authentication rejected")
	ErrCredentialTooLong     = errors.New("socks5: login or password longer than 255 bytes")
	ErrHostTooLong           = errors.New("socks5: target hostname longer than 255 bytes")
	ErrUnsupportedRelayATYP  = errors.New("socks5: udp associate replied with non-IPv4 ATYP")
	ErrShortReply            = errors.New("socks5: short server reply")
)

// ErrSocks5Reply wraps a non-zero REP code so callers can react to specific
// SOCKS5 reply codes (e.g. REP=0x07 = command not supported, REP=0x05 =
// connection refused).
type ErrSocks5Reply struct{ Code byte }

// Error implements the error interface.
func (e ErrSocks5Reply) Error() string {
	return fmt.Sprintf("socks5: server replied with non-zero REP code 0x%02X", e.Code)
}

// Is reports whether target matches this reply error by Code.
func (e ErrSocks5Reply) Is(target error) bool {
	t, ok := target.(ErrSocks5Reply)
	if !ok {
		return false
	}
	return t.Code == e.Code
}

// applyDeadline applies the deadline from ctx (if any) to conn. Returns a
// function to clear the deadline.
func applyDeadline(ctx context.Context, conn net.Conn) {
	if dl, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(dl)
	} else {
		_ = conn.SetDeadline(time.Time{})
	}
}

// joinCtxErr wraps err with ctx.Err() if ctx has been cancelled or expired,
// so that callers see context.Canceled / context.DeadlineExceeded in the
// error chain even when the underlying I/O reported a deadline-based error.
func joinCtxErr(ctx context.Context, err error) error {
	if err == nil {
		return nil
	}
	if cerr := ctx.Err(); cerr != nil {
		return errors.Join(err, cerr)
	}
	return err
}

// socks5Greeting performs the RFC 1928 client greeting on conn.
// useAuth=true sends "05 02 00 02" (offer no-auth and user/pass);
// useAuth=false sends "05 01 00" (offer no-auth only).
func socks5Greeting(ctx context.Context, conn net.Conn, useAuth bool) (method byte, rawReply []byte, err error) {
	applyDeadline(ctx, conn)

	var greet []byte
	if useAuth {
		greet = []byte{0x05, 0x02, 0x00, 0x02}
	} else {
		greet = []byte{0x05, 0x01, 0x00}
	}

	if _, werr := conn.Write(greet); werr != nil {
		return 0, nil, joinCtxErr(ctx, fmt.Errorf("socks5 greeting: write: %w", werr))
	}

	reply := make([]byte, 2)
	n, rerr := io.ReadFull(conn, reply)
	if rerr != nil {
		partial := reply[:n]
		if errors.Is(rerr, io.ErrUnexpectedEOF) || errors.Is(rerr, io.EOF) {
			return 0, partial, joinCtxErr(ctx, fmt.Errorf("socks5 greeting: %w (raw=%x)", ErrShortReply, partial))
		}
		return 0, partial, joinCtxErr(ctx, fmt.Errorf("socks5 greeting: read: %w (raw=%x)", rerr, partial))
	}

	if reply[0] != 0x05 {
		return 0, reply, fmt.Errorf("socks5 greeting: %w (raw=%x)", ErrSocks5BadVersion, reply)
	}
	if reply[1] == 0xFF {
		return reply[1], reply, fmt.Errorf("socks5 greeting: %w (raw=%x)", ErrSocks5RejectedAllAuth, reply)
	}
	return reply[1], reply, nil
}

// socks5Auth performs RFC 1929 user/pass sub-negotiation on conn,
// after greeting selected method 0x02.
func socks5Auth(ctx context.Context, conn net.Conn, login, password string) (rawReply []byte, err error) {
	if len(login) > 255 || len(password) > 255 {
		return nil, ErrCredentialTooLong
	}

	applyDeadline(ctx, conn)

	buf := make([]byte, 0, 3+len(login)+len(password))
	buf = append(buf, 0x01)              // VER
	buf = append(buf, byte(len(login)))  // ULEN
	buf = append(buf, []byte(login)...)  // UNAME
	buf = append(buf, byte(len(password)))
	buf = append(buf, []byte(password)...)

	if _, werr := conn.Write(buf); werr != nil {
		return nil, joinCtxErr(ctx, fmt.Errorf("socks5 auth: write: %w", werr))
	}

	reply := make([]byte, 2)
	n, rerr := io.ReadFull(conn, reply)
	if rerr != nil {
		partial := reply[:n]
		if errors.Is(rerr, io.ErrUnexpectedEOF) || errors.Is(rerr, io.EOF) {
			return partial, joinCtxErr(ctx, fmt.Errorf("socks5 auth: %w (raw=%x)", ErrShortReply, partial))
		}
		return partial, joinCtxErr(ctx, fmt.Errorf("socks5 auth: read: %w (raw=%x)", rerr, partial))
	}

	if reply[0] != 0x01 {
		return reply, fmt.Errorf("socks5 auth: auth subneg version mismatch: got 0x%02X want 0x01 (raw=%x)", reply[0], reply)
	}
	if reply[1] != 0x00 {
		return reply, fmt.Errorf("socks5 auth: %w (raw=%x)", ErrAuthRejected, reply)
	}
	return reply, nil
}

// socks5Connect performs SOCKS5 CONNECT (CMD=01) to host:port using
// ATYP=03 (domain name).
func socks5Connect(ctx context.Context, conn net.Conn, host string, port uint16) (rawReply []byte, err error) {
	if len(host) > 255 {
		return nil, ErrHostTooLong
	}

	applyDeadline(ctx, conn)

	// VER=05 CMD=01 RSV=00 ATYP=03 LEN host port
	req := make([]byte, 0, 7+len(host))
	req = append(req, 0x05, 0x01, 0x00, 0x03)
	req = append(req, byte(len(host)))
	req = append(req, []byte(host)...)
	var portBuf [2]byte
	binary.BigEndian.PutUint16(portBuf[:], port)
	req = append(req, portBuf[:]...)

	if _, werr := conn.Write(req); werr != nil {
		return nil, joinCtxErr(ctx, fmt.Errorf("socks5 connect: write: %w", werr))
	}

	// We always read 10 bytes (assuming ATYP=01 IPv4 reply, the most
	// common case from real proxies). Parsing variable-length BND is
	// out of scope for the diagnostic.
	reply := make([]byte, 10)
	n, rerr := io.ReadFull(conn, reply)
	if rerr != nil {
		partial := reply[:n]
		if errors.Is(rerr, io.ErrUnexpectedEOF) || errors.Is(rerr, io.EOF) {
			return partial, joinCtxErr(ctx, fmt.Errorf("socks5 connect: %w (raw=%x)", ErrShortReply, partial))
		}
		return partial, joinCtxErr(ctx, fmt.Errorf("socks5 connect: read: %w (raw=%x)", rerr, partial))
	}

	if reply[0] != 0x05 {
		return reply, fmt.Errorf("socks5 connect: %w (raw=%x)", ErrSocks5BadVersion, reply)
	}
	if reply[1] != 0x00 {
		return reply, fmt.Errorf("socks5 connect: %w (raw=%x)", ErrSocks5Reply{Code: reply[1]}, reply)
	}
	return reply, nil
}

// socks5UDPAssociate performs SOCKS5 UDP ASSOCIATE (CMD=03) on conn.
func socks5UDPAssociate(ctx context.Context, conn net.Conn) (relay *net.UDPAddr, rawReply []byte, err error) {
	applyDeadline(ctx, conn)

	// VER=05 CMD=03 RSV=00 ATYP=01 DST.ADDR=0.0.0.0 DST.PORT=0
	req := []byte{0x05, 0x03, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	if _, werr := conn.Write(req); werr != nil {
		return nil, nil, joinCtxErr(ctx, fmt.Errorf("socks5 udp-associate: write: %w", werr))
	}

	reply := make([]byte, 10)
	n, rerr := io.ReadFull(conn, reply)
	if rerr != nil {
		partial := reply[:n]
		if errors.Is(rerr, io.ErrUnexpectedEOF) || errors.Is(rerr, io.EOF) {
			return nil, partial, joinCtxErr(ctx, fmt.Errorf("socks5 udp-associate: %w (raw=%x)", ErrShortReply, partial))
		}
		return nil, partial, joinCtxErr(ctx, fmt.Errorf("socks5 udp-associate: read: %w (raw=%x)", rerr, partial))
	}

	if reply[0] != 0x05 {
		return nil, reply, fmt.Errorf("socks5 udp-associate: %w (raw=%x)", ErrSocks5BadVersion, reply)
	}
	if reply[1] != 0x00 {
		return nil, reply, fmt.Errorf("socks5 udp-associate: %w (raw=%x)", ErrSocks5Reply{Code: reply[1]}, reply)
	}
	if reply[3] != 0x01 {
		return nil, reply, fmt.Errorf("socks5 udp-associate: %w (atyp=0x%02X raw=%x)", ErrUnsupportedRelayATYP, reply[3], reply)
	}

	ip := net.IPv4(reply[4], reply[5], reply[6], reply[7])
	port := binary.BigEndian.Uint16(reply[8:10])
	relay = &net.UDPAddr{IP: ip, Port: int(port)}
	return relay, reply, nil
}