package workloadnode

import (
	"bufio"
	"context"
	"errors"
	"io"
	"os"
	"strings"

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

// Cmd is a wrapper for RemoteCommanderInterruptible.
//
// The implementation of Cmd is almost copied from os/exec.Cmd.
// os/exec.Cmd run cmd by using os.Process and we run cmd
// by using resource.RemoteCommanderInterruptible.
type Cmd struct {
	s string

	//
	Stdout io.Writer

	// Commander is the underlying RemoteCommander, once started.
	Commander *resource.RemoteCommanderInterruptible

	// ExitCode is the exit code of the cmd,
	// available after a call to Wait or Run.
	ExitCode         *int
	ExitDueToCtxDone bool

	// FD to be closed after Wait.
	closeAfterWait []io.Closer
	goroutine      []func() error
	errch          chan error // one send per goroutine
	ctx            context.Context
	waitDone       chan struct{}
}

func Command(s string) *Cmd {
	cmd := &Cmd{s: s}
	return cmd
}

func CommandContext(ctx context.Context, s string) *Cmd {
	if ctx == nil {
		panic("nil Context")
	}
	cmd := Command(s)
	cmd.ctx = ctx
	return cmd
}

func (c *Cmd) setDefaultStdout() error {
	stdoutFlag := os.Getenv("ENDLESS_CMD_STDOUT")
	if strings.ToLower(stdoutFlag) == "stdout" {
		// TODO: Due to https://github.com/pingcap/test-infra/issues/298,
		// we can not use os.Stdout as cmd stdout, and we use Pipe temporarily.
		// cmd.Stdout = os.Stdout
		pr, pw, err := os.Pipe()
		if err != nil {
			return err
		}
		c.goroutine = append(c.goroutine, func() error {
			b := bufio.NewReader(pr)
			for {
				line, err := b.ReadString('\n')
				if err != nil {
					if err == io.EOF {
						break
					}
					log.Info("Redirect cmd stdout failed", zap.Error(err))
					return err
				} else {
					print(line)
				}
			}
			log.Info("Redirect cmd stdout finished")
			return nil
		})
		c.Stdout = pw
		c.closeAfterWait = append(c.closeAfterWait, pw)
	} else {
		f, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
		if err != nil {
			return err
		}
		c.Stdout = f
		c.closeAfterWait = append(c.closeAfterWait, f)
	}
	return nil
}

func (c *Cmd) Start(node resource.WorkloadNode) (err error) {
	if c.Commander != nil {
		return errors.New("remote commander already started")
	}
	options := resource.WorkloadNodeExecOptions{Command: c.s}
	// Stdout is redirected to os.DevNull if Stdout is not set.
	if c.Stdout == nil {
		err = c.setDefaultStdout()
		if err != nil {
			return err
		}
	}
	// Start the command on remote node.
	var stdoutWriteCloser io.WriteCloser
	if wc, ok := c.Stdout.(io.WriteCloser); ok {
		stdoutWriteCloser = wc
	} else {
		// Ensure backward compatibility.
		// Only wrap Stdout when it isn't a WriteCloser.
		stdoutWriteCloser = resource.NopCloser(c.Stdout)
	}
	commander, err := node.ExecAsyncInterruptible(stdoutWriteCloser, options)
	if err != nil {
		return err
	}
	// FIXME: there can be credentials in cmd string. For example,
	// user usually set S3 access key and secret in cmd during br restore.
	log.Info("Start remote command",
		zap.String("cmd", c.s),
		zap.String("nodename", node.Handle().Name))
	c.Commander = &commander

	if len(c.goroutine) > 0 {
		c.errch = make(chan error, len(c.goroutine))
		for _, fn := range c.goroutine {
			go func(fn func() error) {
				c.errch <- fn()
			}(fn)
		}
	}

	if c.ctx != nil {
		c.waitDone = make(chan struct{})
		go func() {
			select {
			case <-c.ctx.Done():
				log.Info("Stop remote command since context is done",
					zap.String("cmd", c.s),
					zap.String("nodename", node.Handle().Name))
				c.ExitDueToCtxDone = true
				if err := c.Stop(); err != nil {
					log.Warn("Stop remote command failed",
						zap.Error(err))
				}
			case <-c.waitDone:
			}
		}()
	}
	return nil
}

func (c *Cmd) Wait() error {
	if c.Commander == nil {
		return errors.New("remote commander not started")
	}
	if c.ExitCode != nil {
		return errors.New("wait was already called")
	}
	commanderErr := (*c.Commander).Wait()
	if c.waitDone != nil {
		close(c.waitDone)
	}
	c.ExitCode = &commanderErr.ExitCode
	err := commanderErr.Err
	log.Info("Remote command finished",
		zap.String("cmd", c.s),
		zap.Int("exitcode", *c.ExitCode),
		zap.Error(err))

	// Wait for goroutines finished.
	var copyError error
	for range c.goroutine {
		if err := <-c.errch; err != nil && copyError == nil {
			copyError = err
		}
	}
	// Add a log here so that we know if any goroutine is blocking.
	log.Info("All cmd goroutine are finished", zap.Error(copyError))

	c.closeDescriptors(c.closeAfterWait)

	if err != nil {
		return err
	}
	return copyError
}

func (c *Cmd) closeDescriptors(closers []io.Closer) {
	for _, fd := range closers {
		fd.Close()
	}
}

func (c *Cmd) Run(node resource.WorkloadNode) error {
	err := c.Start(node)
	if err != nil {
		return err
	}
	return c.Wait()
}

func (c *Cmd) Stop() error {
	if c.Commander == nil {
		return errors.New("remote commander not started")
	}
	return (*c.Commander).Interrupt()
}

func (c *Cmd) String() string {
	return c.s
}