package tools

import (
	"context"
	"io"
)

// Progress is where a long-running tool reports output it has produced but not
// finished producing. It carries deltas, not the whole output so far: the
// consumer accumulates if it needs to, and a front end that only wants to show
// the last line does not have to diff two strings to find it.
//
// It may be called from more than one goroutine — a command's stdout and
// stderr are copied by separate ones — so an implementation must be safe for
// concurrent use.
type Progress func(chunk string)

type progressKey struct{}

// WithProgress attaches a sink to ctx. The agent loop installs one around each
// tool call so the tool does not have to know who is watching, or whether
// anyone is.
func WithProgress(ctx context.Context, p Progress) context.Context {
	if p == nil {
		return ctx
	}
	return context.WithValue(ctx, progressKey{}, p)
}

// ProgressFrom returns the sink on ctx, or nil when nothing is watching.
// A tool that finds nothing here should do no extra work at all: the cost of
// streaming is only worth paying when it is going somewhere.
func ProgressFrom(ctx context.Context) Progress {
	p, _ := ctx.Value(progressKey{}).(Progress)
	return p
}

// progressWriter feeds an io.Writer chain into a sink.
type progressWriter struct{ p Progress }

func (w progressWriter) Write(b []byte) (int, error) {
	if len(b) > 0 {
		w.p(string(b))
	}
	return len(b), nil
}

// tee returns w, plus a copy into the sink when there is one. A nil sink gives
// back exactly what was passed in, so the no-observer path allocates nothing
// and copies nothing.
func tee(w io.Writer, p Progress) io.Writer {
	if p == nil {
		return w
	}
	return io.MultiWriter(w, progressWriter{p})
}