package ssh

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

type Session struct {
	Host    string //ip:port
	Session *ssh.Session
	Client  *ssh.Client
	config  *ssh.ClientConfig
}

type Result struct {
	Id         int
	Host       string
	Command    string
	LocalPath  string
	RemotePath string
	Result     string
	StartTime  time.Time
	EndTime    time.Time
	Error      error
}

func NewSession(host string, conf *ssh.ClientConfig) *Session {
	return &Session{Host: host, config: conf}
}

func (s *Session) Connect() error {
	if s.Client != nil {
		return s.createSession()
	}
	client, err := ssh.Dial("tcp", s.Host, s.config)
	if err != nil {
		return err
	}
	s.Client = client
	return s.createSession()
}

func (s *Session) createSession() error {
	if s.Session != nil {
		return nil
	}
	session, err := s.Client.NewSession()
	if err != nil {
		return err
	}
	s.Session = session
	return nil
}

func (s *Session) Close() {
	if s.Session != nil {
		s.Session.Close()
		s.Session = nil
	}
	if s.Client != nil {
		s.Client.Close()
		s.Client = nil
	}
}

func (s *Session) Exec(id int, command string) *Result {

	result := &Result{
		Id:      id,
		Host:    s.Host,
		Command: command,
	}
	if err := s.Connect(); err != nil {
		result.Error = err
		return result
	}

	var b bytes.Buffer
	var b1 bytes.Buffer

	s.Session.Stdout = &b
	s.Session.Stderr = &b1

	start := time.Now()
	if err := s.Session.Run(command); err != nil {
		result.Error = err
		result.Result = b1.String()
		return result
	}
	end := time.Now()
	result.Result = b.String()
	result.StartTime = start
	result.EndTime = end
	return result
}