package server

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/daniel/cursor-adapter/internal/types"
)

// SSEWriter 封裝 http.ResponseWriter 用於 SSE streaming。
type SSEWriter struct {
	w     http.ResponseWriter
	flush http.Flusher
}

// NewSSEWriter 建立 SSEWriter，設定必要的 headers。
func NewSSEWriter(w http.ResponseWriter) *SSEWriter {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	flusher, _ := w.(http.Flusher)

	return &SSEWriter{w: w, flush: flusher}
}

// WriteChunk 寫入一個 SSE chunk。
func (s *SSEWriter) WriteChunk(chunk types.ChatCompletionChunk) error {
	data, err := json.Marshal(chunk)
	if err != nil {
		return fmt.Errorf("marshal chunk: %w", err)
	}
	fmt.Fprintf(s.w, "data: %s\n\n", data)
	if s.flush != nil {
		s.flush.Flush()
	}
	return nil
}

// WriteDone 寫入 SSE 結束標記。
func (s *SSEWriter) WriteDone() {
	fmt.Fprint(s.w, "data: [DONE]\n\n")
	if s.flush != nil {
		s.flush.Flush()
	}
}

// WriteError 寫入 SSE 格式的錯誤。
func (s *SSEWriter) WriteError(errMsg string) {
	stopReason := "stop"
	chunk := types.NewChatCompletionChunk("error", 0, "", types.Delta{Content: &errMsg})
	chunk.Choices[0].FinishReason = &stopReason
	s.WriteChunk(chunk)
	s.WriteDone()
}
