package math

import (
	"crypto/md5"

	"github.com/capitalone/fpe/ff3"
)

type Crypt struct {
	// Key and tweak should be byte arrays. Put your key and tweak here.
	// To make it easier for demo purposes, decode from a hex string here.
	key   []byte
	tweak []byte

	cipher ff3.Cipher
}

func (c *Crypt) Init(skey, stweak string) error {
	c.key = stringToMd5Bytes(skey, 32)
	c.tweak = stringToMd5Bytes(stweak, 8)

	// Create a new FF3 cipher "object"
	// 10 is the radix/base, and 8 is the tweak length.
	FF3, err := ff3.NewCipher(10, c.key, c.tweak)
	if err != nil {
		return err
	}
	c.cipher = FF3
	return nil
}

func stringToMd5Bytes(s string, size int) []byte {
	h := md5.New()
	h.Write([]byte(s))
	bts := h.Sum(nil)
	if size < len(bts) {
		return bts[:size]
	}
	return bts
}

func (c *Crypt) Encrypt(s string) (string, error) {
	// Call the encryption function on an example SSN
	ciphertext, err := c.cipher.Encrypt(s)
	if err != nil {
		return "", err
	}
	return ciphertext, nil
}

func (c *Crypt) Decrypt(s string) (string, error) {
	plaintext, err := c.cipher.Decrypt(s)
	if err != nil {
		return "", err
	}
	return plaintext, nil
}