package symm

import (
	"crypto/cipher"
	"errors"
	"io"
)

func encryptCFB8(block cipher.Block, data, iv []byte) ([]byte, error) {
	if len(iv) != block.BlockSize() {
		return nil, errors.New("iv length must match block size")
	}
	reg := make([]byte, len(iv))
	copy(reg, iv)
	regView := make([]byte, block.BlockSize())
	streamBlock := make([]byte, block.BlockSize())
	out := make([]byte, len(data))
	head := 0

	for i := 0; i < len(data); i++ {
		buildCFB8Register(regView, reg, head)
		block.Encrypt(streamBlock, regView)
		c := data[i] ^ streamBlock[0]
		out[i] = c
		advanceCFB8Register(reg, &head, c)
	}
	return out, nil
}

func decryptCFB8(block cipher.Block, src, iv []byte) ([]byte, error) {
	if len(iv) != block.BlockSize() {
		return nil, errors.New("iv length must match block size")
	}
	reg := make([]byte, len(iv))
	copy(reg, iv)
	regView := make([]byte, block.BlockSize())
	streamBlock := make([]byte, block.BlockSize())
	out := make([]byte, len(src))
	head := 0

	for i := 0; i < len(src); i++ {
		buildCFB8Register(regView, reg, head)
		block.Encrypt(streamBlock, regView)
		p := src[i] ^ streamBlock[0]
		out[i] = p
		advanceCFB8Register(reg, &head, src[i])
	}
	return out, nil
}

func encryptCFB8Stream(block cipher.Block, dst io.Writer, src io.Reader, iv []byte, decrypt bool) error {
	if len(iv) != block.BlockSize() {
		return errors.New("iv length must match block size")
	}
	reg := make([]byte, len(iv))
	copy(reg, iv)
	regView := make([]byte, block.BlockSize())
	streamBlock := make([]byte, block.BlockSize())
	buf := make([]byte, 32*1024)
	out := make([]byte, 32*1024)
	head := 0

	for {
		n, err := src.Read(buf)
		if n > 0 {
			for i := 0; i < n; i++ {
				buildCFB8Register(regView, reg, head)
				block.Encrypt(streamBlock, regView)
				if decrypt {
					out[i] = buf[i] ^ streamBlock[0]
					advanceCFB8Register(reg, &head, buf[i])
				} else {
					c := buf[i] ^ streamBlock[0]
					out[i] = c
					advanceCFB8Register(reg, &head, c)
				}
			}
			if _, werr := dst.Write(out[:n]); werr != nil {
				return werr
			}
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

func buildCFB8Register(dst, reg []byte, head int) {
	first := len(reg) - head
	copy(dst, reg[head:])
	copy(dst[first:], reg[:head])
}

func advanceCFB8Register(reg []byte, head *int, feedback byte) {
	reg[*head] = feedback
	*head = *head + 1
	if *head == len(reg) {
		*head = 0
	}
}