package llm

import "strings"

// thinkFilter removes inline reasoning blocks ("<think>…</think>") that some
// OpenAI-compatible backends (e.g. MiniMax-M2.x on vLLM, seen on Dahl) put into
// the normal content stream instead of a separate reasoning field. Tags may be
// split across chunks. Text inside a block is never streamed to the user.
type thinkFilter struct {
	pending string // undecided tail that may be the start of a tag
	inThink bool
	think   strings.Builder // reasoning text of the current (unclosed) block
	out     strings.Builder // visible content
	sawOpen bool
}

const (
	openTag  = "<think>"
	closeTag = "</think>"
	gap      = "\n\n\n"
)

// push consumes a chunk and returns the text that may be shown now.
func (f *thinkFilter) push(s string) string {
	buf := f.pending + s
	f.pending = ""
	var vis strings.Builder
	for buf != "" {
		tag := openTag
		if f.inThink {
			tag = closeTag
		}
		i, n := strings.Index(buf, tag), len(tag)
		// Dahl/vLLM sometimes swallows "</think>" and leaves only a "\n\n\n" gap;
		// treat that gap as the end of the block so the answer still streams live.
		if f.inThink {
			if j := strings.Index(buf, gap); j >= 0 && (i < 0 || j < i) {
				i, n = j, len(gap)
			}
		}
		if i >= 0 {
			f.emit(&vis, buf[:i])
			buf = buf[i+n:]
			if f.inThink {
				f.think.Reset()
			} else {
				f.sawOpen = true
			}
			f.inThink = !f.inThink
			continue
		}
		// a stray closing tag outside a block is dropped
		if !f.inThink {
			if i := strings.Index(buf, closeTag); i >= 0 {
				f.emit(&vis, buf[:i])
				buf = buf[i+len(closeTag):]
				continue
			}
		}
		keep := partialSuffix(buf, tag)
		if f.inThink {
			if k := partialSuffix(buf, gap); k > keep {
				keep = k
			}
		} else {
			if k := partialSuffix(buf, closeTag); k > keep {
				keep = k
			}
		}
		f.emit(&vis, buf[:len(buf)-keep])
		f.pending = buf[len(buf)-keep:]
		buf = ""
	}
	return vis.String()
}

func (f *thinkFilter) emit(vis *strings.Builder, s string) {
	if s == "" {
		return
	}
	if f.inThink {
		f.think.WriteString(s)
		return
	}
	// drop whitespace directly after a closed block at the start of the answer
	if f.out.Len() == 0 {
		s = strings.TrimLeft(s, " \n\r\t")
		if s == "" {
			return
		}
	}
	f.out.WriteString(s)
	vis.WriteString(s)
}

// finish flushes the remainder. If a block was never closed (observed when the
// backend swallows "</think>"), the answer is recovered from the text after the
// first blank-line gap so the user does not end up with nothing.
func (f *thinkFilter) finish() string {
	var vis strings.Builder
	if f.pending != "" {
		p := f.pending
		f.pending = ""
		f.emit(&vis, p)
	}
	if f.inThink {
		t := f.think.String()
		f.inThink = false
		f.think.Reset()
		ans := ""
		if i := strings.Index(t, "\n\n\n"); i >= 0 {
			ans = t[i:]
		} else if f.out.Len() == 0 {
			ans = t // better a visible reasoning than an empty reply
		}
		f.emit(&vis, strings.TrimSpace(ans))
	}
	return vis.String()
}

// partialSuffix returns the length of the longest suffix of s that is a proper prefix of tag.
func partialSuffix(s, tag string) int {
	for n := len(tag) - 1; n > 0; n-- {
		if len(s) >= n && strings.HasSuffix(s, tag[:n]) {
			return n
		}
	}
	return 0
}