package symm

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

const (
	gcmStreamMagic     = "SCG1"
	gcmStreamChunkSize = 32 * 1024
)

var ErrInvalidGCMStreamChunk = errors.New("invalid gcm stream chunk")

func encryptGCMChunk(aead cipher.AEAD, plain, nonce, aad []byte, chunkIndex uint64) []byte {
	chunkNonce := deriveChunkNonce(nonce, chunkIndex)
	return aead.Seal(nil, chunkNonce, plain, aad)
}

func decryptGCMChunk(aead cipher.AEAD, ciphertext, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	chunkNonce := deriveChunkNonce(nonce, chunkIndex)
	return aead.Open(nil, chunkNonce, ciphertext, aad)
}

func encryptGCMChunkedStream(dst io.Writer, src io.Reader, aead cipher.AEAD, nonce, aad []byte) error {
	if _, err := dst.Write([]byte(gcmStreamMagic)); err != nil {
		return err
	}

	buf := make([]byte, gcmStreamChunkSize)
	lenBuf := make([]byte, 4)
	var chunkIndex uint64

	for {
		n, err := src.Read(buf)
		if n > 0 {
			sealed := encryptGCMChunk(aead, buf[:n], nonce, aad, chunkIndex)
			binary.BigEndian.PutUint32(lenBuf, uint32(len(sealed)))
			if _, werr := dst.Write(lenBuf); werr != nil {
				return werr
			}
			if _, werr := dst.Write(sealed); werr != nil {
				return werr
			}
			chunkIndex++
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

func decryptGCMChunkedOrLegacyStream(dst io.Writer, src io.Reader, aead cipher.AEAD, nonce, aad []byte) error {
	header := make([]byte, len(gcmStreamMagic))
	n, err := io.ReadFull(src, header)
	if err != nil {
		if err == io.EOF {
			return nil
		}
		if err != io.ErrUnexpectedEOF {
			return err
		}
		return decryptGCMLegacyBuffered(dst, io.MultiReader(bytes.NewReader(header[:n]), src), aead, nonce, aad)
	}

	if string(header) != gcmStreamMagic {
		return decryptGCMLegacyBuffered(dst, io.MultiReader(bytes.NewReader(header), src), aead, nonce, aad)
	}
	return decryptGCMChunkedStream(dst, src, aead, nonce, aad)
}

func decryptGCMChunkedStream(dst io.Writer, src io.Reader, aead cipher.AEAD, nonce, aad []byte) error {
	lenBuf := make([]byte, 4)
	maxChunkLen := uint32(gcmStreamChunkSize + aead.Overhead())
	var chunkIndex uint64

	for {
		_, err := io.ReadFull(src, lenBuf)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			if err == io.ErrUnexpectedEOF {
				return io.ErrUnexpectedEOF
			}
			return err
		}

		chunkLen := binary.BigEndian.Uint32(lenBuf)
		if chunkLen < uint32(aead.Overhead()) || chunkLen > maxChunkLen {
			return ErrInvalidGCMStreamChunk
		}

		chunk := make([]byte, chunkLen)
		if _, err := io.ReadFull(src, chunk); err != nil {
			if err == io.ErrUnexpectedEOF {
				return io.ErrUnexpectedEOF
			}
			return err
		}

		plain, err := decryptGCMChunk(aead, chunk, nonce, aad, chunkIndex)
		if err != nil {
			return err
		}
		if _, err := dst.Write(plain); err != nil {
			return err
		}
		chunkIndex++
	}
}

func decryptGCMLegacyBuffered(dst io.Writer, src io.Reader, aead cipher.AEAD, nonce, aad []byte) error {
	enc, err := io.ReadAll(src)
	if err != nil {
		return err
	}
	plain, err := aead.Open(nil, nonce, enc, aad)
	if err != nil {
		return err
	}
	_, err = dst.Write(plain)
	return err
}

func deriveChunkNonce(baseNonce []byte, chunkIndex uint64) []byte {
	nonce := make([]byte, len(baseNonce))
	copy(nonce, baseNonce)
	if len(nonce) < 8 {
		return nonce
	}

	var indexBytes [8]byte
	binary.BigEndian.PutUint64(indexBytes[:], chunkIndex)
	off := len(nonce) - 8
	for i := 0; i < 8; i++ {
		nonce[off+i] ^= indexBytes[i]
	}
	return nonce
}