-- | Secp256k1 cryptographic primitives.

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

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

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

  -- * Signing
  , sign
  , checkSignature
  ) where

import Crypto.Hash (Blake2b_256(..))
import Crypto.Number.Serialize (i2ospOf_, os2ip)
import qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import qualified Crypto.PubKey.ECC.Generate as ECC.Generate
import Crypto.PubKey.ECC.Types (Curve, CurveName(..), Point(..), curveSizeBits, getCurveByName)
import Crypto.Random (MonadRandom, drgNewSeed, seedFromInteger, withDRG)
import Data.ByteArray (ByteArray, ByteArrayAccess)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import Fmt (Buildable, build)
import Test.QuickCheck (Arbitrary(..), vector)

import Michelson.Text
import Tezos.Crypto.Util

curve :: Curve
curve = getCurveByName SEC_p256k1

curveSizeBytes :: Int
curveSizeBytes = curveSizeBits curve `div` 8

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

-- | Secp256k1 public cryptographic key.
data PublicKey = PublicKey
  { unPublicKey :: ECDSA.PublicKey
  , pkBytes :: Maybe ByteString
  -- ^ This is the hack we use to make serialization correct.
  -- Decoding is currently not implemented, so when we have to
  -- decode bytes we remember these bytes and produce some random
  -- public key.
  --
  -- TODO (#18) remove it.
  } deriving stock (Show)

-- TODO (#18): derive it instead once the above hack is removed.
instance Eq PublicKey where
  pk1 == pk2 = publicKeyToBytes @ByteString pk1 == publicKeyToBytes pk2

instance Arbitrary PublicKey where
  arbitrary = toPublic <$> arbitrary

-- | Secp256k1 secret cryptographic key.
newtype SecretKey = SecretKey
  { unSecretKey :: ECDSA.KeyPair
  } deriving stock (Show, Eq)

-- | Deterministicaly generate a secret key from seed.
detSecretKey :: ByteString -> SecretKey
detSecretKey seed = deterministic seed $ detSecretKeyDo

detSecretKeyDo :: MonadRandom m => m SecretKey
detSecretKeyDo = SecretKey <$> do
  (publicKey, privateKey) <- ECC.Generate.generate curve
  return $
    ECDSA.KeyPair curve (ECDSA.public_q publicKey) (ECDSA.private_d privateKey)

instance Arbitrary SecretKey where
  arbitrary = detSecretKey . BS.pack <$> vector 32

-- | Create a public key from a secret key.
toPublic :: SecretKey -> PublicKey
toPublic =
  flip PublicKey Nothing .
  ECDSA.PublicKey curve . (\(ECDSA.KeyPair _ pp _) -> pp) . unSecretKey

-- | Secp256k1 cryptographic signature.
newtype Signature = Signature
  { unSignature :: ECDSA.Signature
  } deriving stock (Show, Eq)

instance Arbitrary Signature where
  arbitrary = do
    seed <- drgNewSeed . seedFromInteger <$> arbitrary
    byteToSign <- arbitrary
    return $ fst $ withDRG seed $ do
      sk <- detSecretKeyDo
      sign sk (one byteToSign)

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

-- | Convert a 'PublicKey' to raw bytes.
--
-- TODO (#18): apparently it uses compressed SEC format as described in
-- <https://www.oreilly.com/library/view/programming-bitcoin/9781492031482/ch04.html>
-- However, it is not tested yet.
publicKeyToBytes :: forall ba. ByteArray ba => PublicKey -> ba
publicKeyToBytes (PublicKey _ (Just bytes)) = BA.convert bytes
publicKeyToBytes (PublicKey (ECDSA.PublicKey _ publicPoint) Nothing) =
  case publicPoint of
    Point x y -> prefix y `BA.append` coordToBytes x
    PointO -> error "PublicKey somehow contains infinity point"
  where
    prefix :: Integer -> ba
    prefix y
      | odd y = BA.singleton 0x03
      | otherwise = BA.singleton 0x02

-- | Make a 'PublicKey' from raw bytes.
--
-- TODO (#18): it should decode from compressed SEC format, but it's left
-- for a future task, so for now we return a constant.
mkPublicKey :: ByteArrayAccess ba => ba -> Either CryptoParseError PublicKey
mkPublicKey ba
  | l == publicKeyLengthBytes =
    Right $ PublicKey (ECDSA.PublicKey curve $ Point 11 12) (Just $ BA.convert ba)
  | otherwise =
    Left $ CryptoParseUnexpectedLength "public key" l
  where
    l = BA.length ba

publicKeyLengthBytes :: Integral n => n
publicKeyLengthBytes = fromIntegral $ 1 + curveSizeBytes

-- | Convert a 'PublicKey' to raw bytes.
--
-- TODO (#18): apparently a signature always has 64 bytes, so this
-- format might be correct, but it is not tested.
signatureToBytes :: ByteArray ba => Signature -> ba
signatureToBytes (Signature (ECDSA.Signature r s)) =
  coordToBytes r <> coordToBytes s

-- | Make a 'Signature' from raw bytes.
--
-- TODO (#18): apparently a signature always has 64 bytes, so this
-- format might be correct, but it is not tested.
mkSignature :: ByteArray ba => ba -> Either CryptoParseError Signature
mkSignature ba
  | l == signatureLengthBytes
  , (rBytes, sBytes) <- BA.splitAt curveSizeBytes ba =
    Right $ Signature (ECDSA.Signature (os2ip rBytes) (os2ip sBytes))
  | otherwise =
    Left $ CryptoParseUnexpectedLength "signature" l
  where
    l = BA.length ba

signatureLengthBytes :: Integral n => n
signatureLengthBytes = fromIntegral $ curveSizeBytes + curveSizeBytes

-- TODO (#18): maybe this function doesn't make sense.
-- We are using `i2ospOf_` because `curveSizeBits` ensures that
-- the number won't have more than that many bytes.
coordToBytes :: ByteArray ba => Integer -> ba
coordToBytes = i2ospOf_ curveSizeBytes

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

publicKeyTag :: ByteString
publicKeyTag = "\003\254\226\086"

signatureTag :: ByteString
signatureTag = "\013\115\101\019\063"

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

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

mformatPublicKey :: PublicKey -> MText
mformatPublicKey = mkMTextUnsafe . formatPublicKey

instance Buildable PublicKey where
  build = build . formatPublicKey

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

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

mformatSignature :: Signature -> MText
mformatSignature = mkMTextUnsafe . 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 :: MonadRandom m => SecretKey -> ByteString -> m Signature
sign (SecretKey keyPair) =
  fmap Signature . ECDSA.sign (ECDSA.toPrivateKey keyPair) Blake2b_256

-- | Check that a sequence of bytes has been signed with a given key.
checkSignature :: PublicKey -> Signature -> ByteString -> Bool
checkSignature (PublicKey pk _) (Signature sig) =
  ECDSA.verify Blake2b_256 pk sig