import {ref} from 'vue'

class ChunkParser {
    constructor(tags, name = 'inference_rag_') {
        this.name = name
        this.tags = tags
        this.chunks = []
    }

    append(chunk, index) {
        this.chunks.push({index: index, chunk: chunk})
    }

    parse(handlers) {
        const buffer = this.buffer()

        this.tags.forEach(tag => {
            const positions = this.position(buffer, tag)
            let string = null

            if(positions.length > 0) {
                const length = `<|${this.name}${tag}|>`.length
                string = buffer.substring(positions[0] + length, positions[1] || buffer.length)
            }

            handlers[tag]?.(string)
        })
    }

    position(input, tag) {
        const regex = new RegExp(`<\\|${this.name}${tag}\\|>`, 'g')
        return [...input.matchAll(regex)].map(match => match.index)
    }

    buffer() {
        if (!this.chunks.length) return '';

        const [{index: firstIndex}] = this.chunks.slice().sort((a, b) => a.index - b.index);

        return this.chunks
            .slice()
            .sort((a, b) => a.index - b.index)
            .filter(({index}, i) => index === firstIndex + i)
            .map(({chunk}) => chunk)
            .join('');
    }

    reset() {
        this.chunks = []
    }
}

export function useInferenceResponse() {
    const status = ref(null)
    const think = ref(null)
    const content = ref(null)

    const parser = new ChunkParser(['think', 'content', 'fields'])

    function append(chunk, index) {
        parser.append(chunk, index)
        parser.parse({
            think: text => think.value = text,
            content: text => content.value = text,
            fields: () => {
            },
        })
    }

    function setStatus(value) {
        status.value = value
    }

    function reset() {
        think.value = null
        content.value = null
        status.value = null

        parser.reset()
    }

    return {think, content, status, append, reset, setStatus}
}