package process

import (
	"io"
	"regexp"
	"sync"
	"time"
)

// This regex matches either newline, or one of the "Erase in Line" escape
// sequences:
// [K Clear from cursor to the end of the line
// [1K Clear from cursor to beginning of the line
// [2K Clear entire line
var lineBreakEscapeRE = regexp.MustCompile(`\n|\x1b\[[012]?K`)

// Timestamper inserts timestamps generated by a callback at the start of each
// line, and inside lines when the line is...printed...very...slowly.
type Timestamper struct {
	out        io.Writer
	formatFunc func(t time.Time) string
	timeout    time.Duration

	mu            sync.Mutex
	startOfLine   bool
	lastTimestamp time.Time
	ansiParser    ansiParser
}

// NewTimestamper sets up a Timestamper outputting to an io.Writer w. Timestamps
// are formatted with f, and non-prefix timestamps are inserted on writes that
// occur after timeout (duration since the previous timestamp).
func NewTimestamper(w io.Writer, f func(t time.Time) string, timeout time.Duration) *Timestamper {
	return &Timestamper{
		out:         w,
		formatFunc:  f,
		startOfLine: true,
		timeout:     timeout,
	}
}

// Write writes the given data, plus any additional timestamps necessary, to the
// Timestamper's output. Timestamps are written at the start of each line, and
// possibly again within lines that are written incrementally over a long period
// of time.
func (p *Timestamper) Write(data []byte) (n int, err error) {
	now := time.Now()

	p.mu.Lock()
	defer p.mu.Unlock()

	written := 0
	for len(data) > 0 {
		// If the previously written data ended in the middle of an ANSI code,
		// keep feeding bytes to the ANSI parser (and writing them out)
		// until we're not
		var codeEnd []byte
		for p.ansiParser.insideCode() && len(data) > len(codeEnd) {
			fed := len(codeEnd)
			p.ansiParser.Write(data[fed : fed+1])
			codeEnd = data[:fed+1]
		}
		if len(codeEnd) > 0 {
			n, err := p.out.Write(codeEnd)
			written += n
			if err != nil {
				return written, err
			}
			data = data[len(codeEnd):]
			if len(data) == 0 {
				break
			}
		}

		// When at the start of a line, or after a long enough time
		// use the callback to format a timestamp and write it.
		if p.startOfLine || now.Sub(p.lastTimestamp) > p.timeout {
			if _, err := p.out.Write([]byte(p.formatFunc(now))); err != nil {
				// Note: timestamps are not included in the written byte total.
				return written, err
			}
			p.startOfLine = false
			p.lastTimestamp = now
		}

		// Find the next newline or line-break escape sequence.
		// FindIndex returns a two-element slice denoting the matched range.
		idx := lineBreakEscapeRE.FindIndex(data)
		if idx == nil {
			// Write all remaining data. The line continues into the next Write
			// call.
			p.ansiParser.Write(data)
			n, err := p.out.Write(data)
			written += n
			return written, err
		}

		// Write up to (and including) the newline / breaking sequence, ready to
		// start a new line in the next iteration or Write.
		chunk := data[:idx[1]]
		p.ansiParser.Write(chunk)
		n, err = p.out.Write(chunk)
		written += n
		if err != nil {
			return written, err
		}
		data = data[idx[1]:]
		p.startOfLine = true
	}
	return written, nil
}
