package checkrunner

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"time"
)

// TCPConnect reports success if a TCP handshake against host:port
// completes within timeout. This is the primitive behind the prober's
// per-port inbound reachability checks (22/80/443/8080).
func TCPConnect(host string, port int, timeout time.Duration) func(ctx context.Context) Result {
	target := net.JoinHostPort(host, strconv.Itoa(port))
	checkType := "tcp-" + strconv.Itoa(port)
	return run(checkType, target, func(ctx context.Context) error {
		d := net.Dialer{Timeout: timeout}
		conn, err := d.DialContext(ctx, "tcp", target)
		if err != nil {
			return err
		}
		return conn.Close()
	})
}

// SSHBanner performs a TCP connect to host:22 and additionally verifies the
// remote sends an "SSH-2.0-" banner, without performing any auth handshake.
// Used for the optional ssh check type.
func SSHBanner(host string, timeout time.Duration) func(ctx context.Context) Result {
	target := net.JoinHostPort(host, "22")
	return run("ssh", target, func(ctx context.Context) error {
		d := net.Dialer{Timeout: timeout}
		conn, err := d.DialContext(ctx, "tcp", target)
		if err != nil {
			return err
		}
		defer conn.Close()

		conn.SetReadDeadline(time.Now().Add(timeout))
		buf := make([]byte, 8)
		n, err := conn.Read(buf)
		if err != nil {
			return fmt.Errorf("read banner: %w", err)
		}
		if n < 8 || string(buf[:8]) != "SSH-2.0-" {
			return fmt.Errorf("unexpected banner prefix %q", string(buf[:n]))
		}
		return nil
	})
}