package crypto

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"io"
	"math/big"

	"github.com/samber/oops"
	"github.com/sirupsen/logrus"

	"golang.org/x/crypto/openpgp/elgamal"
)

var elgp = new(big.Int).SetBytes([]byte{
	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xC9, 0x0F, 0xDA, 0xA2, 0x21, 0x68, 0xC2, 0x34,
	0xC4, 0xC6, 0x62, 0x8B, 0x80, 0xDC, 0x1C, 0xD1, 0x29, 0x02, 0x4E, 0x08, 0x8A, 0x67, 0xCC, 0x74,
	0x02, 0x0B, 0xBE, 0xA6, 0x3B, 0x13, 0x9B, 0x22, 0x51, 0x4A, 0x08, 0x79, 0x8E, 0x34, 0x04, 0xDD,
	0xEF, 0x95, 0x19, 0xB3, 0xCD, 0x3A, 0x43, 0x1B, 0x30, 0x2B, 0x0A, 0x6D, 0xF2, 0x5F, 0x14, 0x37,
	0x4F, 0xE1, 0x35, 0x6D, 0x6D, 0x51, 0xC2, 0x45, 0xE4, 0x85, 0xB5, 0x76, 0x62, 0x5E, 0x7E, 0xC6,
	0xF4, 0x4C, 0x42, 0xE9, 0xA6, 0x37, 0xED, 0x6B, 0x0B, 0xFF, 0x5C, 0xB6, 0xF4, 0x06, 0xB7, 0xED,
	0xEE, 0x38, 0x6B, 0xFB, 0x5A, 0x89, 0x9F, 0xA5, 0xAE, 0x9F, 0x24, 0x11, 0x7C, 0x4B, 0x1F, 0xE6,
	0x49, 0x28, 0x66, 0x51, 0xEC, 0xE4, 0x5B, 0x3D, 0xC2, 0x00, 0x7C, 0xB8, 0xA1, 0x63, 0xBF, 0x05,
	0x98, 0xDA, 0x48, 0x36, 0x1C, 0x55, 0xD3, 0x9A, 0x69, 0x16, 0x3F, 0xA8, 0xFD, 0x24, 0xCF, 0x5F,
	0x83, 0x65, 0x5D, 0x23, 0xDC, 0xA3, 0xAD, 0x96, 0x1C, 0x62, 0xF3, 0x56, 0x20, 0x85, 0x52, 0xBB,
	0x9E, 0xD5, 0x29, 0x07, 0x70, 0x96, 0x96, 0x6D, 0x67, 0x0C, 0x35, 0x4E, 0x4A, 0xBC, 0x98, 0x04,
	0xF1, 0x74, 0x6C, 0x08, 0xCA, 0x18, 0x21, 0x7C, 0x32, 0x90, 0x5E, 0x46, 0x2E, 0x36, 0xCE, 0x3B,
	0xE3, 0x9E, 0x77, 0x2C, 0x18, 0x0E, 0x86, 0x03, 0x9B, 0x27, 0x83, 0xA2, 0xEC, 0x07, 0xA2, 0x8F,
	0xB5, 0xC5, 0x5D, 0xF0, 0x6F, 0x4C, 0x52, 0xC9, 0xDE, 0x2B, 0xCB, 0xF6, 0x95, 0x58, 0x17, 0x18,
	0x39, 0x95, 0x49, 0x7C, 0xEA, 0x95, 0x6A, 0xE5, 0x15, 0xD2, 0x26, 0x18, 0x98, 0xFA, 0x05, 0x10,
	0x15, 0x72, 0x8E, 0x5A, 0x8A, 0xAC, 0xAA, 0x68, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
})

var (
	one  = big.NewInt(1)
	elgg = big.NewInt(2)
)

var (
	ElgDecryptFail   = oops.Errorf("failed to decrypt elgamal encrypted data")
	ElgEncryptTooBig = oops.Errorf("failed to encrypt data, too big for elgamal")
)

// generate an elgamal key pair
func ElgamalGenerate(priv *elgamal.PrivateKey, rand io.Reader) (err error) {
	log.Debug("Generating ElGamal key pair")
	priv.P = elgp
	priv.G = elgg
	xBytes := make([]byte, priv.P.BitLen()/8)
	_, err = io.ReadFull(rand, xBytes)
	if err == nil {
		// set private key
		priv.X = new(big.Int).SetBytes(xBytes)
		// compute public key
		priv.Y = new(big.Int).Exp(priv.G, priv.X, priv.P)
		log.Debug("ElGamal key pair generated successfully")
	} else {
		log.WithError(err).Error("Failed to generate ElGamal key pair")
	}
	return
}

type elgDecrypter struct {
	k *elgamal.PrivateKey
}

func (elg *elgDecrypter) Decrypt(data []byte) (dec []byte, err error) {
	log.WithField("data_length", len(data)).Debug("Decrypting ElGamal data")
	dec, err = elgamalDecrypt(elg.k, data, true) // TODO(psi): should this be true or false?
	if err != nil {
		log.WithError(err).Error("Failed to decrypt ElGamal data")
	} else {
		log.WithField("decrypted_length", len(dec)).Debug("ElGamal data decrypted successfully")
	}
	return
}

