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.
//
// Delta-only compression: once initial compression has run, subsequent compressions
// only process the new (delta) content using context-independent layers (entropy,
// ngram), then append the result to the existing compressed output. This avoids
// re-compressing the entire accumulated content each time the threshold is crossed.
type StreamCompressor struct {
	mu                sync.RWMutex
	raw               []string // accumulated raw segments
	compressed        string   // accumulated 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
	lastCompressedIdx int      // raw[] index up to which content has been compressed
	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. Only the new delta (content added
// since the last compression) is compressed using context-independent layers,
// then merged with the existing compressed output.
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 delta (segments not yet compressed) under lock.
	var deltaCopy string
	var rawLen int
	if shouldCompress {
		rawLen = len(sc.raw)
		deltaIdx := sc.lastCompressedIdx
		if deltaIdx < rawLen {
			deltaCopy = strings.Join(sc.raw[deltaIdx:], "\n")
		} else {
			// Nothing new to compress; don't start a goroutine.
			sc.compressing = false
			shouldCompress = false
		}
	}
	sc.mu.Unlock()

	if shouldCompress {
		sc.wg.Add(1)
		go sc.backgroundCompress(deltaCopy, rawLen)
	}
}

// backgroundCompress runs delta-only compression on the given content and
// merges the result with the existing compressed output.
func (sc *StreamCompressor) backgroundCompress(delta string, compressedUpTo int) {
	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:
	}

	// Compress only the delta content using the user's options.
	compressed, deltaStats := Compress(delta, sc.opts...)

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

	// Merge: accumulate compressed output rather than replacing it.
	if sc.compressed == "" {
		sc.compressed = compressed
	} else {
		sc.compressed = sc.compressed + "\n" + compressed
	}

	// Accumulate stats across delta compressions.
	sc.stats.OriginalTokens += deltaStats.OriginalTokens
	sc.stats.FinalTokens += deltaStats.FinalTokens
	sc.stats.TokensSaved += deltaStats.TokensSaved
	if sc.stats.OriginalTokens > 0 {
		sc.stats.ReductionPercent = float64(sc.stats.TokensSaved) / float64(sc.stats.OriginalTokens) * 100
	}

	sc.lastCompressedIdx = compressedUpTo
	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()
}

// Reset clears all accumulated content and resets the compressor to its
// initial state, allowing it to be reused.
func (sc *StreamCompressor) Reset() {
	sc.mu.Lock()
	defer sc.mu.Unlock()

	sc.raw = nil
	sc.compressed = ""
	sc.stats = Stats{}
	sc.dirty = false
	sc.compressing = false
	sc.lastCompressedIdx = 0
}

// 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())
}