package ssh

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"strconv"
)

type SSHConfig struct {
	User string
	Auth []ssh.AuthMethod
	Host string
	Port int
}

type SSHClient struct {
	session *ssh.Session
	conf    *SSHConfig
}

func New(ip, port, user, pw string) (*SSHClient, error) {
	p, _ := strconv.Atoi(port)
	conf := &SSHConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pw),
		},
		Host: ip,
		Port: p,
	}

	return &SSHClient{conf: conf}, nil
}

func (cli *SSHClient) Session() error {
	sshConfig := &ssh.ClientConfig{
		User:            cli.conf.User,
		Auth:            cli.conf.Auth,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	connection, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", cli.conf.Host, cli.conf.Port), sshConfig)
	if err != nil {
		return err
	}
	session, err := connection.NewSession()
	if err != nil {
		return err
	}

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	if err := session.RequestPty("xterm", 80, 40, modes); err != nil {
		session.Close()
		return err
	}

	cli.session = session

	return nil
}

func (cli *SSHClient) RunCommand(command string) ([]byte, error) {

	if err := cli.Session(); err != nil {
		return nil, err
	}

	var b bytes.Buffer
	cli.session.Stdout = &b

	err := cli.session.Run(command)
	if err != nil {
		return nil, err
	}

	cli.session.Close()

	return b.Bytes(), nil
}