package stream

import (
	"bufio"
	"bytes"
	"io"
	"net/http"
	"strings"
)

// SSEEvent is a single SSE data frame.
type SSEEvent struct {
	Data string // the JSON payload of the `data:` line
	Done bool   // true when the payload is [DONE]
}

// ReadSSE reads SSE frames from r, calling fn for each `data:` line.
// It is used both for reading upstream streams and, via a pipe, for writing
// converted streams to the client.
func ReadSSE(r io.Reader, fn func(SSEEvent) error) error {
	br := bufio.NewReader(r)
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
		line = trimCRLF(line)
		if !bytes.HasPrefix([]byte(line), []byte("data:")) {
			continue
		}
		data := line[len("data:"):]
		data = strings.TrimPrefix(data, " ")
		if len(data) == 0 {
			continue
		}
		ev := SSEEvent{Data: data, Done: bytes.Equal([]byte(data), []byte("[DONE]"))}
		if err := fn(ev); err != nil {
			return err
		}
	}
}

// Write writes an SSE event to w and flushes it.
func Write(w http.ResponseWriter, ev SSEEvent) error {
	if ev.Done {
		if _, err := io.WriteString(w, "data: [DONE]\n\n"); err != nil {
			return err
		}
	} else {
		if _, err := io.WriteString(w, "data: "+ev.Data+"\n\n"); err != nil {
			return err
		}
	}
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
	return nil
}

// WriteRaw writes a raw SSE frame string (with trailing newlines) and flushes.
func WriteRaw(w http.ResponseWriter, frame []byte) error {
	if _, err := w.Write(frame); err != nil {
		return err
	}
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
	return nil
}

func trimCRLF(s string) string {
	for len(s) > 0 && (s[len(s)-1] == '\n' || s[len(s)-1] == '\r') {
		s = s[:len(s)-1]
	}
	return s
}