package httpserver

import (
	"bytes"
	"fmt"
	"io"
)

// LimitBuffer buf with limit
type LimitBuffer interface {
	io.Writer
	String() string
	Remain() int
}

// limitBuffer buf with limit
type limitBuffer struct {
	buf    *bytes.Buffer
	limit  int
	remain int
	length int
}

// NewLimitBuffer ..
func NewLimitBuffer(limit int) LimitBuffer {
	buf := bytes.NewBuffer(make([]byte, 0, limit))
	return &limitBuffer{buf: buf, limit: limit, remain: limit}
}

func (b *limitBuffer) Write(p []byte) (n int, err error) {
	defer func() {
		b.length += n
	}()

	// discord
	if b.remain <= 0 {
		return len(p), nil
	}

	// write remain
	if len(p) > b.remain {
		n, err = b.buf.Write(p[:b.remain])
		b.remain -= n

		// alway return all writed length
		n = len(p)
		return
	}

	// write all
	n, err = b.buf.Write(p)
	b.remain -= n
	return
}

// String ..
func (b *limitBuffer) String() string {
	if b.length > b.limit {
		return b.buf.String() + fmt.Sprintf("...(total %dB)", b.length)
	}

	return b.buf.String()
}

// Remain 剩余多少个字节没有写满
func (b *limitBuffer) Remain() int {
	return b.remain
}

type teeReadCloser struct {
	tee io.Reader
	r   io.Closer
}

// Read implement io.Read interface
func (t *teeReadCloser) Read(p []byte) (n int, err error) {
	return t.tee.Read(p)
}

// Close implement io.Close interface
func (t *teeReadCloser) Close() error {
	return t.r.Close()
}

// TeeReadCloser TeeReader with limit
func TeeReadCloser(r io.ReadCloser, w io.Writer) io.ReadCloser {
	return &teeReadCloser{r: r, tee: io.TeeReader(r, w)}
}