-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

{-# LANGUAGE NoPolyKinds #-}

-- | Type-safe operations with @bytes@-like data.
module Lorentz.Bytes
  ( BytesLike (..)

  -- * Packed
  , Packed (..)

  -- * Signatures
  , TSignature (..)
  , lSign

  -- * Hashes
  , Hash (..)
  , DHashAlgorithm
  , KnownHashAlgorithm (..)
  , toHashHs
  , Sha256
  , Sha512
  , Blake2b
  , Sha3
  , Keccak

  -- * Typed Chest
  , ChestT (..)
  , OpenChestT (..)
  , openChestT
  ) where

import Crypto.Random (MonadRandom)
import Fmt (Buildable(..))
import Morley.Util.Markdown
import qualified Type.Reflection as Refl

import Lorentz.Annotation
import Lorentz.Base
import Lorentz.Constraints.Scopes
import Lorentz.Doc
import Lorentz.Value
import qualified Morley.Michelson.Typed as T
import Morley.Tezos.Crypto
import qualified Morley.Tezos.Crypto.Hash as Crypto

-- | Everything which is represented as bytes inside.
class (KnownValue bs, ToT bs ~ ToT ByteString) => BytesLike bs where
  toBytes :: bs -> ByteString

instance BytesLike ByteString where
  toBytes = id

----------------------------------------------------------------------------
-- Packing
----------------------------------------------------------------------------

-- | Represents a 'ByteString' resulting from packing a value of type @a@.
--
-- This is /not/ guaranteed to keep some packed value, and @unpack@ can fail.
-- We do so because often we need to accept values of such type from user,
-- and also because there is no simple way to check validity of packed data
-- without performing full unpack.
-- So this wrapper is rather a hint for users.
newtype Packed a = Packed { unPacked :: ByteString }
  deriving stock (Show, Eq, Ord, Generic)
  deriving newtype (IsoValue, HasAnnotation, BytesLike)

instance Buildable (Packed a) where
  build = build . toVal

instance TypeHasDoc a => TypeHasDoc (Packed a) where
  typeDocMdDescription = [md|
    Packed value of the given type.
    This exactly matches the result of Michelson `PACK` instruction application
    to the given value.
    |]
  typeDocMdReference = poly1TypeDocMdReference
  typeDocDependencies p =
    genericTypeDocDependencies p <>
    [ dTypeDep @a
    , dTypeDep @MText, dTypeDep @Integer  -- for examples below
    ]
  typeDocHaskellRep = concreteTypeDocHaskellRep @(Packed (MText, Integer))
  typeDocMichelsonRep = concreteTypeDocMichelsonRep @(Packed (MText, Integer))

----------------------------------------------------------------------------
-- Signatures
----------------------------------------------------------------------------

-- | Represents a signature, where signed data has given type.
--
-- Since we usually sign a packed data, a common pattern for this type is
-- @TSignature ('Packed' signedData)@.
-- If you don't want to use 'Packed', use plain @TSignature ByteString@ instead.
newtype TSignature a = TSignature { unTSignature :: Signature }
  deriving stock (Show, Generic)
  deriving newtype (IsoValue, HasAnnotation)

instance Buildable (TSignature a) where
  build = build . toVal

instance TypeHasDoc a => TypeHasDoc (TSignature a) where
  typeDocMdDescription = "Signature for data of the given type."
  typeDocMdReference = poly1TypeDocMdReference
  typeDocDependencies p =
    genericTypeDocDependencies p <>
    [ dTypeDep @a
    , dTypeDep @MText, dTypeDep @Integer  -- for examples below
    ]
  typeDocHaskellRep = concreteTypeDocHaskellRep @(TSignature (MText, Integer))
  typeDocMichelsonRep = concreteTypeDocMichelsonRep @(TSignature (MText, Integer))

-- | Sign data using 'SecretKey'
lSign :: (MonadRandom m, BytesLike a) => SecretKey -> a -> m (TSignature a)
lSign sk (toBytes -> bs) = TSignature <$> sign sk bs

----------------------------------------------------------------------------
-- Hashes
----------------------------------------------------------------------------

-- | Open kind for hash algorithms, to make it more difficult to apply type
-- arguments incorrectly.
type HashAlgorithmKind = HashAlgoTag -> Type
data HashAlgoTag

-- | Hash of type @t@ evaluated from data of type @a@.
newtype Hash (alg :: HashAlgorithmKind) a = UnsafeHash { unHash :: ByteString }
  deriving stock (Show, Eq, Ord, Generic)
  deriving newtype (IsoValue, HasAnnotation, BytesLike)

instance Buildable (Hash alg a) where
  build = build . toVal

instance (KnownHashAlgorithm alg, TypeHasDoc a) => TypeHasDoc (Hash alg a) where
  typeDocMdDescription = [md|
    Hash of a value.

    First type argument denotes algorithm used to compute the hash, and the second
    argument describes the data being hashed.
    |]

  typeDocMdReference tp wp =
    T.applyWithinParens wp $ mconcat
      [ mdLocalRef (mdTicked "Hash") (docItemRef (DType tp))
      , " "
      , hashAlgorithmMdRef (Proxy @alg)
      , " "
      , typeDocMdReference (Proxy @a) (T.WithinParens True)
      ]

  typeDocDependencies p =
    genericTypeDocDependencies p <>
    [ SomeDocDefinitionItem (DHashAlgorithm (Proxy @alg)), dTypeDep @a
    , SomeDocDefinitionItem (DHashAlgorithm (Proxy @Blake2b)), dTypeDep @ByteString
        --- ^ for examples below
    ]
  typeDocHaskellRep = concreteTypeDocHaskellRep @(Hash Blake2b ByteString)
  typeDocMichelsonRep = concreteTypeDocMichelsonRep @(Hash Blake2b ByteString)

-- | Hash algorithm used in Tezos.
class Typeable alg => KnownHashAlgorithm (alg :: HashAlgorithmKind) where
  hashAlgorithmName :: Proxy alg -> Text
  hashAlgorithmName _ =
    toText . Refl.tyConName . Refl.typeRepTyCon $ Refl.typeRep @alg

  computeHash :: ByteString -> ByteString
  toHash :: BytesLike bs => bs : s :-> Hash alg bs : s

-- | Evaluate hash in Haskell world.
toHashHs
  :: forall alg bs. (BytesLike bs, KnownHashAlgorithm alg)
  => bs -> Hash alg bs
toHashHs = UnsafeHash . computeHash @alg . toBytes

-- | Documentation item for hash algorithms.
data DHashAlgorithm where
  DHashAlgorithm :: KnownHashAlgorithm alg => Proxy alg -> DHashAlgorithm

instance Eq DHashAlgorithm where
  a == b = (a `compare` b) == EQ
instance Ord DHashAlgorithm where
  DHashAlgorithm a `compare` DHashAlgorithm b =
    hashAlgorithmName a `compare` hashAlgorithmName b

instance DocItem DHashAlgorithm where
  type DocItemPlacement DHashAlgorithm = 'DocItemInDefinitions
  type DocItemReferenced DHashAlgorithm = 'True
  docItemPos = 5310
  docItemSectionName = Just "Referenced hash algorithms"

  docItemRef (DHashAlgorithm alg) =
    DocItemRef $ DocItemId ("hash-alg-" <> hashAlgorithmName alg)

  docItemToMarkdown _ (DHashAlgorithm alg) =
    "* " <> build (hashAlgorithmName alg)

-- Creates a reference to given hash algorithm description.
hashAlgorithmMdRef :: KnownHashAlgorithm alg => Proxy alg -> Markdown
hashAlgorithmMdRef alg =
  mdLocalRef
    (mdTicked . build $ hashAlgorithmName alg)
    (docItemRef (DHashAlgorithm alg))

data Sha256 :: HashAlgorithmKind

instance KnownHashAlgorithm Sha256 where
  computeHash = Crypto.sha256
  toHash = I T.SHA256

data Sha512 :: HashAlgorithmKind

instance KnownHashAlgorithm Sha512 where
  computeHash = Crypto.sha512
  toHash = I T.SHA512

data Blake2b :: HashAlgorithmKind

instance KnownHashAlgorithm Blake2b where
  computeHash = Crypto.blake2b
  toHash = I T.BLAKE2B

data Sha3 :: HashAlgorithmKind

instance KnownHashAlgorithm Sha3 where
  computeHash = Crypto.sha3
  toHash = I T.SHA3

data Keccak :: HashAlgorithmKind

instance KnownHashAlgorithm Keccak where
  computeHash = Crypto.keccak
  toHash = I T.KECCAK

----------------------------------------------------------------------------
-- Typed Chest
----------------------------------------------------------------------------

newtype ChestT a = ChestT { unChestT :: Chest }
  deriving newtype (IsoValue, HasAnnotation)
  deriving stock Generic

instance TypeHasDoc a => TypeHasDoc (ChestT a) where
  typeDocMdDescription = [md|
    Timelock puzzle chest containing a typed value.

    In Lorentz, use `openChestT` instead of `openChest` to open it.
    |]
  typeDocMdReference = poly1TypeDocMdReference
  typeDocDependencies p =
    genericTypeDocDependencies p <>
    [ dTypeDep @a
    , dTypeDep @MText, dTypeDep @Integer  -- for examples below
    ]
  typeDocHaskellRep = concreteTypeDocHaskellRep @(ChestT (Packed (MText, Integer)))
  typeDocMichelsonRep = concreteTypeDocMichelsonRep @(ChestT (Packed (MText, Integer)))

data OpenChestT a = ChestContentT a | ChestOpenFailedT Bool
  deriving stock (Generic, Show, Eq)
  deriving anyclass (T.IsoValue, HasAnnotation)

instance (TypeHasDoc a) => TypeHasDoc (OpenChestT a) where
  typeDocMdDescription = "Typed result of opening a typed timelocked chest."
  typeDocMdReference = poly1TypeDocMdReference
  typeDocDependencies _ =
    [ dTypeDep @a
    , dTypeDep @MText, dTypeDep @Integer  -- for examples below
    ]
  typeDocHaskellRep _ _ = Nothing
  typeDocMichelsonRep = concreteTypeDocMichelsonRep @(OpenChestT (Packed (MText, Integer)))

openChestT :: BytesLike a => ChestKey : ChestT a : Natural : s :-> OpenChestT a : s
openChestT = I T.OPEN_CHEST