package probe

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

const (
	socks5Version        = 0x05
	socks5NoAuth         = 0x00
	socks5UserPassword   = 0x02
	socks5NoAcceptable   = 0xff
	socks5ConnectCommand = 0x01
	socks5AddressIPv4    = 0x01
	socks5AddressDomain  = 0x03
	socks5AddressIPv6    = 0x04
)

var errSOCKS5Authentication = errors.New("SOCKS5 authentication failed")

type socks5Dialer struct {
	proxyAddress string
	username     string
	password     string
	dialer       net.Dialer
}

func newSOCKS5Dialer(proxyURL *url.URL) socks5Dialer {
	username := ""
	password := ""
	if proxyURL.User != nil {
		username = proxyURL.User.Username()
		password, _ = proxyURL.User.Password()
	}
	return socks5Dialer{proxyAddress: proxyURL.Host, username: username, password: password}
}

func (dialer socks5Dialer) DialContext(ctx context.Context, _ string, address string) (net.Conn, error) {
	connection, err := dialer.dialer.DialContext(ctx, "tcp", dialer.proxyAddress)
	if err != nil {
		return nil, err
	}
	if deadline, hasDeadline := ctx.Deadline(); hasDeadline {
		if err := connection.SetDeadline(deadline); err != nil {
			_ = connection.Close()
			return nil, err
		}
		defer connection.SetDeadline(time.Time{})
	}
	if err := dialer.connect(connection, address); err != nil {
		_ = connection.Close()
		return nil, err
	}
	return connection, nil
}

func (dialer socks5Dialer) connect(connection net.Conn, address string) error {
	method := byte(socks5NoAuth)
	if dialer.username != "" || dialer.password != "" {
		method = socks5UserPassword
	}
	if _, err := connection.Write([]byte{socks5Version, 0x01, method}); err != nil {
		return err
	}
	selection := make([]byte, 2)
	if _, err := io.ReadFull(connection, selection); err != nil {
		return err
	}
	if selection[0] != socks5Version || selection[1] == socks5NoAcceptable {
		return fmt.Errorf("invalid SOCKS5 method selection")
	}
	if selection[1] == socks5UserPassword {
		if method != socks5UserPassword {
			return errSOCKS5Authentication
		}
		if err := socks5Authenticate(connection, dialer.username, dialer.password); err != nil {
			return err
		}
	} else if selection[1] != method {
		return fmt.Errorf("unsupported SOCKS5 authentication method %d", selection[1])
	}
	request, err := socks5ConnectRequest(address)
	if err != nil {
		return err
	}
	if _, err := connection.Write(request); err != nil {
		return err
	}
	return readSOCKS5ConnectReply(connection)
}

func socks5Authenticate(connection net.Conn, username, password string) error {
	if len(username) == 0 || len(username) > 255 || len(password) == 0 || len(password) > 255 {
		return errSOCKS5Authentication
	}
	request := make([]byte, 0, len(username)+len(password)+3)
	request = append(request, 0x01, byte(len(username)))
	request = append(request, username...)
	request = append(request, byte(len(password)))
	request = append(request, password...)
	if _, err := connection.Write(request); err != nil {
		return err
	}
	reply := make([]byte, 2)
	if _, err := io.ReadFull(connection, reply); err != nil {
		return err
	}
	if reply[0] != 0x01 || reply[1] != 0x00 {
		return errSOCKS5Authentication
	}
	return nil
}

func socks5ConnectRequest(address string) ([]byte, error) {
	host, rawPort, err := net.SplitHostPort(address)
	if err != nil || host == "" {
		return nil, errors.New("invalid SOCKS5 target address")
	}
	port, err := strconv.ParseUint(rawPort, 10, 16)
	if err != nil || port == 0 {
		return nil, errors.New("invalid SOCKS5 target port")
	}
	request := []byte{socks5Version, socks5ConnectCommand, 0x00}
	if parsed := net.ParseIP(host); parsed != nil {
		if ipv4 := parsed.To4(); ipv4 != nil {
			request = append(request, socks5AddressIPv4)
			request = append(request, ipv4...)
		} else {
			request = append(request, socks5AddressIPv6)
			request = append(request, parsed.To16()...)
		}
	} else {
		if len(host) > 255 {
			return nil, errors.New("SOCKS5 target hostname is too long")
		}
		request = append(request, socks5AddressDomain, byte(len(host)))
		request = append(request, host...)
	}
	encodedPort := make([]byte, 2)
	binary.BigEndian.PutUint16(encodedPort, uint16(port))
	return append(request, encodedPort...), nil
}

func readSOCKS5ConnectReply(reader io.Reader) error {
	header := make([]byte, 4)
	if _, err := io.ReadFull(reader, header); err != nil {
		return err
	}
	if header[0] != socks5Version || header[1] != 0x00 || header[2] != 0x00 {
		return fmt.Errorf("SOCKS5 CONNECT failed with status %d", header[1])
	}
	addressLength := 0
	switch header[3] {
	case socks5AddressIPv4:
		addressLength = net.IPv4len
	case socks5AddressIPv6:
		addressLength = net.IPv6len
	case socks5AddressDomain:
		length := make([]byte, 1)
		if _, err := io.ReadFull(reader, length); err != nil {
			return err
		}
		addressLength = int(length[0])
		if addressLength == 0 {
			return errors.New("invalid SOCKS5 bound hostname")
		}
	default:
		return errors.New("invalid SOCKS5 bound address type")
	}
	boundAddressAndPort := make([]byte, addressLength+2)
	_, err := io.ReadFull(reader, boundAddressAndPort)
	return err
}