package telnet

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

// TelnetClientSession Telnet客户端会话对象
type TelnetClientSession struct {
	Client net.Conn
}

// Close 关闭会话
func (s *TelnetClientSession) Close() {
	if s.Client != nil {
		s.Client.Close()
	}
}

// WindowChange informs the remote host about a terminal window dimension change to h rows and w columns.
func (s *TelnetClientSession) WindowChange(h, w int) error {
	if s.Client == nil {
		return fmt.Errorf("client is nil to content write failed")
	}
	// 需要确保接收方理解并正确处理发送窗口大小设置命令
	s.Client.Write([]byte{255, 251, 31})
	s.Client.Write([]byte{255, 250, 31, byte(w >> 8), byte(w & 0xFF), byte(h >> 8), byte(h & 0xFF), 255, 240})
	return nil
}

// Write 写入命令 不带回车(\n)也会执行根据客户端情况
func (s *TelnetClientSession) Write(cmd string) (int, error) {
	if s.Client == nil {
		return 0, fmt.Errorf("client is nil to content write failed")
	}
	return s.Client.Write([]byte(cmd))
}

// Read 读取结果 等待一会才有结果
func (s *TelnetClientSession) Read() []byte {
	if s.Client == nil {
		return []byte{}
	}

	buf := make([]byte, 1024)
	// 设置读取超时时间为100毫秒
	s.Client.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	n, err := s.Client.Read(buf)
	if err != nil {
		return []byte{}
	}
	return buf[:n]
}

// CombinedOutput 发送命令带结果返回
func (s *TelnetClientSession) CombinedOutput(cmd string) (string, error) {
	n, err := s.Write(cmd)
	if n == 0 || err != nil {
		return "", err
	}

	var buf bytes.Buffer
	tmp := make([]byte, 1024)
	for {
		// 设置读取超时时间为1000毫秒
		s.Client.SetReadDeadline(time.Now().Add(1000 * time.Millisecond))
		n, err := s.Client.Read(tmp)
		if err != nil {
			// 判断是否是超时错误
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				break
			}
			break
		}
		if n == 0 {
			break
		}
		buf.Write(tmp[:n])
	}
	defer buf.Reset()

	return buf.String(), nil
}
