package vmess

import (
	"crypto/cipher"
	"encoding/binary"
	"errors"
	"io"

	"github.com/Dreamacro/clash/common/pool"
)

type aeadWriter struct {
	io.Writer
	cipher.AEAD
	nonce [32]byte
	count uint16
	iv    []byte
}

func newAEADWriter(w io.Writer, aead cipher.AEAD, iv []byte) *aeadWriter {
	return &aeadWriter{Writer: w, AEAD: aead, iv: iv}
}

func (w *aeadWriter) Write(b []byte) (n int, err error) {
	buf := pool.Get(pool.RelayBufferSize)
	defer pool.Put(buf)
	length := len(b)
	for {
		if length == 0 {
			break
		}
		readLen := chunkSize - w.Overhead()
		if length < readLen {
			readLen = length
		}
		payloadBuf := buf[lenSize : lenSize+chunkSize-w.Overhead()]
		copy(payloadBuf, b[n:n+readLen])

		binary.BigEndian.PutUint16(buf[:lenSize], uint16(readLen+w.Overhead()))
		binary.BigEndian.PutUint16(w.nonce[:2], w.count)
		copy(w.nonce[2:], w.iv[2:12])

		w.Seal(payloadBuf[:0], w.nonce[:w.NonceSize()], payloadBuf[:readLen], nil)
		w.count++

		_, err = w.Writer.Write(buf[:lenSize+readLen+w.Overhead()])
		if err != nil {
			break
		}
		n += readLen
		length -= readLen
	}
	return
}

type aeadReader struct {
	io.Reader
	cipher.AEAD
	nonce   [32]byte
	buf     []byte
	offset  int
	iv      []byte
	sizeBuf []byte
	count   uint16
}

func newAEADReader(r io.Reader, aead cipher.AEAD, iv []byte) *aeadReader {
	return &aeadReader{Reader: r, AEAD: aead, iv: iv, sizeBuf: make([]byte, lenSize)}
}

func (r *aeadReader) Read(b []byte) (int, error) {
	if r.buf != nil {
		n := copy(b, r.buf[r.offset:])
		r.offset += n
		if r.offset == len(r.buf) {
			pool.Put(r.buf)
			r.buf = nil
		}
		return n, nil
	}

	_, err := io.ReadFull(r.Reader, r.sizeBuf)
	if err != nil {
		return 0, err
	}

	size := int(binary.BigEndian.Uint16(r.sizeBuf))
	if size > maxSize {
		return 0, errors.New("Buffer is larger than standard")
	}

	buf := pool.Get(size)
	_, err = io.ReadFull(r.Reader, buf[:size])
	if err != nil {
		pool.Put(buf)
		return 0, err
	}

	binary.BigEndian.PutUint16(r.nonce[:2], r.count)
	copy(r.nonce[2:], r.iv[2:12])

	_, err = r.Open(buf[:0], r.nonce[:r.NonceSize()], buf[:size], nil)
	r.count++
	if err != nil {
		return 0, err
	}
	realLen := size - r.Overhead()
	n := copy(b, buf[:realLen])
	if len(b) >= realLen {
		pool.Put(buf)
		return n, nil
	}

	r.offset = n
	r.buf = buf[:realLen]
	return n, nil
}