package stream

import "strings"

type Parser struct {
	inThink bool
	buffer  string
}

const (
	thinkOpenTag  = "<think>"
	thinkCloseTag = "</think>"
)

func (p *Parser) Accept(delta string) (visible string, reasoning string) {
	p.buffer += delta
	for p.buffer != "" {
		if p.inThink {
			idx := strings.Index(p.buffer, thinkCloseTag)
			if idx >= 0 {
				reasoning += p.buffer[:idx]
				p.buffer = p.buffer[idx+len(thinkCloseTag):]
				p.inThink = false
				continue
			}
			keep := tagPrefixSuffixLen(p.buffer, thinkCloseTag)
			if len(p.buffer) > keep {
				reasoning += p.buffer[:len(p.buffer)-keep]
				p.buffer = p.buffer[len(p.buffer)-keep:]
			}
			return visible, reasoning
		}

		idx := strings.Index(p.buffer, thinkOpenTag)
		if idx >= 0 {
			visible += p.buffer[:idx]
			p.buffer = p.buffer[idx+len(thinkOpenTag):]
			p.inThink = true
			continue
		}
		keep := tagPrefixSuffixLen(p.buffer, thinkOpenTag)
		if len(p.buffer) > keep {
			visible += p.buffer[:len(p.buffer)-keep]
			p.buffer = p.buffer[len(p.buffer)-keep:]
		}
		return visible, reasoning
	}
	return visible, reasoning
}

func (p *Parser) Flush() (visible string, reasoning string) {
	if p.inThink {
		reasoning = p.buffer
	} else {
		visible = p.buffer
	}
	p.buffer = ""
	p.inThink = false
	return visible, reasoning
}

func tagPrefixSuffixLen(text, tag string) int {
	limit := len(tag) - 1
	if len(text) < limit {
		limit = len(text)
	}
	for i := limit; i > 0; i-- {
		if strings.HasPrefix(tag, text[len(text)-i:]) {
			return i
		}
	}
	return 0
}