package host

import (
	"bytes"
	"io"
	"os"
	"time"

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

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

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) (*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.workload.Handle().Name))
	go func() {
		stdout, stderr, exitCode, err := w.workload.Exec(resource.WorkloadNodeExecOptions{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
	}
}

func (w *WorkloadHost) Run(cmd *workload.CMD) error {
	o, err := w.Exec(cmd.Command, cmd.Timeout)
	if err == TimeoutErr {
		return errors.Trace(err)
	}

	if err != nil {
		log.Error("run workload failed", zap.Error(err), zap.String("stderr", o.Stderr), zap.String("stdout", o.Stdout), zap.Int("exit code", o.ExitCode))
		return errors.Trace(err)
	}

	if o.ExitCode != 0 {
		log.Error("run workload failed", zap.String("stderr", o.Stderr), zap.String("stdout", o.Stdout), zap.Int("exit code", o.ExitCode))
		return errors.Errorf("run workload failed, ExitCode: %+v", o.ExitCode)
	}
	if o.Stderr != "" {
		logger.Warn("run workload stderr is not empty", zap.String("stderr", o.Stderr))
	}

	return nil
}

// ExecAsync exec cmd and return immediately
func (w *WorkloadHost) ExecAsync(cmd *workload.CMD, timeout time.Duration) error {
	// done := make(chan *command.ExecOutput, 1)
	logger.Info("Execute command", zap.String("command", cmd.Command), zap.Duration("timeout", timeout),
		zap.String("resource name", w.workload.Handle().Name))
	var stdout, stderr bytes.Buffer
	_, err := w.workload.ExecAsync(resource.NopCloser(io.MultiWriter(os.Stdout, &stdout)), resource.NopCloser(io.MultiWriter(os.Stdout, &stderr)), resource.WorkloadNodeExecOptions{Command: cmd.Command})
	if err != nil {
		return err
	}
	return nil
}