package protocol

import (
	"context"
	"crypto/tls"
	"net"
	"time"
)

// Dialer used by probes to establish a connection.
var Dialer net.Dialer

// Probe a network service by reading its banner and running protocol detection.
func Probe(network, address string) (proto *Protocol, confidence float64, err error) {
	return ProbeContext(context.Background(), network, address)
}

// ProbeContext is like [Probe] with a [context.Context].
func ProbeContext(ctx context.Context, network, address string) (proto *Protocol, confidence float64, err error) {
	var conn net.Conn
	if conn, err = Dialer.DialContext(ctx, network, address); err != nil {
		return
	}
	defer func() { _ = conn.Close() }()
	return probeConn(conn)
}

// ProbeTimeout is like [Probe] but with a fixed timeout.
func ProbeTimeout(network, address string, timeout time.Duration) (proto *Protocol, confidence float64, err error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	return ProbeContext(ctx, network, address)
}

// ProbeTLS is like [Probe] but first establishes a TLS connection.
func ProbeTLS(network, address string, tlsConfig *tls.Config) (proto *Protocol, confidence float64, err error) {
	return ProbeTLSContext(context.Background(), network, address, tlsConfig)
}

// ProbeTLSContext is like [ProbeTLS] with a [context.Context].
func ProbeTLSContext(ctx context.Context, network, address string, tlsConfig *tls.Config) (proto *Protocol, confidence float64, err error) {
	var conn net.Conn
	if conn, err = Dialer.DialContext(ctx, network, address); err != nil {
		return
	}
	defer func() { _ = conn.Close() }()
	secure := tls.Client(conn, tlsConfig)
	if err = secure.Handshake(); err != nil {
		return
	}
	return probeConn(secure)
}

// ProbeTLSTimeout is like [ProbeTLS] but with a fixed timeout.
func ProbeTLSTimeout(network, address string, tlsConfig *tls.Config, timeout time.Duration) (proto *Protocol, confidence float64, err error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	return ProbeTLSContext(ctx, network, address, tlsConfig)
}

func probeConn(conn net.Conn) (proto *Protocol, confidence float64, err error) {
	var (
		data = make([]byte, 2048)
		n    int
	)
	if n, err = conn.Read(data); err != nil {
		return
	}
	return Detect(Client, data[:n], getPortFromAddr(conn.LocalAddr()), getPortFromAddr(conn.RemoteAddr()))
}