package ssh

import (
	"bytes"
	"fmt"
	"io"
	"sync"

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

// SSHClientSession SSH客户端会话对象
type SSHClientSession struct {
	Stdin   io.WriteCloser
	Stdout  *singleWriter
	Session *gossh.Session
}

// Close 关闭会话
func (s *SSHClientSession) Close() {
	if s.Stdin != nil {
		s.Stdin.Close()
	}
	if s.Stdout != nil {
		s.Stdout = nil
	}
	if s.Session != nil {
		s.Session.Close()
	}
}

// Write 写入命令 回车(\n)才会执行
func (s *SSHClientSession) Write(cmd string) (int, error) {
	if s.Stdin == nil {
		return 0, fmt.Errorf("ssh client session is nil to content write failed")
	}
	return s.Stdin.Write([]byte(cmd))
}

// Read 读取结果
func (s *SSHClientSession) Read() []byte {
	if s.Stdout == nil {
		return []byte{}
	}
	bs := s.Stdout.Bytes()
	if len(bs) > 0 {
		s.Stdout.Reset()
		return bs
	}
	return []byte{}
}

// singleWriter SSH客户端会话消息
type singleWriter struct {
	b  bytes.Buffer
	mu sync.Mutex
}

func (w *singleWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.b.Write(p)
}
func (w *singleWriter) Bytes() []byte {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.b.Bytes()
}
func (w *singleWriter) Reset() {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.b.Reset()
}