package paddingx

import (
	"bytes"
	"errors"
	"strings"
)

const (
	PKCS5    = "PKCS5"
	PKCS7    = "PKCS7"
	ZERO     = "ZERO"
	ANSIX923 = "ANSIX923"
)

func Pad(data []byte, blockSize int, mode string) ([]byte, error) {
	if blockSize <= 0 {
		return nil, errors.New("block size must be greater than zero")
	}
	switch normalizeMode(mode) {
	case "", PKCS7:
		return PKCS7Padding(data, blockSize), nil
	case PKCS5:
		// Compatibility mode: historically PKCS5 was used generically in this project.
		return PKCS7Padding(data, blockSize), nil
	case ZERO:
		return zeroPadding(data, blockSize), nil
	case ANSIX923:
		return ansiX923Padding(data, blockSize), nil
	default:
		return nil, errors.New("padding type not supported")
	}
}

func Unpad(data []byte, blockSize int, mode string) ([]byte, error) {
	if blockSize <= 0 {
		return nil, errors.New("block size must be greater than zero")
	}
	switch normalizeMode(mode) {
	case "", PKCS7:
		return PKCS7Unpadding(data, blockSize)
	case PKCS5:
		// Compatibility mode: historically PKCS5 was used generically in this project.
		return PKCS7Unpadding(data, blockSize)
	case ZERO:
		return zeroUnpadding(data)
	case ANSIX923:
		return ansiX923Unpadding(data, blockSize)
	default:
		return nil, errors.New("padding type not supported")
	}
}

func PKCS7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

func PKCS7Unpadding(data []byte, blockSize int) ([]byte, error) {
	if len(data) == 0 || len(data)%blockSize != 0 {
		return nil, errors.New("invalid PKCS7 padding")
	}
	padding := int(data[len(data)-1])
	if padding <= 0 || padding > blockSize || padding > len(data) {
		return nil, errors.New("invalid PKCS7 padding")
	}
	for i := len(data) - padding; i < len(data); i++ {
		if int(data[i]) != padding {
			return nil, errors.New("invalid PKCS7 padding")
		}
	}
	return data[:len(data)-padding], nil
}

func PKCS5Padding(data []byte) []byte {
	return PKCS7Padding(data, 8)
}

func PKCS5Unpadding(data []byte) ([]byte, error) {
	return PKCS7Unpadding(data, 8)
}

func zeroPadding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	if padding == blockSize {
		return data
	}
	return append(data, bytes.Repeat([]byte{0x00}, padding)...)
}

func zeroUnpadding(data []byte) ([]byte, error) {
	idx := len(data)
	for idx > 0 && data[idx-1] == 0x00 {
		idx--
	}
	return data[:idx], nil
}

func ansiX923Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	if padding == 0 {
		padding = blockSize
	}
	pad := make([]byte, padding)
	pad[len(pad)-1] = byte(padding)
	return append(data, pad...)
}

func ansiX923Unpadding(data []byte, blockSize int) ([]byte, error) {
	if len(data) == 0 || len(data)%blockSize != 0 {
		return nil, errors.New("invalid ANSI X9.23 padding")
	}
	padding := int(data[len(data)-1])
	if padding <= 0 || padding > blockSize || padding > len(data) {
		return nil, errors.New("invalid ANSI X9.23 padding")
	}
	for i := len(data) - padding; i < len(data)-1; i++ {
		if data[i] != 0x00 {
			return nil, errors.New("invalid ANSI X9.23 padding")
		}
	}
	return data[:len(data)-padding], nil
}

func normalizeMode(mode string) string {
	return strings.ToUpper(strings.TrimSpace(mode))
}