package main

import (
	"bytes"
	"os"
	"time"

	"github.com/silenceper/pool"
	log "github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
)

const SshDialTimeout = time.Duration(10 * time.Second)
const SshInactivityTimeout = time.Duration(time.Minute)

type Ssh struct {
	name   string
	signer ssh.Signer
	config *ssh.ClientConfig
	client *ssh.Client
}

func NewSsh(name, addr, user, key string) (*Ssh, error) {
	log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("done")

	s := &Ssh{
		name: name,
	}

	k, err := os.ReadFile(key)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "os.ReadFile", "error": err}).Errorf("")
		return s, err
	}

	parsedKey, err := ssh.ParseRawPrivateKey(k)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "ssh.ParseRawPrivateKey", "error": err}).Errorf("")
		return s, err
	}

	s.signer, err = ssh.NewSignerFromKey(parsedKey)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "ssh.NewSignerFromKey", "error": err}).Errorf("")
		return s, err
	}

	s.config = &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(s.signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         SshDialTimeout,
	}

	s.client, err = ssh.Dial("tcp", addr, s.config)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "ssh.Dial", "error": err}).Errorf("")
		return s, err
	}

	return s, nil

}

func (s *Ssh) Close() error {
	log.WithFields(log.Fields{"name": s.name}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": s.name}).Debugf("done")

	return s.client.Close()
}

func NewSshPool(name, addr, user, key string) (pool.Pool, error) {
	log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("done")

	//factory Specify the method to create the connection
	factory := func() (interface{}, error) { return NewSsh(name, addr, user, key) }

	// close Specify the method to close the connection
	close := func(v interface{}) error { return v.(*Ssh).Close() }

	// Create a connection pool: Initialize the number of connections to 0, the maximum idle connection is 2, and the maximum concurrent connection is 25
	poolConfig := &pool.Config{
		InitialCap: 0,
		MaxIdle:    2,
		MaxCap:     25,
		Factory:    factory,
		Close:      close,
		//Ping:       ping,
		//The maximum idle time of the connection, the connection exceeding this time will be closed, which can avoid the problem of automatic failure when connecting to EOF when idle
		IdleTimeout: SshInactivityTimeout,
	}

	return pool.NewChannelPool(poolConfig)
}

func (s *Ssh) Exec(cmd string) (string, error) {
	log.WithFields(log.Fields{"name": s.name, "cmd": cmd}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": s.name, "cmd": cmd}).Debugf("done")

	session, err := s.client.NewSession()
	if err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "client.NewSession", "error": err}).Errorf("")
		return "", err
	}
	defer session.Close()

	if session.Setenv("TZ", cfg.Timezone); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.Setenv", "error": err}).Errorf("")
		return "", err
	}

	var bufout, buferr bytes.Buffer
	session.Stdout = &bufout
	session.Stderr = &buferr

	err = session.Run(cmd)
	if err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.Run", "error": err, "stderr": buferr.String()}).Errorf("")
		return "", err
	}

	return bufout.String(), nil
}