package ssh

import (
	"bytes"
	"fmt"
	"net"
	"time"

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

// Client 封装 SSH 客户端连接
type Client struct {
	client  *ssh.Client
	timeout time.Duration
}

// ClientConfig SSH 连接参数
type ClientConfig struct {
	Host     string
	Port     int
	Username string
	Password string
	Timeout  time.Duration
}

// NewClient 创建新的 SSH 连接
func NewClient(cfg ClientConfig) (*Client, error) {
	if cfg.Port == 0 {
		cfg.Port = 22
	}
	if cfg.Timeout == 0 {
		cfg.Timeout = 30 * time.Second
	}

	sshConfig := &ssh.ClientConfig{
		User: cfg.Username,
		Auth: []ssh.AuthMethod{
			ssh.Password(cfg.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         cfg.Timeout,
	}

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

	return &Client{
		client:  client,
		timeout: cfg.Timeout,
	}, nil
}

// Close 关闭 SSH 连接
func (c *Client) Close() error {
	if c.client != nil {
		return c.client.Close()
	}
	return nil
}

// Exec 执行单条命令并返回输出
func (c *Client) Exec(cmd string) (string, error) {
	session, err := c.client.NewSession()
	if err != nil {
		return "", fmt.Errorf("创建 SSH 会话失败: %w", err)
	}
	defer session.Close()

	var stdout, stderr bytes.Buffer
	session.Stdout = &stdout
	session.Stderr = &stderr

	if err := session.Run(cmd); err != nil {
		return stdout.String(), fmt.Errorf("执行命令失败: %w, stderr: %s", err, stderr.String())
	}

	return stdout.String(), nil
}

// ExecBatch 批量执行命令，返回每条命令的输出
func (c *Client) ExecBatch(cmds []string) (map[string]string, error) {
	results := make(map[string]string, len(cmds))

	for _, cmd := range cmds {
		output, err := c.Exec(cmd)
		if err != nil {
			results[cmd] = fmt.Sprintf("ERROR: %v", err)
			continue
		}
		results[cmd] = output
	}

	return results, nil
}

// ExecInteractive 交互式执行命令（用于需要 enable 模式的设备）
func (c *Client) ExecInteractive(cmds []string) (map[string]string, error) {
	session, err := c.client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("创建 SSH 会话失败: %w", err)
	}
	defer session.Close()

	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if err := session.RequestPty("xterm", 80, 200, modes); err != nil {
		return nil, fmt.Errorf("请求 PTY 失败: %w", err)
	}

	stdinPipe, err := session.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("获取 stdin 失败: %w", err)
	}

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf

	if err := session.Shell(); err != nil {
		return nil, fmt.Errorf("启动 Shell 失败: %w", err)
	}

	results := make(map[string]string, len(cmds))
	marker := "===CMD_END_MARKER==="

	for _, cmd := range cmds {
		// 写入命令和唯一标记，用于分隔输出
		fullCmd := fmt.Sprintf("%s; echo '%s'\n", cmd, marker)
		stdinPipe.Write([]byte(fullCmd))

		// 等待标记出现
		time.Sleep(500 * time.Millisecond)

		output := stdoutBuf.String()
		stdoutBuf.Reset()

		// 提取标记之前的内容作为命令输出
		if idx := bytes.LastIndex(stdoutBuf.Bytes(), []byte(marker)); idx >= 0 {
			results[cmd] = output[:idx]
		} else {
			results[cmd] = output
		}
	}

	stdinPipe.Write([]byte("exit\n"))
	session.Wait()

	return results, nil
}

// CheckConnection 检测目标主机 SSH 端口是否可达
func CheckConnection(host string, port int, timeout time.Duration) error {
	if port == 0 {
		port = 22
	}
	if timeout == 0 {
		timeout = 5 * time.Second
	}

	addr := fmt.Sprintf("%s:%d", host, port)
	conn, err := net.DialTimeout("tcp", addr, timeout)
	if err != nil {
		return fmt.Errorf("连接 %s 失败: %w", addr, err)
	}
	conn.Close()
	return nil
}
