package client

import (
	"fmt"
	"io"
	"net"
	"strconv"
)

// writeProxyHeader writes the optional PROXY protocol header to a local
// target. The destination is the unspecified address of the source family on
// port 443, matching the reference clients (the real public address is not
// known to the client).
func writeProxyHeader(w io.Writer, version, peer, sni string) error {
	switch version {
	case "v1":
		_, err := w.Write(proxyV1(peer))
		return err
	case "v2":
		_, err := w.Write(proxyV2(peer, sni))
		return err
	default:
		return fmt.Errorf("unknown proxy protocol %q", version)
	}
}

func proxyV1(peer string) []byte {
	ip, port, ok := parsePeer(peer)
	if !ok {
		return []byte("PROXY UNKNOWN\r\n")
	}
	if v4 := ip.To4(); v4 != nil {
		return []byte(fmt.Sprintf("PROXY TCP4 %s 0.0.0.0 %d 443\r\n", v4, port))
	}
	return []byte(fmt.Sprintf("PROXY TCP6 %s :: %d 443\r\n", ip, port))
}

func proxyV2(peer, sni string) []byte {
	signature := []byte("\r\n\r\n\x00\r\nQUIT\n")
	ip, port, ok := parsePeer(peer)

	command := byte(0x21) // PROXY
	family := byte(0x00)  // UNSPEC
	var address []byte
	if ok {
		if v4 := ip.To4(); v4 != nil {
			family = 0x11 // TCP4
			address = append(address, v4...)
			address = append(address, 0, 0, 0, 0)
		} else {
			family = 0x21 // TCP6
			address = append(address, ip.To16()...)
			address = append(address, make([]byte, 16)...)
		}
		address = append(address, byte(port>>8), byte(port), 0x01, 0xbb) // destination port 443
	} else {
		command = 0x20 // LOCAL
	}

	var tlv []byte
	if sni != "" && len(sni) <= 255 {
		tlv = []byte{0x02, 0x00, byte(len(sni))}
		tlv = append(tlv, sni...)
	}
	length := len(address) + len(tlv)

	header := make([]byte, 0, len(signature)+4+length)
	header = append(header, signature...)
	header = append(header, command, family, byte(length>>8), byte(length))
	header = append(header, address...)
	header = append(header, tlv...)
	return header
}

func parsePeer(peer string) (net.IP, int, bool) {
	host, portText, err := net.SplitHostPort(peer)
	if err != nil {
		host, portText = peer, "0"
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return nil, 0, false
	}
	if v4 := ip.To4(); v4 != nil {
		ip = v4
	}
	port, _ := strconv.Atoi(portText)
	if port < 0 || port > 65535 {
		port = 0
	}
	return ip, port, true
}