package anyreader

import (
	"git.milar.in/milarin/ds"
	"git.milar.in/milarin/slices"
)

type Reader[T any] struct {
	buf     ds.Stack[T]
	indices ds.Stack[uint64]
	index   uint64
	src     func() (T, error)
}

func NewReaderFromSlice[T any](s []T) *Reader[T] {
	return NewReaderFromErrorFunc(sliceToFunc(s))
}

func NewReaderFromFunc[T any](src func() T) *Reader[T] {
	return NewReaderFromErrorFunc(func() (T, error) { return src(), nil })
}

func NewReaderFromErrorFunc[T any](src func() (T, error)) *Reader[T] {
	return &Reader[T]{
		src:     src,
		buf:     ds.NewArrayStack[T](),
		index:   0,
		indices: ds.NewArrayStack[uint64](),
	}
}

func (r *Reader[T]) Read() (T, error) {
	v, err := r.src()
	if err == nil {
		r.buf.Push(v)
		r.index++
	}
	return v, err
}

func (r *Reader[T]) Unread() error {
	if r.buf.Empty() {
		return ErrNothingToUnread.New()
	}

	v := r.buf.Pop()
	r.index--

	returned := false
	oldSrc := r.src
	r.src = func() (T, error) {
		if returned {
			return oldSrc()
		}

		returned = true
		return v, nil
	}

	return nil
}

func (r *Reader[T]) UnreadN(n int) error {
	for i := 0; i < n; i++ {
		err := r.Unread()
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *Reader[T]) Peek() (T, error) {
	value, err := r.Read()
	if err != nil {
		return *new(T), err
	}

	if err := r.Unread(); err != nil {
		return *new(T), err
	}

	return value, nil
}

func (r *Reader[T]) ReadWhile(f ...func(T) bool) ([]T, error) {
	res := make([]T, 0, 10)

	var value T
	var err error
	for value, err = r.Read(); err == nil && findFirstTrue(value, f); value, err = r.Read() {
		res = append(res, value)
	}

	return res, err
}

func (r *Reader[T]) ReadUntil(f ...func(T) bool) ([]T, error) {
	return r.ReadWhile(func(v T) bool { return !findFirstTrue(v, f) })
}

func (r *Reader[T]) SkipUntil(f ...func(T) bool) error {
	_, err := r.ReadUntil(f...)
	if err != nil {
		return err
	}
	return r.Unread()
}

func (r *Reader[T]) SkipWhile(f ...func(T) bool) error {
	_, err := r.ReadWhile(f...)
	if err != nil {
		return err
	}
	return r.Unread()
}

func (r *Reader[T]) Expect(f ...func(T) bool) (bool, error) {
	value, err := r.Read()
	if err != nil {
		return false, err
	}
	return findFirstTrue(value, f), nil
}

func (r *Reader[T]) Push() {
	r.indices.Push(r.index)
}

func (r *Reader[T]) Pop() ([]T, error) {
	if r.indices.Empty() {
		return nil, ErrPopFailed.New()
	}

	lastIndex := r.indices.Pop()
	currentIndex := r.index
	if lastIndex < currentIndex {
		values := make([]T, 0, int(currentIndex-lastIndex))
		for i := 0; i < int(currentIndex-lastIndex); i++ {
			err := r.Unread()
			if err != nil {
				return nil, err
			}

			value, err := r.Peek()
			if err != nil {
				return nil, err
			}

			values = append(values, value)
		}
		return slices.Reverse(values), nil
	} else if lastIndex > currentIndex {
		values := make([]T, 0, int(lastIndex-currentIndex))
		for i := 0; i < int(lastIndex-currentIndex); i++ {
			value, err := r.Read()
			if err != nil {
				return nil, err
			}
			values = append(values, value)
		}
		return values, nil
	}

	return []T{}, nil
}