package main

import (
	"bytes"
	"encoding/json"
)

var streamError = []byte("data: {\"error\":{\"message\":\"Batch stream ended without a completed response\",\"type\":\"batch_error\"}}\n\n")

type events struct {
	pending  []byte
	complete bool
	done     bool
}

func (e *events) feed(body []byte, end bool) []byte {
	e.pending = append(e.pending, body...)
	var output []byte
	for !e.done {
		cut, size := bytes.Index(e.pending, []byte("\n\n")), 2
		if cr := bytes.Index(e.pending, []byte("\r\n\r\n")); cr >= 0 && (cut < 0 || cr < cut) {
			cut, size = cr, 4
		}
		if cut < 0 {
			break
		}
		event := e.pending[:cut+size]
		var data []byte
		for line := range bytes.SplitSeq(e.pending[:cut], []byte("\n")) {
			if bytes.HasPrefix(line, []byte("data:")) {
				if len(data) > 0 {
					data = append(data, '\n')
				}
				data = append(data, bytes.TrimSpace(line[5:])...)
			}
		}
		if bytes.Equal(data, []byte("[DONE]")) {
			e.done = true
			if !e.complete {
				output = append(output, streamError...)
			} else {
				output = append(output, event...)
			}
		} else {
			var chunk struct {
				Choices []struct {
					Finish *string `json:"finish_reason"`
				} `json:"choices"`
			}
			if json.Unmarshal(data, &chunk) == nil {
				for _, choice := range chunk.Choices {
					if choice.Finish != nil && *choice.Finish != "" {
						e.complete = true
					}
				}
			}
			output = append(output, event...)
		}
		e.pending = e.pending[cut+size:]
	}
	if end && !e.done {
		output = append(output, streamError...)
		e.done = true
	}
	return output
}