-- | Utilities shared by multiple cryptographic primitives.

module Tezos.Crypto.Util
  ( CryptoParseError (..)
  , encodeBase58Check
  , decodeBase58Check
  , B58CheckWithPrefixError (..)
  , decodeBase58CheckWithPrefix
  , formatImpl
  , parseImpl
  , firstRight
  , deterministic
  ) where

import Crypto.Error (CryptoError)
import Crypto.Number.Serialize (os2ip)
import Crypto.Random (ChaChaDRG, MonadPseudoRandom, drgNewSeed, seedFromInteger, withDRG)
import qualified Data.ByteArray as ByteArray
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base58 as Base58
import Fmt (Buildable, Builder, build, fmt, hexF)

import Tezos.Crypto.Hash

-- | Error that can happen during parsing of cryptographic primitive types.
data CryptoParseError
  = CryptoParseWrongBase58Check
  | CryptoParseWrongTag ByteString
  | CryptoParseCryptoError CryptoError
  | CryptoParseUnexpectedLength Builder Int
  deriving stock (Show, Eq)

instance Buildable CryptoParseError where
  build =
    \case
      CryptoParseWrongBase58Check -> "Wrong base58check encoding of bytes"
      CryptoParseWrongTag tag -> "Prefix is wrong tag: " <> fmt (hexF tag)
      CryptoParseCryptoError err ->
        "Cryptographic library reported an error: " <>
        build (displayException err)
      CryptoParseUnexpectedLength what l ->
        "Unexpected length of " <> what <> ": " <> build l

-- | Encode a bytestring in Base58Check format.
encodeBase58Check :: ByteString -> Text
encodeBase58Check =
  decodeUtf8 . Base58.encodeBase58 Base58.bitcoinAlphabet . withCheckSum
  where
    withCheckSum :: ByteString -> ByteString
    withCheckSum bs = bs <> checkSum bs

-- | Decode a bytestring from Base58Check format.
decodeBase58Check :: Text -> Maybe ByteString
decodeBase58Check text = do
  bytes <- Base58.decodeBase58 Base58.bitcoinAlphabet (encodeUtf8 text)
  let (payload, chk) = BS.splitAt (length bytes - 4) bytes
  guard $ chk == checkSum payload
  return payload

checkSum :: ByteString -> ByteString
checkSum = BS.take 4 . (sha256 . sha256)

data B58CheckWithPrefixError
  = B58CheckWithPrefixWrongPrefix ByteString
  | B58CheckWithPrefixWrongEncoding
  deriving stock (Show)

-- | Parse a base58check encoded value expecting some prefix. If the
-- actual prefix matches the expected one, it's stripped of and the
-- resulting payload is returned.
decodeBase58CheckWithPrefix ::
  ByteString -> Text -> Either B58CheckWithPrefixError ByteString
decodeBase58CheckWithPrefix prefix text =
  case decodeBase58Check text of
    Nothing -> Left B58CheckWithPrefixWrongEncoding
    Just bs ->
      let (actualPrefix, payload) = BS.splitAt (length prefix) bs
       in if actualPrefix == prefix
          then Right payload
          else Left (B58CheckWithPrefixWrongPrefix actualPrefix)

-- | Template for 'format*' functions.
formatImpl :: ByteArray.ByteArrayAccess x => ByteString -> x -> Text
formatImpl tag = encodeBase58Check . mappend tag . ByteArray.convert

-- | Template for 'parse*' functions.
parseImpl
  :: ByteString
  -> (ByteString -> Either CryptoParseError res)
  -> Text
  -> Either CryptoParseError res
parseImpl expectedTag constructor text = do
  let convertErr :: B58CheckWithPrefixError -> CryptoParseError
      convertErr =
        \case B58CheckWithPrefixWrongPrefix prefix -> CryptoParseWrongTag prefix
              B58CheckWithPrefixWrongEncoding -> CryptoParseWrongBase58Check
  payload <- first convertErr $ decodeBase58CheckWithPrefix expectedTag text
  constructor payload

-- | Returns first encountered 'Right' in a list. If there are none,
-- returns arbitrary 'Left'.
-- It is useful to implement parsing.
firstRight :: NonEmpty (Either e a) -> Either e a
firstRight (h :| rest) =
  case h of
    Left e -> maybe (Left e) firstRight $ nonEmpty rest
    Right a -> Right a

-- | Do randomized action using specified seed.
deterministic :: ByteString -> MonadPseudoRandom ChaChaDRG a -> a
deterministic seed = fst . withDRG chachaSeed
  where
    chachaSeed = drgNewSeed . seedFromInteger . os2ip $ seed