package sshclient

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

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

// Config SSH 连接配置（不含连接状态）
type Config struct {
	User     string
	Password string
	KeyFile  string
	Port     int
	Timeout  time.Duration
	SudoPwd  string // 非 root 用户执行特权命令(sudo)时的密码
}

// 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 cleanOutput(out), fmt.Errorf("命令执行失败: %w", err)
	}
	return cleanOutput(out), nil
}

// isNoiseLine 判断某行是否为登录横幅/欢迎语等噪声（fnOS 的 admin 用户会把
// "Could not chdir to home directory" 之类写入 stdout，污染命令解析）。
func isNoiseLine(line string) bool {
	s := strings.TrimSpace(line)
	if s == "" {
		return false
	}
	prefixes := []string{
		"Could not chdir to",
		"Last login:",
		"Welcome to",
		"/etc/motd",
		"Permission denied",
		"bash: line 1:",
	}
	for _, p := range prefixes {
		if strings.HasPrefix(s, p) {
			return true
		}
	}
	// 某些横幅是 "Could: command not found" 这类由前面噪声行产生，
	// 形如 "xxx: command not found" 且前半段是噪声关键词的也过滤
	if strings.Contains(s, ": command not found") {
		return true
	}
	return false
}

// cleanOutput 过滤掉登录横幅等噪声行，仅保留真实命令输出
func cleanOutput(raw []byte) string {
	lines := strings.Split(string(raw), "\n")
	var kept []string
	for _, ln := range lines {
		if isNoiseLine(ln) {
			continue
		}
		kept = append(kept, ln)
	}
	return strings.Join(kept, "\n")
}

// EnsureHome 确保远端登录用户的 HOME 目录存在，消除 fnOS "Could not chdir to
// home directory" 登录横幅（该横幅会写入 stdout 污染命令输出）。
func EnsureHome(client *ssh.Client) {
	out, err := RunCommand(client, "echo $HOME")
	home := strings.TrimSpace(out)
	if err != nil || home == "" || home == "/" {
		return
	}
	RunCommand(client, fmt.Sprintf("mkdir -p %s", home))
}

// 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
}