// decrypt an elgamal encrypted message, i2p style
func elgamalDecrypt(priv *elgamal.PrivateKey, data []byte, zeroPadding bool) (decrypted []byte, err error) {
	log.WithFields(logrus.Fields{
		"data_length":  len(data),
		"zero_padding": zeroPadding,
	}).Debug("Decrypting ElGamal data")

	a := new(big.Int)
	b := new(big.Int)
	idx := 0
	if zeroPadding {
		idx++
	}
	a.SetBytes(data[idx : idx+256])
	if zeroPadding {
		idx++
	}
	b.SetBytes(data[idx+256:])

	// decrypt
	m := new(big.Int).Mod(new(big.Int).Mul(b, new(big.Int).Exp(a, new(big.Int).Sub(new(big.Int).Sub(priv.P, priv.X), one), priv.P)), priv.P).Bytes()

	// check digest
	d := sha256.Sum256(m[33:255])
	good := 0
	if subtle.ConstantTimeCompare(d[:], m[1:33]) == 1 {
		// decryption successful
		good = 1
		log.Debug("ElGamal decryption successful")
	} else {
		// decrypt failed
		err = ElgDecryptFail
		log.WithError(err).Error("ElGamal decryption failed")
	}
	// copy result
	decrypted = make([]byte, 222)
	subtle.ConstantTimeCopy(good, decrypted, m[33:255])

	if good == 0 {
		// if decrypt failed nil out decrypted slice
		decrypted = nil
	}
	return
}

type ElgamalEncryption struct {
	p, a, b1 *big.Int
}

func (elg *ElgamalEncryption) Encrypt(data []byte) (enc []byte, err error) {
	log.WithField("data_length", len(data)).Debug("Encrypting data with ElGamal")
	return elg.EncryptPadding(data, true)
}

func (elg *ElgamalEncryption) EncryptPadding(data []byte, zeroPadding bool) (encrypted []byte, err error) {
	log.WithFields(logrus.Fields{
		"data_length":  len(data),
		"zero_padding": zeroPadding,
	}).Debug("Encrypting data with ElGamal padding")

	if len(data) > 222 {
		err = ElgEncryptTooBig
		return
	}
	mbytes := make([]byte, 255)
	mbytes[0] = 0xFF
	copy(mbytes[33:], data)
	// do sha256 of payload
	d := sha256.Sum256(mbytes[33 : len(data)+33])
	copy(mbytes[1:], d[:])
	m := new(big.Int).SetBytes(mbytes)
	// do encryption
	b := new(big.Int).Mod(new(big.Int).Mul(elg.b1, m), elg.p).Bytes()

	if zeroPadding {
		encrypted = make([]byte, 514)
		copy(encrypted[1:], elg.a.Bytes())
		copy(encrypted[258:], b)
	} else {
		encrypted = make([]byte, 512)
		copy(encrypted, elg.a.Bytes())
		copy(encrypted[256:], b)
	}

	log.WithField("encrypted_length", len(encrypted)).Debug("Data encrypted successfully with ElGamal")
	return
}

// create an elgamal public key from byte slice
func createElgamalPublicKey(data []byte) (k *elgamal.PublicKey) {
	log.WithField("data_length", len(data)).Debug("Creating ElGamal public key")
	if len(data) == 256 {
		k = &elgamal.PublicKey{
			G: elgg,
			P: elgp,
			Y: new(big.Int).SetBytes(data),
		}
		log.Debug("ElGamal public key created successfully")
	} else {
		log.Warn("Invalid data length for ElGamal public key")
	}

	return
}

// create an elgamal private key from byte slice
func createElgamalPrivateKey(data []byte) (k *elgamal.PrivateKey) {
	log.WithField("data_length", len(data)).Debug("Creating ElGamal private key")
	if len(data) == 256 {
		x := new(big.Int).SetBytes(data)
		y := new(big.Int).Exp(elgg, x, elgp)
		k = &elgamal.PrivateKey{
			PublicKey: elgamal.PublicKey{
				Y: y,
				G: elgg,
				P: elgp,
			},
			X: x,
		}
		log.Debug("ElGamal private key created successfully")
	} else {
		log.Warn("Invalid data length for ElGamal private key")
	}
	return
}

// create a new elgamal encryption session
func createElgamalEncryption(pub *elgamal.PublicKey, rand io.Reader) (enc *ElgamalEncryption, err error) {
	log.Debug("Creating ElGamal encryption session")
	kbytes := make([]byte, 256)
	k := new(big.Int)
	for err == nil {
		_, err = io.ReadFull(rand, kbytes)
		k = new(big.Int).SetBytes(kbytes)
		k = k.Mod(k, pub.P)
		if k.Sign() != 0 {
			break
		}
	}
	if err == nil {
		enc = &ElgamalEncryption{
			p:  pub.P,
			a:  new(big.Int).Exp(pub.G, k, pub.P),
			b1: new(big.Int).Exp(pub.Y, k, pub.P),
		}
		log.Debug("ElGamal encryption session created successfully")
	} else {
		log.WithError(err).Error("Failed to create ElGamal encryption session")
	}
	return
}

type (
	ElgPublicKey  [256]byte
	ElgPrivateKey [256]byte
)

func (elg ElgPublicKey) Len() int {
	return len(elg)
}

func (elg ElgPublicKey) Bytes() []byte {
	return elg[:]
}

func (elg ElgPublicKey) NewEncrypter() (enc Encrypter, err error) {
	log.Debug("Creating new ElGamal encrypter")
	k := createElgamalPublicKey(elg[:])
	enc, err = createElgamalEncryption(k, rand.Reader)
	if err != nil {
		log.WithError(err).Error("Failed to create ElGamal encrypter")
	} else {
		log.Debug("ElGamal encrypter created successfully")
	}
	return
}

func (elg ElgPrivateKey) Len() int {
	return len(elg)
}

func (elg ElgPrivateKey) NewDecrypter() (dec Decrypter, err error) {
	log.Debug("Creating new ElGamal decrypter")
	dec = &elgDecrypter{
		k: createElgamalPrivateKey(elg[:]),
	}
	log.Debug("ElGamal decrypter created successfully")
	return
}