package bzsftp

import (
	"fmt"
	"gitee.com/breezeHub/bzv2/pkg/code/bzerr"
	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"io"
	"net"
	"os"
	"time"
)

type config struct {
	Host       string //ip
	Port       int    // 端口
	User       string //用户名
	Password   string //密码
	SSHTimeout time.Duration
}

type Client struct {
	config     *config
	sshClient  *ssh.Client  //ssh client
	sftpClient *sftp.Client //sftp client
}

func New(host string, port int, user, password string, opts ...Option) (*Client, error) {
	c := &Client{
		config: &config{
			Host:       host,
			Port:       port,
			User:       user,
			Password:   password,
			SSHTimeout: 10 * time.Second,
		},
	}

	for _, opt := range opts {
		opt(c)
	}

	// 连接 ssh
	conf := ssh.ClientConfig{
		User: c.config.User,
		Auth: []ssh.AuthMethod{ssh.Password(c.config.Password)},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
		Timeout: c.config.SSHTimeout,
	}
	addr := fmt.Sprintf("%s:%d", c.config.Host, c.config.Port)
	sshClient, err := ssh.Dial("tcp", addr, &conf)
	if err != nil {
		return c, bzerr.Wrap(err, "ssh dial error")
	}
	c.sshClient = sshClient

	// 使用sshClient 构建 sftpClient
	sftpClient, err := sftp.NewClient(sshClient)
	if err != nil {
		return c, bzerr.Wrap(err, "sftp dial error")
	}
	c.sftpClient = sftpClient

	return c, nil
}

func (c *Client) RunShell(shell string) (string, error) {
	//获取session，这个session是用来远程执行操作的
	session, err := c.sshClient.NewSession()
	if err != nil {
		return "", bzerr.Wrap(err, "new session error")
	}

	//执行shell
	output, err := session.CombinedOutput(shell)
	if err != nil {
		return "", bzerr.Wrap(err, "run shell error")
	}

	return string(output), nil
}

func (c *Client) Upload(srcPath, dstPath string, bufSizes ...int) error {
	srcFile, err := os.Open(srcPath) //本地
	if err != nil {
		return bzerr.Wrap(err, "read file error")
	}
	defer srcFile.Close()

	dstFile, err := c.sftpClient.Create(dstPath) //远程
	if err != nil {
		return bzerr.Wrap(err, "create remote file error")
	}
	defer dstFile.Close()

	bufSize := 1024
	if len(bufSizes) > 0 {
		bufSize = bufSizes[0]
	}

	buf := make([]byte, bufSize)
	for {
		n, err := srcFile.Read(buf)
		if err != nil {
			if err != io.EOF {
				return bzerr.Wrap(err, "copy to remote file error")
			}
			break
		}
		if _, err = dstFile.Write(buf[:n]); err != nil {
			return bzerr.Wrap(err, "copy to remote file error")
		}
	}
	return nil
}

func (c *Client) Download(srcPath, dstPath string) error {
	srcFile, err := c.sftpClient.Open(srcPath) //远程
	if err != nil {
		return bzerr.Wrap(err, "read file error")
	}
	defer srcFile.Close()

	dstFile, err := os.Create(dstPath) //本地
	if err != nil {
		return bzerr.Wrap(err, "create remote file error")
	}
	defer dstFile.Close()

	if _, err := srcFile.WriteTo(dstFile); err != nil {
		return bzerr.Wrap(err, "download error")
	}
	return nil
}