package executor

import (
	"fmt"
	"io"
	"time"

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

type SSHExecutor struct {
	client *ssh.Client
}

func NewSSHExecutor(host string, port int, user, pass string) (*SSHExecutor, error) {
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pass),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         5 * time.Second,
	}

	cli, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", host, port), config)
	if err != nil {
		return nil, err
	}

	return &SSHExecutor{client: cli}, nil
}

func (s *SSHExecutor) Exec(cmd string, opts ...ExecOption) (*ExecResult, error) {
	session, err := s.client.NewSession()
	if err != nil {
		return nil, err
	}
	defer session.Close()

	stdout, _ := session.StdoutPipe()
	stderr, _ := session.StderrPipe()

	err = session.Start(cmd)
	if err != nil {
		return nil, err
	}

	outBytes, _ := io.ReadAll(stdout)
	errBytes, _ := io.ReadAll(stderr)

	waitErr := session.Wait()

	exitCode := 0
	if waitErr != nil {
		if exit, ok := waitErr.(*ssh.ExitError); ok {
			exitCode = exit.ExitStatus()
		}
	}

	return &ExecResult{
		Command:  cmd,
		Stdout:   string(outBytes),
		Stderr:   string(errBytes),
		ExitCode: exitCode,
	}, waitErr
}

// Close closes the SSH client connection
func (s *SSHExecutor) Close() error {
	if s.client != nil {
		return s.client.Close()
	}
	return nil
}