package proxy

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

type sseEventTracker struct {
	buf           []byte
	recognized    bool
	terminal      bool
	afterTerminal 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 t.terminal {
			t.afterTerminal = true
		}
		recognized, terminal := classifySSEEvent(event)
		t.recognized = t.recognized || recognized
		if terminal {
			t.terminal = true
		}
	}
}

func (t *sseEventTracker) Complete() bool {
	return len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal)
}

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

func classifySSEEvent(event []byte) (recognized, terminal bool) {
	var data strings.Builder
	normalized := strings.ReplaceAll(string(event), "\r\n", "\n")
	for _, line := range strings.Split(strings.ReplaceAll(normalized, "\r", "\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, true
	}
	var envelope struct {
		Type       string          `json:"type"`
		Choices    json.RawMessage `json:"choices"`
		Candidates []struct {
			FinishReason string `json:"finishReason"`
		} `json:"candidates"`
	}
	if json.Unmarshal([]byte(payload), &envelope) != nil {
		return false, false
	}
	if envelope.Type != "" {
		return true, envelope.Type == "message_stop" || envelope.Type == "response.completed"
	}
	if envelope.Choices != nil {
		return true, false
	}
	if envelope.Candidates != nil {
		for _, candidate := range envelope.Candidates {
			if candidate.FinishReason != "" {
				return true, true
			}
		}
		return true, false
	}
	return false, false
}