package socks5

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

// AssociateUDP opens a TCP control conn to the upstream SOCKS5 proxy,
// runs greeting + (optional) auth + UDP ASSOCIATE (CMD=03), and returns:
//
//   - the relay UDP endpoint (host:port the proxy bound for our datagrams)
//   - the kept-open control TCP (caller MUST keep open for the lifetime
//     of the UDP association — closing it tears down the relay on the
//     proxy side per RFC 1928 §6).
//
// The given ctx bounds dial + handshake; once AssociateUDP returns,
// ctrl has its deadline cleared.
//
// If the proxy replies BND.ADDR == 0.0.0.0 (some implementations do
// this to mean "use the same IP you connected to"), we substitute the
// proxy host's resolved IP.
func AssociateUDP(ctx context.Context, cfg Config) (relay *net.UDPAddr, ctrl net.Conn, err error) {
	d := net.Dialer{}
	conn, err := d.DialContext(ctx, "tcp", cfg.ProxyAddr)
	if err != nil {
		return nil, nil, fmt.Errorf("dial proxy: %w", err)
	}
	if dl, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(dl)
	}
	defer func() {
		if err != nil {
			conn.Close()
		}
	}()

	// Greeting (same as TCP CONNECT path)
	if cfg.UseAuth {
		if _, werr := conn.Write([]byte{0x05, 0x02, 0x00, 0x02}); werr != nil {
			return nil, nil, fmt.Errorf("greet write: %w", werr)
		}
	} else {
		if _, werr := conn.Write([]byte{0x05, 0x01, 0x00}); werr != nil {
			return nil, nil, fmt.Errorf("greet write: %w", werr)
		}
	}
	var rep [2]byte
	if _, rerr := io.ReadFull(conn, rep[:]); rerr != nil {
		return nil, nil, fmt.Errorf("greet read: %w", rerr)
	}
	if rep[0] != 0x05 {
		return nil, nil, fmt.Errorf("greet: server version %#x is not SOCKS5", rep[0])
	}
	if rep[1] == 0xff {
		return nil, nil, errors.New("greet: proxy rejected all offered auth methods")
	}
	method := rep[1]

	if method == 0x02 {
		if !cfg.UseAuth {
			return nil, nil, errors.New("proxy requires auth but Config.UseAuth is false")
		}
		if len(cfg.Login) > 255 || len(cfg.Password) > 255 {
			return nil, nil, errors.New("login or password too long")
		}
		buf := make([]byte, 0, 3+len(cfg.Login)+len(cfg.Password))
		buf = append(buf, 0x01, byte(len(cfg.Login)))
		buf = append(buf, []byte(cfg.Login)...)
		buf = append(buf, byte(len(cfg.Password)))
		buf = append(buf, []byte(cfg.Password)...)
		if _, werr := conn.Write(buf); werr != nil {
			return nil, nil, fmt.Errorf("auth write: %w", werr)
		}
		if _, rerr := io.ReadFull(conn, rep[:]); rerr != nil {
			return nil, nil, fmt.Errorf("auth read: %w", rerr)
		}
		if rep[1] != 0x00 {
			return nil, nil, errors.New("auth: invalid login or password")
		}
	}

	// UDP ASSOCIATE request: 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, fmt.Errorf("udp-associate write: %w", werr)
	}

	// We accept ATYP=01 (IPv4) replies only — sufficient for our use
	// case (mihomo + standard proxies). Reading 10 bytes covers exactly
	// that case: VER REP RSV ATYP BND.ADDR(4) BND.PORT(2).
	reply := make([]byte, 10)
	if _, rerr := io.ReadFull(conn, reply); rerr != nil {
		return nil, nil, fmt.Errorf("udp-associate read: %w", rerr)
	}
	if reply[0] != 0x05 {
		return nil, nil, fmt.Errorf("udp-associate: server version %#x is not SOCKS5", reply[0])
	}
	if reply[1] != 0x00 {
		return nil, nil, fmt.Errorf("udp-associate: REP=%#02x", reply[1])
	}
	if reply[3] != 0x01 {
		return nil, nil, fmt.Errorf("udp-associate: unsupported BND.ATYP=%#02x (need IPv4)", reply[3])
	}

	bndIP := net.IPv4(reply[4], reply[5], reply[6], reply[7]).To4()
	bndPort := binary.BigEndian.Uint16(reply[8:10])

	// Per RFC 1928 §6 / common practice: BND.ADDR=0.0.0.0 means "use
	// the same address you used to reach me". Substitute proxy host's
	// IP from the established TCP conn's RemoteAddr.
	if bndIP.Equal(net.IPv4zero.To4()) {
		if ra, ok := conn.RemoteAddr().(*net.TCPAddr); ok && ra.IP != nil {
			if v4 := ra.IP.To4(); v4 != nil {
				bndIP = v4
			}
		}
	}

	// Clear deadline so caller can use ctrl as-is (keepalive only).
	_ = conn.SetDeadline(time.Time{})

	return &net.UDPAddr{IP: bndIP, Port: int(bndPort)}, conn, nil
}

// EncapUDPv4 wraps an outbound UDP payload in the SOCKS5 UDP datagram
// envelope (RFC 1928 §7) for ATYP=01 (IPv4). The returned buffer has
// the form:
//
//	RSV(2)=0000 | FRAG(1)=00 | ATYP(1)=01 | DST.ADDR(4) | DST.PORT(2) | DATA
//
// The 10-byte prefix tells the relay where to forward the datagram.
// Returns an error if dstIP is not IPv4.
func EncapUDPv4(dstIP net.IP, dstPort uint16, payload []byte) ([]byte, error) {
	v4 := dstIP.To4()
	if v4 == nil {
		return nil, errors.New("EncapUDPv4: dst must be IPv4")
	}
	out := make([]byte, 10+len(payload))
	out[0] = 0x00 // RSV
	out[1] = 0x00 // RSV
	out[2] = 0x00 // FRAG (no fragmentation)
	out[3] = 0x01 // ATYP IPv4
	copy(out[4:8], v4)
	binary.BigEndian.PutUint16(out[8:10], dstPort)
	copy(out[10:], payload)
	return out, nil
}

// DecapUDPv4 parses an inbound SOCKS5 UDP datagram (RFC 1928 §7) for
// ATYP=01 (IPv4). On the inbound path the relay puts the ORIGIN's
// addr/port in DST.ADDR/DST.PORT — i.e. for us, the original DST that
// answered (e.g. the Discord voice server). The returned (srcIP,
// srcPort) reflect that origin; payload is the original UDP body.
//
// Errors when:
//   - buf shorter than 10 bytes (truncated header)
//   - FRAG != 0 (we don't reassemble fragments)
//   - ATYP != 1 (we only handle IPv4 in this path)
func DecapUDPv4(buf []byte) (srcIP net.IP, srcPort uint16, payload []byte, err error) {
	if len(buf) < 10 {
		return nil, 0, nil, errors.New("DecapUDPv4: truncated header")
	}
	if buf[2] != 0x00 {
		return nil, 0, nil, fmt.Errorf("DecapUDPv4: FRAG=%d not supported", buf[2])
	}
	if buf[3] != 0x01 {
		return nil, 0, nil, fmt.Errorf("DecapUDPv4: ATYP=%#02x not IPv4", buf[3])
	}
	srcIP = net.IPv4(buf[4], buf[5], buf[6], buf[7])
	srcPort = binary.BigEndian.Uint16(buf[8:10])
	payload = buf[10:]
	return srcIP, srcPort, payload, nil
}