package bls

import (
	"encoding/hex"
	bls "github.com/consensys/gnark-crypto/ecc/bls12-381"
	"github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
	"math/big"
)

const dstG1 = "BLS_SIG_BLS12381G1_XMD:SHA-256_SSWU_RO_NUL_"

var g2, g2Gen bls.G2Affine

func init() {
	_, _, _, g2Gen = bls.Generators()
	g2.Neg(&g2Gen)
}

type PublicKey bls.G2Affine

// PublicKeyFromBytes returns a PublicKey from a byte slice.
func PublicKeyFromBytes(b []byte) (*PublicKey, error) {
	var publicKey bls.G2Affine
	if err := publicKey.Unmarshal(b); err != nil {
		return nil, err
	}
	return (*PublicKey)(&publicKey), nil
}

// PublicKeyFromHexString returns a PublicKey from a hex string.
func PublicKeyFromHexString(s string) (*PublicKey, error) {
	b, err := hex.DecodeString(s)
	if err != nil {
		return nil, err
	}
	return PublicKeyFromBytes(b)
}

type SecretKey fr.Element

// NewSecretKeyByCSPRNG returns a new SecretKey generated by CSPRNG.
func NewSecretKeyByCSPRNG() *SecretKey {
	var secretKey fr.Element
	if _, err := secretKey.SetRandom(); err != nil {
		return nil
	}
	return (*SecretKey)(&secretKey)
}

func (sk *SecretKey) PublicKey() *PublicKey {
	element := fr.Element(*sk)
	return (*PublicKey)(g2Gen.ScalarMultiplication(
		&g2Gen,
		element.BigInt(big.NewInt(0)),
	))
}

func (sk *SecretKey) Sign(msg []byte) (*Signature, error) {
	element := fr.Element(*sk)
	g1, err := bls.HashToG1(msg, []byte(dstG1))
	if err != nil {
		return nil, err
	}
	return (*Signature)(g1.ScalarMultiplication(
		&g1,
		element.BigInt(big.NewInt(0)),
	)), nil
}

type Signature bls.G1Affine

// SignatureFromBytes returns a Signature from a byte slice.
func SignatureFromBytes(b []byte) (*Signature, error) {
	var signature bls.G1Affine
	if _, err := signature.SetBytes(b); err != nil {
		return nil, err
	}
	return (*Signature)(&signature), nil
}

// SignatureFromHexString returns a Signature from a hex string.
func SignatureFromHexString(s string) (*Signature, error) {
	b, err := hex.DecodeString(s)
	if err != nil {
		return nil, err
	}
	return SignatureFromBytes(b)
}

func (sig *Signature) Verify(pk *PublicKey, msg []byte) bool {
	g1, err := bls.HashToG1(msg, []byte(dstG1))
	if err != nil {
		return false
	}
	valid, err := bls.PairingCheck(
		[]bls.G1Affine{bls.G1Affine(*sig), g1},
		[]bls.G2Affine{g2, bls.G2Affine(*pk)},
	)
	if err != nil {
		return false
	}
	return valid
}
