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

{-# 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 Type.Reflection qualified as Refl

import Lorentz.Annotation
import Lorentz.Base
import Lorentz.Constraints.Scopes
import Lorentz.Doc
import Lorentz.Value
import Morley.AsRPC (HasRPCRepr(..))
import Morley.Michelson.Typed qualified as T
import Morley.Tezos.Crypto hiding (Hash)
import Morley.Tezos.Crypto.Hash qualified 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 HasRPCRepr (Packed a) where
  type AsRPC (Packed a) = Packed a

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 HasRPCRepr (TSignature a) where
  type AsRPC (TSignature a) = TSignature a

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 HasRPCRepr (Hash alg a) where
  type AsRPC (Hash alg a) = Hash alg a

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 HasRPCRepr (ChestT a) where
  type AsRPC (ChestT a) = ChestT a

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 HasRPCRepr a => HasRPCRepr (OpenChestT a) where
  type AsRPC (OpenChestT a) = OpenChestT (AsRPC a)

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