package tiup

import (
	"context"
	"fmt"
	"time"

	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/test-infra/sdk/resource"
)

type workloadHost struct {
	workload resource.WorkloadNode
}

func NewWorkloadHost(workload resource.WorkloadNode) *workloadHost {
	return &workloadHost{workload: workload}
}

func (w *workloadHost) Exec(cmd string, timeout time.Duration) (stdout string, err error) {
	var stderr string
	var exitCode int
	done := make(chan struct{})
	go func() {
		stdout, stderr, exitCode, err = w.workload.Exec(resource.WorkloadNodeExecOptions{Command: cmd})
		close(done)
	}()

	select {
	case <-time.After(timeout):
		return "", fmt.Errorf("%s exec %s timeout", w.workload.Handle().Name, cmd)
	case <-done:
		if exitCode != 0 {
			return "", fmt.Errorf("%s exec %s exit code %d stderr %s", w.workload.Handle().Name, cmd, exitCode, stderr)
		}

		return stdout, err
	}
}

func (w *workloadHost) Run(cmd string, timeout time.Duration) error {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	workloadCmd := workloadnode.CommandContext(ctx, cmd)
	err := workloadCmd.Run(w.workload)
	return err
}