package client

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/nixc/reverse-ssh-traefik/internal/sshutil"
	"golang.org/x/crypto/ssh"
)

// LoadPrivateKey loads an SSH private key from a file path or raw PEM content.
func LoadPrivateKey(keyOrPath string) (ssh.Signer, error) {
	var keyBytes []byte

	if isFilePath(keyOrPath) {
		// #nosec G304 -- operator-supplied SSH private key path.
		data, err := os.ReadFile(keyOrPath)
		if err != nil {
			return nil, fmt.Errorf("read key file %s: %w", keyOrPath, err)
		}
		keyBytes = data
	} else {
		keyBytes = []byte(keyOrPath)
	}

	signer, err := ssh.ParsePrivateKey(keyBytes)
	if err != nil {
		return nil, fmt.Errorf("parse private key: %w", err)
	}

	return signer, nil
}

func isFilePath(v string) bool {
	if strings.HasPrefix(v, "/") || strings.HasPrefix(v, "./") || strings.HasPrefix(v, "~") {
		return true
	}
	return !strings.Contains(v, "-----BEGIN")
}

// Connect establishes an SSH connection to the tunnel server.
func Connect(addr string, signer ssh.Signer, hostKeyCallback ssh.HostKeyCallback) (*ssh.Client, error) {
	config := &ssh.ClientConfig{
		User: "tunnel",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		Config:            sshutil.SecureConfig(),
		HostKeyAlgorithms: sshutil.HostKeyAlgorithms(),
		HostKeyCallback:   hostKeyCallback,
		Timeout:           15 * time.Second,
	}

	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return nil, fmt.Errorf("SSH dial %s: %w", addr, err)
	}

	return client, nil
}