package host

import (
	"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/kafka"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
)

type KafkaHost struct {
	kafka resource.KafkaCluster
}

func NewKafkaHost(kafka resource.KafkaCluster) *KafkaHost {
	return &KafkaHost{
		kafka: kafka,
	}
}

func (k *KafkaHost) Exec(cmd string, timeout time.Duration) (*command.ExecOutput, error) {
	done := make(chan *command.ExecOutput, 1)
	logger.Info("Execute command", zap.String("command", cmd))
	go func() {
		stdout, stderr, exitCode, err := k.kafka.Exec(resource.KafkaClusterExecOptions{Target: resource.Kafka, 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 (k *KafkaHost) Run(cmd *kafka.CMD) error {
	o, err := k.Exec(cmd.Command, cmd.Timeout)
	if err == TimeoutErr {
		return err
	}

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

	if o.ExitCode != 0 {
		log.Error("run kafka cmd 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 kafka cmd topic", zap.String("stderr", o.Stderr))
	}

	return nil
}

// ConsumerHost is for kafka consumer, treat it as a workload should be enough.
type ConsumerHost struct {
	WorkloadHost
}

func NewKafkaConsumerHost(consumer resource.WorkloadNode) *ConsumerHost {
	return &ConsumerHost{WorkloadHost{workload: consumer}}
}