package proxy

import (
	"bytes"
	"encoding/json"
	"strings"
)

type sseEventTracker struct {
	buf      []byte
	terminal bool
}

func (t *sseEventTracker) Write(p []byte) {
	t.buf = append(t.buf, p...)
	for {
		end, separator := completeSSEEvent(t.buf)
		if end < 0 {
			return
		}
		event := t.buf[:end]
		t.buf = t.buf[end+separator:]
		if terminalSSEEvent(event) {
			t.terminal = true
		}
	}
}

func (t *sseEventTracker) Complete() bool { return len(t.buf) == 0 }

func completeSSEEvent(buf []byte) (int, int) {
	lf := bytes.Index(buf, []byte("\n\n"))
	crlf := bytes.Index(buf, []byte("\r\n\r\n"))
	if crlf >= 0 && (lf < 0 || crlf < lf) {
		return crlf, 4
	}
	if lf >= 0 {
		return lf, 2
	}
	return -1, 0
}

func terminalSSEEvent(event []byte) bool {
	var data strings.Builder
	for _, line := range strings.Split(strings.ReplaceAll(string(event), "\r\n", "\n"), "\n") {
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		if data.Len() > 0 {
			data.WriteByte('\n')
		}
		data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
	}
	payload := data.String()
	if payload == "[DONE]" {
		return true
	}
	var envelope struct {
		Type string `json:"type"`
	}
	if json.Unmarshal([]byte(payload), &envelope) != nil {
		return false
	}
	return envelope.Type == "message_stop" || envelope.Type == "response.completed"
}