package ndjson

import (
	"bufio"
	"encoding/json"
	"io"
	"sync"
)

type Writer struct {
	w   io.Writer
	buf *bufio.Writer
	mu  sync.Mutex
}

func NewWriter(w io.Writer) *Writer {
	return &Writer{
		w:   w,
		buf: bufio.NewWriter(w),
	}
}

func (w *Writer) Write(v interface{}) error {
	w.mu.Lock()
	defer w.mu.Unlock()

	data, err := json.Marshal(v)
	if err != nil {
		return err
	}

	if _, err := w.buf.Write(data); err != nil {
		return err
	}

	if err := w.buf.WriteByte('\n'); err != nil {
		return err
	}

	return w.buf.Flush()
}

func (w *Writer) WriteRaw(data []byte) error {
	w.mu.Lock()
	defer w.mu.Unlock()

	if _, err := w.buf.Write(data); err != nil {
		return err
	}

	if err := w.buf.WriteByte('\n'); err != nil {
		return err
	}

	return w.buf.Flush()
}

func (w *Writer) Flush() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.buf.Flush()
}

type StreamEvent struct {
	Type    string      `json:"type"`
	Block   interface{} `json:"block,omitempty"`
	BlockID string      `json:"blockId,omitempty"`
	Chunk   string      `json:"chunk,omitempty"`
	Patch   interface{} `json:"patch,omitempty"`
	Data    interface{} `json:"data,omitempty"`
}

func WriteBlock(w *Writer, block interface{}) error {
	return w.Write(StreamEvent{
		Type:  "block",
		Block: block,
	})
}

func WriteTextChunk(w *Writer, blockID, chunk string) error {
	return w.Write(StreamEvent{
		Type:    "textChunk",
		BlockID: blockID,
		Chunk:   chunk,
	})
}

func WriteUpdateBlock(w *Writer, blockID string, patch interface{}) error {
	return w.Write(StreamEvent{
		Type:    "updateBlock",
		BlockID: blockID,
		Patch:   patch,
	})
}

func WriteResearchComplete(w *Writer) error {
	return w.Write(StreamEvent{
		Type: "researchComplete",
	})
}

func WriteMessageEnd(w *Writer) error {
	return w.Write(StreamEvent{
		Type: "messageEnd",
	})
}

func WriteError(w *Writer, err error) error {
	return w.Write(StreamEvent{
		Type: "error",
		Data: err.Error(),
	})
}