package collector

import (
	"bytes"
	"fmt"
	"strconv"

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

// SSHCollector SSH采集器
type SSHCollector struct {
	client *ssh.Client
}

// SSHConfig SSH连接配置
type SSHConfig struct {
	Host     string
	Port     int
	Username string
	Password string
	KeyPath  string
}

// NewSSHCollector 创建新的SSH采集器
func NewSSHCollector(cfg *SSHConfig) (*SSHCollector, error) {
	// 构建SSH客户端配置
	sshConfig := &ssh.ClientConfig{
		User:            cfg.Username,
		Auth:            []ssh.AuthMethod{},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // 生产环境中应该使用更安全的HostKeyCallback
	}

	// 添加认证方式
	if cfg.Password != "" {
		sshConfig.Auth = append(sshConfig.Auth, ssh.Password(cfg.Password))
	}

	// 连接SSH服务器
	addr := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)
	client, err := ssh.Dial("tcp", addr, sshConfig)
	if err != nil {
		return nil, err
	}

	return &SSHCollector{
		client: client,
	}, nil
}

// Close 关闭SSH连接
func (c *SSHCollector) Close() error {
	return c.client.Close()
}

// executeCommand 执行SSH命令
func (c *SSHCollector) executeCommand(cmd string) (string, error) {
	// 创建SSH会话
	session, err := c.client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	// 执行命令并获取输出
	var stdout bytes.Buffer
	session.Stdout = &stdout

	if err := session.Run(cmd); err != nil {
		return "", err
	}

	return stdout.String(), nil
}

// GetCPUUsage 获取CPU使用率
func (c *SSHCollector) GetCPUUsage() (float64, error) {
	// 使用top命令获取CPU使用率
	cmd := "top -bn1 | grep 'Cpu(s)' | awk '{print $2 + $4}'"

	output, err := c.executeCommand(cmd)
	if err != nil {
		return 0, err
	}

	// 解析输出
	usage, err := strconv.ParseFloat(output, 64)
	if err != nil {
		return 0, err
	}

	return usage, nil
}

// GetMemoryUsage 获取内存使用率
func (c *SSHCollector) GetMemoryUsage() (float64, error) {
	// 使用free命令获取内存使用情况
	cmd := "free | grep Mem | awk '{print ($3/$2)*100}'"

	output, err := c.executeCommand(cmd)
	if err != nil {
		return 0, err
	}

	// 解析输出
	usage, err := strconv.ParseFloat(output, 64)
	if err != nil {
		return 0, err
	}

	return usage, nil
}

// GetDiskUsage 获取磁盘使用率
func (c *SSHCollector) GetDiskUsage(partition string) (float64, error) {
	// 使用df命令获取磁盘使用情况
	cmd := fmt.Sprintf("df -h %s | tail -1 | awk '{print $5}' | sed 's/%%//'", partition)

	output, err := c.executeCommand(cmd)
	if err != nil {
		return 0, err
	}

	// 解析输出
	usage, err := strconv.ParseFloat(output, 64)
	if err != nil {
		return 0, err
	}

	return usage, nil
}