package sshclient

import (
	"fmt"
	"net"
	"os"
	"time"

	"golang.org/x/crypto/ssh"
)

// Config SSH 连接配置（不含连接状态）
type Config struct {
	User     string
	Password string
	KeyFile  string
	Port     int
	Timeout  time.Duration
}

// NewConfig 创建配置
func NewConfig(user, password, keyFile string, port int, timeout time.Duration) *Config {
	return &Config{
		User:     user,
		Password: password,
		KeyFile:  keyFile,
		Port:     port,
		Timeout:  timeout,
	}
}

// Connect 建立 SSH 连接，由调用方负责 Close
func Connect(ip string, cfg *Config) (*ssh.Client, error) {
	sshCfg, err := cfg.buildSSHConfig()
	if err != nil {
		return nil, err
	}

	addr := net.JoinHostPort(ip, fmt.Sprintf("%d", cfg.Port))
	client, err := ssh.Dial("tcp", addr, sshCfg)
	if err != nil {
		return nil, fmt.Errorf("SSH 连接 %s 失败: %w", addr, err)
	}
	return client, nil
}

// RunCommand 在已有 SSH 连接上执行命令，返回输出（实现 stress.Executor 接口）
func RunCommand(client *ssh.Client, command string) (string, error) {
	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("创建会话失败: %w", err)
	}
	defer session.Close()

	out, err := session.CombinedOutput(command)
	if err != nil {
		return string(out), fmt.Errorf("命令执行失败: %w", err)
	}
	return string(out), nil
}

// buildSSHConfig 构建认证配置
func (c *Config) buildSSHConfig() (*ssh.ClientConfig, error) {
	var authMethods []ssh.AuthMethod

	if c.Password != "" {
		authMethods = append(authMethods, ssh.Password(c.Password))
	}

	if c.KeyFile != "" {
		key, err := os.ReadFile(c.KeyFile)
		if err != nil {
			return nil, fmt.Errorf("读取密钥文件失败: %w", err)
		}
		signer, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return nil, fmt.Errorf("解析私钥失败: %w", err)
		}
		authMethods = append(authMethods, ssh.PublicKeys(signer))
	}

	if len(authMethods) == 0 {
		return nil, fmt.Errorf("未提供认证方式（密码或密钥）")
	}

	return &ssh.ClientConfig{
		User:            c.User,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         c.Timeout,
	}, nil
}