-- SPDX-FileCopyrightText: 2023 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | BLS-MinPk cryptographic primitives.

module Morley.Tezos.Crypto.BLS
  ( -- * Cryptographic primitive types
    PublicKey (..)
  , SecretKey
  , Signature (..)
  , detSecretKey
  , randomSecretKey
  , toPublic

  -- * Raw bytes (no checksums, tags or anything)
  , publicKeyToBytes
  , mkPublicKey
  , publicKeyLengthBytes
  , signatureToBytes
  , mkSignature
  , signatureLengthBytes

  -- * Formatting and parsing
  , formatPublicKey
  , mformatPublicKey
  , parsePublicKey
  , formatSecretKey
  , parseSecretKey
  , formatSignature
  , mformatSignature
  , parseSignature

  -- * Signing
  , sign
  , checkSignature
  ) where

import Crypto.BLST qualified as BLST
import Crypto.Random (MonadRandom(getRandomBytes))
import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, convert)
import Data.ByteArray.Sized (SizedByteArray(..), sizedByteArray)
import Fmt (Buildable, build)

import Morley.Michelson.Text
import Morley.Tezos.Crypto.Util

----------------------------------------------------------------------------
-- Types, instances, conversions
----------------------------------------------------------------------------

type MinPk = 'BLST.G1

-- | Domain separation tag used by tezos
dst :: Maybe ByteString
-- see https://gitlab.com/nomadic-labs/cryptography/ocaml-bls12-381-signature/-/blob/1.0.0/src/bls12_381_signature.ml#L350
-- https://gitlab.com/tezos/tezos/-/blob/e9a9b61969b7a44749bed9bd9cdbcb4f2283220f/src/lib_crypto/bls.ml#L333
dst = Just "BLS_SIG_BLS12381G2_XMD:SHA-256_SSWU_RO_AUG_"

-- | BLS-MinPk public cryptographic key.
newtype PublicKey = PublicKey
  { unPublicKey :: BLST.PublicKey MinPk
  } deriving stock (Show, Eq, Generic)
    deriving anyclass NFData

instance Ord PublicKey where
  compare = compare `on` (publicKeyToBytes :: PublicKey -> ByteString)

-- | BLS-MinPk secret cryptographic key.
newtype SecretKey = SecretKey
  { unSecretKey :: BLST.SecretKey
  } deriving stock (Show, Eq, Generic)
    deriving anyclass NFData

-- | Deterministicaly generate a secret key from seed.
detSecretKey :: ByteString -> SecretKey
detSecretKey seed = SecretKey $ deterministic seed do
  bs <- fromMaybe (error "impossible") . sizedByteArray @32 @Bytes <$> getRandomBytes 32
  pure $ BLST.keygen bs

-- | Generate a random secret key.
randomSecretKey :: MonadRandom m => m SecretKey
randomSecretKey = detSecretKey <$> getRandomBytes 32

-- | Create a public key from a secret key.
toPublic :: SecretKey -> PublicKey
toPublic = PublicKey . BLST.skToPk . unSecretKey

-- | BLS-MinPk cryptographic signature.
newtype Signature = Signature
  { unSignature :: BLST.Signature MinPk 'BLST.Hash
  } deriving stock (Show, Eq, Generic)
    deriving anyclass NFData

----------------------------------------------------------------------------
-- Conversion to/from raw bytes (no checksums, tags or anything)
----------------------------------------------------------------------------

-- | Convert a 'PublicKey' to raw bytes.
publicKeyToBytes :: ByteArray ba => PublicKey -> ba
publicKeyToBytes = convert . unSizedByteArray . BLST.compressPk . unPublicKey

-- | Convert a 'PublicKey' to raw bytes.
secretKeyToBytes :: ByteArray ba => SecretKey -> ba
secretKeyToBytes = convert . unSizedByteArray . BLST.serializeSk . unSecretKey

toSized
  :: forall n a. (KnownNat n, ByteArrayAccess a)
  => LText -> a -> Either CryptoParseError (SizedByteArray n a)
toSized what =
  maybeToRight (CryptoParseUnexpectedLength what $ fromIntegralOverflowing $ natVal @n Proxy)
  . sizedByteArray

