stream.go

package tok

import (
	"strings"
	"sync"

	"github.com/GrayCodeAI/tok/internal/core"
)

// StreamCompressor maintains a background-compressed version of accumulating content.
// As new content is appended, it re-compresses in the background so a compressed
// snapshot is always available without blocking.
type StreamCompressor struct {
	mu          sync.RWMutex
	raw         []string  // accumulated raw segments
	compressed  string    // latest compressed output (ready to read)
	stats       Stats     // stats from last compression
	opts        []Option  // compression options
	threshold   int       // re-compress when raw tokens exceed this
	dirty       bool      // new content since last compression
	compressing bool      // background compression in progress
	done        chan struct{}
	wg          sync.WaitGroup // tracks in-progress compression goroutines
}

// NewStreamCompressor creates a background compressor that keeps compressed
// output ready at all times. Threshold is the token count that triggers
// background re-compression. If threshold <= 0, it defaults to 2000 tokens.
func NewStreamCompressor(threshold int, opts ...Option) *StreamCompressor {
	if threshold <= 0 {
		threshold = 2000
	}
	return &StreamCompressor{
		threshold: threshold,
		opts:      opts,
		done:      make(chan struct{}),
	}
}

// Append adds new content. If accumulated tokens exceed the threshold,
// triggers background re-compression.
func (sc *StreamCompressor) Append(content string) {
	if content == "" {
		return
	}

	sc.mu.Lock()
	sc.raw = append(sc.raw, content)
	sc.dirty = true
	shouldCompress := !sc.compressing && sc.tokenCountLocked() >= sc.threshold
	if shouldCompress {
		sc.compressing = true
	}
	// Take a snapshot of the raw content under lock for the goroutine.
	var rawCopy string
	if shouldCompress {
		rawCopy = sc.joinRawLocked()
	}
	sc.mu.Unlock()

	if shouldCompress {
		sc.wg.Add(1)
		go sc.backgroundCompress(rawCopy)
	}
}

// backgroundCompress runs compression on the given raw content and updates
// the compressed snapshot when done.
func (sc *StreamCompressor) backgroundCompress(raw string) {
	defer sc.wg.Done()

	// Check if we've been shut down before starting.
	select {
	case <-sc.done:
		sc.mu.Lock()
		sc.compressing = false
		sc.mu.Unlock()
		return
	default:
	}

	compressed, stats := Compress(raw, sc.opts...)

	sc.mu.Lock()
	defer sc.mu.Unlock()

	sc.compressed = compressed
	sc.stats = stats
	sc.dirty = false
	sc.compressing = false
}

// Snapshot returns the current compressed output without blocking.
// If compression hasn't run yet, returns the raw content joined.
func (sc *StreamCompressor) Snapshot() (string, Stats) {
	sc.mu.RLock()
	defer sc.mu.RUnlock()

	if sc.compressed == "" {
		return sc.joinRawLocked(), Stats{}
	}
	return sc.compressed, sc.stats
}

// Raw returns all accumulated raw content.
func (sc *StreamCompressor) Raw() string {
	sc.mu.RLock()
	defer sc.mu.RUnlock()
	return sc.joinRawLocked()
}

// TokenCount returns estimated token count of raw content.
func (sc *StreamCompressor) TokenCount() int {
	sc.mu.RLock()
	defer sc.mu.RUnlock()
	return sc.tokenCountLocked()
}

// Close shuts down the background compressor and waits for any in-progress
// compression to finish.
func (sc *StreamCompressor) Close() {
	close(sc.done)
	sc.wg.Wait()
}

// joinRawLocked joins raw segments. Caller must hold at least a read lock.
func (sc *StreamCompressor) joinRawLocked() string {
	return strings.Join(sc.raw, "\n")
}

// tokenCountLocked estimates the token count of raw content.
// Caller must hold at least a read lock.
func (sc *StreamCompressor) tokenCountLocked() int {
	return core.EstimateTokens(sc.joinRawLocked())
}