package core

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"slices"
	"strings"
)

type StreamFormat string

const (
	StreamFormatSSE   StreamFormat = "sse"
	StreamFormatEmpty StreamFormat = ""
)

const (
	defaultMaxBufSize = 64 * 1024 // 64KB
)

// defaultStreamDelimiter is the default stream delimiter used to split messages.
const defaultStreamDelimiter = '\n'

// Stream represents a stream of messages sent from a server.
type Stream[T any] struct {
	reader streamReader
	closer io.Closer
}

// StreamOption adapts the behavior of the Stream.
type StreamOption func(*streamOptions)

// WithDelimiter overrides the delimiter for the Stream.
//
// By default, the Stream is newline-delimited.
func WithDelimiter(delimiter string) StreamOption {
	return func(opts *streamOptions) {
		opts.delimiter = delimiter
	}
}

// WithPrefix overrides the prefix for the Stream.
//
// By default, the Stream doesn't have a prefix.
func WithPrefix(prefix string) StreamOption {
	return func(opts *streamOptions) {
		opts.prefix = prefix
	}
}

// WithTerminator overrides the terminator for the Stream.
//
// By default, the Stream terminates on EOF.
func WithTerminator(terminator string) StreamOption {
	return func(opts *streamOptions) {
		opts.terminator = terminator
	}
}

// WithFormat overrides the isSSE flag for the Stream.
//
// By default, the Stream is not SSE.
func WithFormat(format StreamFormat) StreamOption {
	return func(opts *streamOptions) {
		opts.format = format
	}
}

// NewStream constructs a new Stream from the given *http.Response.
func NewStream[T any](response *http.Response, opts ...StreamOption) *Stream[T] {
	options := new(streamOptions)
	for _, opt := range opts {
		opt(options)
	}
	return &Stream[T]{
		reader: newStreamReader(response.Body, options),
		closer: response.Body,
	}
}

// Recv reads a message from the stream, returning io.EOF when
// all the messages have been read.
func (s Stream[T]) Recv() (T, error) {
	var value T
	bytes, err := s.reader.ReadFromStream()
	if err != nil {
		return value, err
	}
	if err := json.Unmarshal(bytes, &value); err != nil {
		return value, err
	}
	return value, nil
}

// Close closes the Stream.
func (s Stream[T]) Close() error {
	return s.closer.Close()
}

// streamReader reads data from a stream.
type streamReader interface {
	ReadFromStream() ([]byte, error)
}

// newStreamReader returns a new streamReader based on the given
// delimiter.
//
// By default, the streamReader uses a simple a *bufio.Reader
// which splits on newlines, and otherwise use a *bufio.Scanner to
// split on custom delimiters.
func newStreamReader(reader io.Reader, options *streamOptions) streamReader {
	if !options.isEmpty() {
		if options.maxBufSize == 0 {
			options.maxBufSize = defaultMaxBufSize
		}
		if options.delimiter == "" {
			options.delimiter = string(defaultStreamDelimiter)
		}
		if options.format == StreamFormatSSE {
			return newSseStreamReader(reader, options)
		}
		return newScannerStreamReader(reader, options)
	}
	return newBufferStreamReader(reader)
}

// BufferStreamReader reads data from a *bufio.Reader, which splits
// on newlines.
type BufferStreamReader struct {
	reader *bufio.Reader
}

func newBufferStreamReader(reader io.Reader) *BufferStreamReader {
	return &BufferStreamReader{
		reader: bufio.NewReader(reader),
	}
}

func (b *BufferStreamReader) ReadFromStream() ([]byte, error) {
	return b.reader.ReadBytes(defaultStreamDelimiter)
}

// ScannerStreamReader reads data from a *bufio.Scanner, which allows for
// configurable delimiters.
type ScannerStreamReader struct {
	scanner *bufio.Scanner
	options *streamOptions
}

func newScannerStreamReader(
	reader io.Reader,
	options *streamOptions,
) *ScannerStreamReader {
	scanner := bufio.NewScanner(reader)
	stream := &ScannerStreamReader{
		scanner: scanner,
		options: options,
	}
	scanner.Split(func(bytes []byte, atEOF bool) (int, []byte, error) {
		if atEOF && len(bytes) == 0 {
			return 0, nil, nil
		}
		n, data, err := stream.parse(bytes)
		if stream.isTerminated(data) {
			return 0, nil, io.EOF
		}
		return n, data, err
	})
	return stream
}

func (s *ScannerStreamReader) ReadFromStream() ([]byte, error) {
	if s.scanner.Scan() {
		return s.scanner.Bytes(), nil
	}
	if err := s.scanner.Err(); err != nil {
		return nil, err
	}
	return nil, io.EOF
}

