package blockmode

import (
	"fmt"

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/internal/subtle"
	"xdx.jelly/xgcl/internal/xor"
)

// Counter (CTR) mode.

// CTR converts a block cipher into a stream cipher by
// repeatedly encrypting an incrementing counter and
// xoring the resulting stream of data with the input.

// See NIST SP 800-38A, pp 13-15

type ctr struct {
	b       EcbBlockMode
	ctr     []byte //counter
	out     []byte // out[outUsed:] is the unused encrypted counters.
	outUsed int
}

const streamBufferSize = 4096

// ctrAble is an interface implemented by ciphers that have a specific optimized
// implementation of CTR, like crypto/aes. NewCTR will check for this interface
// and return the specific Stream if found.
type ctrAble interface {
	NewCTR(iv []byte) (TernaryStream, error)
}

// NewCTR returns a Stream which encrypts/decrypts using the given Block in
// counter mode. The length of iv must be the same as the Block's block size.
func NewCTR(block EcbBlockMode, iv []byte) (TernaryStream, error) {
	if ctr, ok := block.(ctrAble); ok {
		return ctr.NewCTR(iv)
	}
	if len(iv) != block.BlockSize() {
		return nil, gerrors.WithAnnotatingf(ErrInvalidIV, "input IV length(%d) must be %d bytes", len(iv), block.BlockSize())
	}
	bufSize := streamBufferSize
	if bufSize < block.BlockSize() {
		bufSize = block.BlockSize()
	}
	return &ctr{
		b:       block,
		ctr:     dup(iv),
		out:     make([]byte, 0, bufSize),
		outUsed: 0,
	}, nil
}

func (x *ctr) reset(iv []byte) error {
	if len(iv) != x.b.BlockSize() {
		return gerrors.WithAnnotatingf(ErrInvalidIV, "input IV length(%d) must be %d bytes", len(iv), x.b.BlockSize())
	}
	copy(x.ctr, iv)
	x.out = x.out[:0]
	x.outUsed = 0
	return nil
}

// refill 填充x.out.
func (x *ctr) refill() {
	begin := len(x.out) - x.outUsed
	copy(x.out, x.out[x.outUsed:])
	x.out = x.out[:cap(x.out)]
	bs := x.b.BlockSize()
	end := begin
	for end <= len(x.out)-bs {
		copy(x.out[end:], x.ctr)
		end += bs
		for i := len(x.ctr) - 1; i >= 0; i-- {
			x.ctr[i]++
			if x.ctr[i] != 0 {
				break
			}
		}
	}
	_ = x.b.EcbEncCryptBlocks(x.out[begin:end], x.out[begin:end])
	x.out = x.out[:end]
	x.outUsed = 0
}

func (x *ctr) XORKeyStream(dst, src []byte) {
	if len(dst) < len(src) {
		// By definition of cipher.Stream, if len(dst) < len(src), XORKeyStream should panic.
		panic(fmt.Sprintf("length of output buf(%d) less than the input(%d)", len(dst), len(src)))
	}
	if subtle.InexactOverlap(dst[:len(src)], src) {
		panic("ctr.XORKeyStream: dst must be exact overlap with src or non-overlap with src")
	}
	for len(src) > 0 {
		if x.outUsed >= len(x.out)-x.b.BlockSize() {
			x.refill()
		}
		n := xor.XorBytes(dst, src, x.out[x.outUsed:])
		dst = dst[n:]
		src = src[n:]
		x.outUsed += n
	}
}

var _ TernaryStream = &ctr{}

// EncryptInit implements TernaryStream
func (x *ctr) EncryptInit(iv []byte) error {
	return x.reset(iv)
}

// Encrypt implements TernaryStream
func (x *ctr) Encrypt(dst []byte, in []byte) ([]byte, error) {
	ret, out := sliceForAppend(dst, len(in))
	x.XORKeyStream(out, in)
	return ret, nil
}

// EncryptFinal implements TernaryStream
func (x *ctr) EncryptFinal(dst []byte) ([]byte, error) {
	return dst, nil
}

// EncryptUpdate implements TernaryStream
func (x *ctr) EncryptUpdate(dst []byte, in []byte) ([]byte, error) {
	return x.Encrypt(dst, in)
}

// Decrypt implements TernaryStream
func (x *ctr) Decrypt(dst []byte, in []byte) ([]byte, error) {
	return x.Encrypt(dst, in)
}

// DecryptFinal implements TernaryStream
func (x *ctr) DecryptFinal(dst []byte) ([]byte, error) {
	return dst, nil
}

// DecryptInit implements TernaryStream
func (x *ctr) DecryptInit(iv []byte) error {
	return x.reset(iv)
}

// DecryptUpdate implements TernaryStream
func (x *ctr) DecryptUpdate(dst []byte, in []byte) ([]byte, error) {
	return x.Encrypt(dst, in)
}