-- | Make a 'PublicKey' from raw bytes.
mkPublicKey :: ByteArrayAccess ba => ba -> Either CryptoParseError PublicKey
mkPublicKey = toSized "public key" >=> bimap CryptoParseBLSTError PublicKey . BLST.decompressPk

publicKeyLengthBytes :: (Integral n, CheckIntSubType Int n) => n
publicKeyLengthBytes = fromIntegral $ BLST.byteSize @'BLST.Compress @(BLST.PublicKey MinPk)

-- | Convert a 'Signature' to raw bytes.
signatureToBytes :: ByteArray ba => Signature -> ba
signatureToBytes = convert . unSizedByteArray . BLST.compressSignature . unSignature

-- | Make a 'Signature' from raw bytes.
mkSignature :: ByteArrayAccess ba => ba -> Either CryptoParseError Signature
mkSignature = toSized "signature" >=>
  bimap CryptoParseBLSTError Signature . BLST.decompressSignature

signatureLengthBytes :: (Integral n, CheckIntSubType Int n) => n
signatureLengthBytes = fromIntegral $ BLST.byteSize @'BLST.Compress @(BLST.Signature MinPk 'BLST.Hash)

mkSecretKey :: ByteArrayAccess ba => ba -> Either CryptoParseError SecretKey
mkSecretKey = toSized "secret key" >=>
  bimap CryptoParseBLSTError SecretKey . pure . BLST.deserializeSk

----------------------------------------------------------------------------
-- Magic bytes
----------------------------------------------------------------------------

-- https://gitlab.com/tezos/tezos/-/blob/4b0dd9e9715ce82ac6429571d8843ab681522daf/src/lib_crypto/base58.ml#L428

publicKeyTag :: ByteString
publicKeyTag = "\006\149\135\204" -- BLpk(76)

secretKeyTag :: ByteString
secretKeyTag = "\003\150\192\040" -- BLsk(54)

signatureTag :: ByteString
signatureTag = "\040\171\064\207" -- BLsig(142)

----------------------------------------------------------------------------
-- Formatting
----------------------------------------------------------------------------

formatPublicKey :: PublicKey -> Text
formatPublicKey = formatImpl publicKeyTag . publicKeyToBytes @Bytes

mformatPublicKey :: PublicKey -> MText
mformatPublicKey = unsafe . mkMText . formatPublicKey

instance Buildable PublicKey where
  build = build . formatPublicKey

parsePublicKey :: Text -> Either CryptoParseError PublicKey
parsePublicKey = parseImpl publicKeyTag mkPublicKey

formatSecretKey :: SecretKey -> Text
formatSecretKey = formatImpl secretKeyTag . secretKeyToBytes @Bytes

instance Buildable SecretKey where
  build = build . formatSecretKey

parseSecretKey :: Text -> Either CryptoParseError SecretKey
parseSecretKey = parseImpl secretKeyTag mkSecretKey

formatSignature :: Signature -> Text
formatSignature = formatImpl signatureTag . signatureToBytes @Bytes

mformatSignature :: Signature -> MText
mformatSignature = unsafe . mkMText . formatSignature

instance Buildable Signature where
  build = build . formatSignature

parseSignature :: Text -> Either CryptoParseError Signature
parseSignature = parseImpl signatureTag mkSignature

----------------------------------------------------------------------------
-- Signing
----------------------------------------------------------------------------

-- | Sign a message using the secret key.
sign :: SecretKey -> ByteString -> Signature
sign sk msg = Signature $ BLST.sign (unSecretKey sk)
  (publicKeyToBytes (toPublic sk) <> msg) dst
-- upstream does this concatenation with compressed public key
-- https://gitlab.com/nomadic-labs/cryptography/ocaml-bls12-381-signature/-/blob/531ed1f509a974f5067f431b6797b9246518520c/src/bls12_381_signature.ml#L662

-- | Check that a sequence of bytes has been signed with a given key.
checkSignature :: PublicKey -> Signature -> ByteString -> Bool
checkSignature pkg@(PublicKey pk) (Signature sig) bytes =
  BLST.BlstSuccess == BLST.verify sig pk (publicKeyToBytes pkg <> bytes) dst