package fpe

import (
	"crypto/cipher"
	"errors"
	"math"
	"math/big"

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

func rev[T any](X []T) {
	l := len(X) / 2
	for i := 0; i < l; i++ {
		t := X[i]
		X[i] = X[l-i-1]
		X[l-i-1] = t
	}
}

// in = in0||in1||in2...
// out0 = Enc(iv ^ in0)
// out1 = Enc(out0 ^ in1)
// ...
// if iv == nil <==> iv = 0^16.
func prf(res []byte, b cipher.Block, iv []byte, in []byte) []byte {
	xor.XorBytes(res, iv, in[:16])
	b.Encrypt(res, res)
	in = in[16:]

	for len(in) > 0 {
		xor.XorBytes(res, res, in[:16])
		b.Encrypt(res, res)
		in = in[16:]
	}
	return res
}

type FF1 struct {
	cipher.Block
	Alphabet
}

func NewFF1(b cipher.Block, a Alphabet) FPE {
	return &FF1{
		Block:    b,
		Alphabet: a,
	}
}

func (f *FF1) round(P []byte, Q []byte, S []byte, A []Numeral, B []Numeral, i int, b int, d int, module *big.Int, isenc bool) {
	Q[len(Q)-b-1] = byte(i)
	f.Num(B).FillBytes(Q[len(Q)-b:])
	prf(S, f.Block, P, Q)
	for j := 1; j < (d+15)/16; j++ {
		Sj := S[16*j : 16*j+16]
		S[16*j+12] = byte(j >> 24)
		S[16*j+13] = byte(j >> 16)
		S[16*j+14] = byte(j >> 8)
		S[16*j+15] = byte(j)
		xor.XorBytes(Sj, S[:16], Sj)
		f.Block.Encrypt(Sj, Sj)
	}
	y := new(big.Int).SetBytes(S[:d])
	c := f.Num(A)
	if isenc {
		c.Add(c, y).Mod(c, module)
	} else {
		c.Sub(c, y).Mod(c, module)
	}
	f.Str(A, c)
}

func (f *FF1) endecrypt(T []byte, X []Numeral, isEnc bool) ([]Numeral, error) {
	n := len(X)
	if n <= 2 {
		return nil, errors.New("input numeral string X must at least 2 characters")
	}

	radix := f.Alphabet.Radix()
	t := len(T)
	u := n / 2
	v := n - u
	b := (int(math.Ceil(float64(v)*math.Log2(float64(radix)))) + 7) >> 3
	d := (b + 7) & (^3)
	P := []byte{
		1, 2, 1,
		0, byte(radix >> 8), byte(radix),
		10, byte(u),
		byte(n >> 24), byte(n >> 16), byte(n >> 8), byte(n),
		byte(t >> 24), byte(t >> 16), byte(t >> 8), byte(t),
	}
	f.Block.Encrypt(P, P)
	// make a copy of X
	res := make([]Numeral, len(X))
	copy(res, X)
	A := res[:u]
	B := res[u:]

	// Q := append(make([]byte, 0, (b+t+1+15)&(^15)), T...)
	Q := make([]byte, (b+t+1+15)&(^15))
	copy(Q, T)
	S := make([]byte, (d+15)&(^15))

	var moduleA, moduleB *big.Int
	radixBig := big.NewInt(int64(radix))
	// 如果输入X很大，这里的module会很大。
	moduleA = new(big.Int).Exp(radixBig, big.NewInt(int64(u)), nil)
	moduleB = moduleA
	if v > u {
		moduleB = radixBig.Mul(moduleB, radixBig)
	}

	if isEnc {
		for i := 0; i < 10; i += 2 {
			f.round(P, Q, S, A, B, i, b, d, moduleA, true)
			f.round(P, Q, S, B, A, i+1, b, d, moduleB, true)
		}
	} else {
		for i := 9; i >= 0; i -= 2 {
			f.round(P, Q, S, B, A, i, b, d, moduleB, false)
			f.round(P, Q, S, A, B, i-1, b, d, moduleA, false)
		}
	}
	return res, nil
}

func (f *FF1) Encrypt(T []byte, X []Numeral) ([]Numeral, error) {
	return f.endecrypt(T, X, true)
}

func (f FF1) Decrypt(T []byte, X []Numeral) ([]Numeral, error) {
	return f.endecrypt(T, X, false)
}