func (s *ScannerStreamReader) parse(bytes []byte) (int, []byte, error) {
	var startIndex int
	if s.options != nil && s.options.prefix != "" {
		if i := strings.Index(string(bytes), s.options.prefix); i >= 0 {
			startIndex = i + len(s.options.prefix)
		}
	}
	data := bytes[startIndex:]
	delimIndex := strings.Index(string(data), s.options.delimiter)
	if delimIndex < 0 {
		return startIndex + len(data), data, nil
	}
	endIndex := delimIndex + len(s.options.delimiter)
	parsedData := data[:endIndex]
	n := startIndex + endIndex
	return n, parsedData, nil
}

func (s *ScannerStreamReader) isTerminated(bytes []byte) bool {
	if s.options == nil || s.options.terminator == "" {
		return false
	}
	return strings.Contains(string(bytes), s.options.terminator)
}

type streamOptions struct {
	delimiter  string
	prefix     string
	terminator string
	format     StreamFormat
	maxBufSize int
}

func (s *streamOptions) isEmpty() bool {
	return s.delimiter == "" && s.prefix == "" && s.terminator == "" && s.format == StreamFormatEmpty
}

// SseStreamReader reads data from a *bufio.Scanner, which allows for
// configurable delimiters.
type SseStreamReader struct {
	scanner *bufio.Scanner
	options *streamOptions
}

func newSseStreamReader(
	reader io.Reader,
	options *streamOptions,
) *SseStreamReader {
	scanner := bufio.NewScanner(reader)
	stream := &SseStreamReader{
		scanner: scanner,
		options: options,
	}
	scanner.Buffer(make([]byte, slices.Min([]int{4096, options.maxBufSize})), options.maxBufSize)
	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		if atEOF && len(data) == 0 {
			return 0, nil, nil
		}
		if i := strings.Index(string(data), stream.options.delimiter+stream.options.delimiter); i >= 0 {
			return i + 1, data[0:i], nil
		}

		if atEOF || stream.isTerminated(data) {
			return len(data), data, nil
		}
		return 0, nil, nil
	})
	return stream
}

func (s *SseStreamReader) isTerminated(bytes []byte) bool {
	if s.options == nil || s.options.terminator == "" {
		return false
	}
	return strings.Contains(string(bytes), s.options.terminator)
}

func (s *SseStreamReader) ReadFromStream() ([]byte, error) {

	event, err := s.nextEvent()
	if err != nil {
		return nil, err
	}
	return event.data, nil
}

func (s *SseStreamReader) nextEvent() (*SseEvent, error) {

	event := SseEvent{}
	if s.scanner.Scan() {
		rawEvent := s.scanner.Bytes()

		lines := strings.Split(string(rawEvent), s.options.delimiter)
		for _, line := range lines {
			s.parseSseLine([]byte(line), &event)
		}

		if event.size() > s.options.maxBufSize {
			return nil, errors.New("SseStreamReader.ReadFromStream: buffer limit exceeded")
		}
		return &event, nil
	}
	return &event, io.EOF
}

func (s *SseStreamReader) parseSseLine(_bytes []byte, event *SseEvent) error {
	if bytes.HasPrefix(_bytes, sseDataPrefix) {
		if len(event.data) > 0 {
			event.data = append(event.data, s.options.delimiter...)
		}
		event.data = append(event.data, _bytes[len(sseDataPrefix):]...)
	} else if bytes.HasPrefix(_bytes, sseIdPrefix) {
		event.id = append(event.id, _bytes[len(sseIdPrefix):]...)
	} else if bytes.HasPrefix(_bytes, sseEventPrefix) {
		event.event = append(event.event, _bytes[len(sseEventPrefix):]...)
	} else if bytes.HasPrefix(_bytes, sseRetryPrefix) {
		event.retry = append(event.retry, _bytes[len(sseRetryPrefix):]...)
	}
	return nil
}

func (event *SseEvent) size() int {
	return len(event.id) + len(event.data) + len(event.event) + len(event.retry)
}

func (event *SseEvent) String() string {
	return fmt.Sprintf("SseEvent{id: %q, event: %q, data: %q, retry: %q}", event.id, event.event, event.data, event.retry)
}

type SseEvent struct {
	id    []byte
	data  []byte
	event []byte
	retry []byte
}

var (
	sseIdPrefix    = []byte("id: ")
	sseDataPrefix  = []byte("data: ")
	sseEventPrefix = []byte("event: ")
	sseRetryPrefix = []byte("retry: ")
)
