package pipeline

import (
	"regexp"
	"strings"
	"time"
	"unicode/utf8"
)

type TextChunker struct {
	MinChars     int
	MaxChars     int
	StartAfter   int
	FirstTimeout time.Duration
	Punctuation  bool
	buf          strings.Builder
	firstAt      time.Time
}

func NewTextChunker(minChars, maxChars, startAfter int, firstTimeout time.Duration, punctuation bool) *TextChunker {
	if minChars <= 0 {
		minChars = 30
	}
	if maxChars < minChars {
		maxChars = 160
	}
	if startAfter <= 0 {
		startAfter = 60
	}
	if firstTimeout <= 0 {
		firstTimeout = 1200 * time.Millisecond
	}
	return &TextChunker{MinChars: minChars, MaxChars: maxChars, StartAfter: startAfter, FirstTimeout: firstTimeout, Punctuation: punctuation}
}

func (c *TextChunker) Add(delta string, final bool) []string {
	delta = normalizeForSpeech(delta)
	if delta != "" {
		if c.buf.Len() == 0 {
			c.firstAt = time.Now()
		}
		c.buf.WriteString(delta)
	}
	var out []string
	for {
		s := strings.TrimSpace(c.buf.String())
		if s == "" {
			c.buf.Reset()
			return out
		}
		emitAt := c.emitIndex(s, final)
		if emitAt <= 0 {
			return out
		}
		chunk := strings.TrimSpace(s[:emitAt])
		out = append(out, chunk)
		rest := strings.TrimSpace(s[emitAt:])
		c.buf.Reset()
		c.buf.WriteString(rest)
		if rest == "" {
			return out
		}
	}
}

func (c *TextChunker) Flush() []string { return c.Add("", true) }

func (c *TextChunker) emitIndex(s string, final bool) int {
	if final {
		return len(s)
	}
	if c.Punctuation && utf8.RuneCountInString(s) >= c.MinChars {
		if idx := lastSentenceBoundary(s, c.MaxChars); idx > 0 {
			return idx
		}
	}
	if utf8.RuneCountInString(s) >= c.MaxChars {
		return byteIndexByRunes(s, c.MaxChars)
	}
	if utf8.RuneCountInString(s) >= c.StartAfter && time.Since(c.firstAt) >= c.FirstTimeout {
		return len(s)
	}
	return 0
}

func lastSentenceBoundary(s string, maxRunes int) int {
	limit := byteIndexByRunes(s, maxRunes)
	if limit <= 0 || limit > len(s) {
		limit = len(s)
	}
	last := -1
	for i, r := range s[:limit] {
		if r == '.' || r == '?' || r == '!' || r == '…' {
			last = i + len(string(r))
		}
	}
	return last
}

func byteIndexByRunes(s string, n int) int {
	if n <= 0 {
		return 0
	}
	i := 0
	for pos := range s {
		if i == n {
			return pos
		}
		i++
	}
	return len(s)
}

var markdownPattern = regexp.MustCompile(`[*_` + "`" + `#>\[\]]`)

func normalizeForSpeech(s string) string {
	s = markdownPattern.ReplaceAllString(s, "")
	s = strings.ReplaceAll(s, "\n", " ")
	return s
}