package host

import (
	"time"

	"github.com/pingcap/test-infra/caselib/pkg/command"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
)

// var TimeoutErr = errors.New("Exec command timeout")

type DrainerHost struct {
	drainer resource.TiDBDrainer
}

func NewDrainerHost(drainer resource.TiDBDrainer) *DrainerHost {
	return &DrainerHost{
		drainer: drainer,
	}
}

func (w *DrainerHost) Exec(cmd string, timeout time.Duration) (*command.ExecOutput, error) {
	done := make(chan *command.ExecOutput, 1)
	logger.Info("Execute command", zap.String("command", cmd), zap.Duration("timeout", timeout),
		zap.String("resource name", w.drainer.Handle().Name))
	go func() {
		stdout, stderr, exitCode, err := w.drainer.Exec(resource.TiDBDrainerExecOptions{Command: cmd})
		done <- &command.ExecOutput{Err: err, Stdout: stdout, Stderr: stderr, ExitCode: exitCode}
	}()

	select {
	case <-time.After(timeout):
		return nil, TimeoutErr
	case output := <-done:
		return output, output.Err
	}
}