package socks5

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

// Config carries connection-time SOCKS5 settings.
type Config struct {
	ProxyAddr string // "host:port"
	UseAuth   bool
	Login     string
	Password  string
}

// Dial opens a TCP connection to the SOCKS5 proxy, runs the greeting,
// optionally authenticates with username/password (RFC 1929), and
// issues a CONNECT to host:port (sent as ATYP=03 domain so the proxy
// resolves on its side). Returns the established net.Conn ready for
// bidirectional traffic.
//
// The given ctx bounds dial + handshake; once Dial returns, the conn
// has its own deadline-free I/O state.
func Dial(ctx context.Context, cfg Config, host string, port uint16) (net.Conn, error) {
	d := net.Dialer{}
	conn, err := d.DialContext(ctx, "tcp", cfg.ProxyAddr)
	if err != nil {
		return nil, fmt.Errorf("dial proxy: %w", err)
	}
	if dl, ok := ctx.Deadline(); ok {
		conn.SetDeadline(dl)
	}
	if err := handshake(conn, cfg, host, port); err != nil {
		conn.Close()
		return nil, err
	}
	conn.SetDeadline(time.Time{})
	return conn, nil
}

func handshake(conn net.Conn, cfg Config, host string, port uint16) error {
	// Greeting
	if cfg.UseAuth {
		if _, err := conn.Write([]byte{0x05, 0x02, 0x00, 0x02}); err != nil {
			return fmt.Errorf("greet write: %w", err)
		}
	} else {
		if _, err := conn.Write([]byte{0x05, 0x01, 0x00}); err != nil {
			return fmt.Errorf("greet write: %w", err)
		}
	}
	var rep [2]byte
	if _, err := io.ReadFull(conn, rep[:]); err != nil {
		return fmt.Errorf("greet read: %w", err)
	}
	if rep[0] != 0x05 {
		return fmt.Errorf("greet: server version %#x is not SOCKS5", rep[0])
	}
	if rep[1] == 0xff {
		return errors.New("greet: proxy rejected all offered auth methods")
	}
	method := rep[1]

	// Auth subneg
	if method == 0x02 {
		if !cfg.UseAuth {
			return errors.New("proxy requires auth but Config.UseAuth is false")
		}
		if len(cfg.Login) > 255 || len(cfg.Password) > 255 {
			return 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 _, err := conn.Write(buf); err != nil {
			return fmt.Errorf("auth write: %w", err)
		}
		if _, err := io.ReadFull(conn, rep[:]); err != nil {
			return fmt.Errorf("auth read: %w", err)
		}
		if rep[1] != 0x00 {
			return errors.New("auth: invalid login or password")
		}
	}

	// CONNECT
	if len(host) > 255 {
		return errors.New("host too long")
	}
	req := make([]byte, 0, 7+len(host))
	req = append(req, 0x05, 0x01, 0x00, 0x03, byte(len(host)))
	req = append(req, []byte(host)...)
	pBuf := make([]byte, 2)
	binary.BigEndian.PutUint16(pBuf, port)
	req = append(req, pBuf...)
	if _, err := conn.Write(req); err != nil {
		return fmt.Errorf("connect write: %w", err)
	}
	var creply [10]byte
	if _, err := io.ReadFull(conn, creply[:]); err != nil {
		return fmt.Errorf("connect read: %w", err)
	}
	if creply[0] != 0x05 {
		return fmt.Errorf("connect: server version %#x is not SOCKS5", creply[0])
	}
	if creply[1] != 0x00 {
		return fmt.Errorf("connect: REP=%#02x", creply[1])
	}
	return nil
}