-- | 'UStore' definition and common type-level stuff.
module Lorentz.UStore.Types
  ( -- * UStore and related type definitions
    UStore (..)
  , type (|~>)(..)
  , UStoreFieldExt (..)
  , UStoreField
  , UStoreMarkerType
  , UMarkerPlainField

    -- ** Extras
  , KnownUStoreMarker (..)
  , mkFieldMarkerUKeyL
  , mkFieldUKey
  , UStoreSubmapKey
  , UStoreSubmapKeyT

    -- ** Type-lookup-by-name
  , GetUStoreKey
  , GetUStoreValue
  , GetUStoreField
  , GetUStoreFieldMarker

    -- ** Marked fields
  , PickMarkedFields

   -- * Internals
  , ElemSignature (..)
  , GetUStore
  , MSKey
  , MSValue
  , FSValue
  , FSMarker
  ) where

import Data.Default (Default)
import Control.Lens (Wrapped)
import Data.Vinyl.Derived (Label)
import qualified Data.Kind as Kind
import Data.Type.Equality (type (==))
import GHC.Generics ((:*:)(..), (:+:)(..))
import qualified GHC.Generics as G
import GHC.TypeLits (ErrorMessage(..), Symbol, TypeError, KnownSymbol)
import Test.QuickCheck (Arbitrary)

import Lorentz.Pack
import Lorentz.Doc
import Lorentz.Polymorphic
import Michelson.Typed.T
import Michelson.Typed.Haskell.Value
import Lorentz.UStore.Common
import Lorentz.Value
import Util.Type

-- | Gathers multple fields and 'BigMap's under one object.
--
-- Type argument of this datatype stands for a "store template" -
-- a datatype with one constructor and multiple fields, each containing
-- an object of type 'UStoreField' or '|~>' and corresponding to single
-- virtual field or 'BigMap' respectively.
-- It's also possible to parameterize it with a larger type which is
-- a product of types satisfying the above property.
newtype UStore (a :: Kind.Type) = UStore
  { unUStore :: BigMap ByteString ByteString
  } deriving stock (Eq, Show, Generic)
    deriving newtype (Default, Semigroup, Monoid, IsoValue,
                      MemOpHs, GetOpHs, UpdOpHs)

instance Wrapped (UStore a)

-- | Describes one virtual big map in the storage.
newtype k |~> v = UStoreSubMap { unUStoreSubMap :: Map k v }
  deriving stock (Show, Eq)
  deriving newtype (Default, Arbitrary)

-- | Describes plain field in the storage.
newtype UStoreFieldExt (m :: UStoreMarkerType) (v :: Kind.Type) = UStoreField { unUStoreField :: v }
  deriving stock (Show, Eq)
  deriving newtype Arbitrary

-- | Just a servant type.
data UStoreMarker

-- | Specific kind used to designate markers for 'UStoreFieldExt'.
--
-- We suggest that fields may serve different purposes and so annotated with
-- special markers accordingly. See example below.
--
-- This kind is implemented like that because we want markers to differ from all
-- other types in kind; herewith 'UStoreMarkerType' is still an open kind
-- (has potentially infinite number of inhabitants).
type UStoreMarkerType = UStoreMarker -> Kind.Type

-- | Just a plain field used as data.
type UStoreField = UStoreFieldExt UMarkerPlainField
data UMarkerPlainField :: UStoreMarkerType

instance Typeable template => TypeHasDoc (UStore template) where
  typeDocName _ = "Upgradeable storage"
  typeDocMdDescription =
    "Storage with not hardcoded structure, which allows upgrading the contract \
    \in place. UStore is capable of storing simple fields and multiple submaps. \
    \For simple fields key is serialized field name. For submap element big_map \
    \key is serialized `(submapName, keyValue)`."
  typeDocMdReference tp =
    customTypeDocMdReference ("UStore", DType tp) []
  typeDocHaskellRep = homomorphicTypeDocHaskellRep
  typeDocMichelsonRep = homomorphicTypeDocMichelsonRep

-- | What do we serialize when constructing big_map key for accessing
-- an UStore submap.
type UStoreSubmapKey k = (MText, k)
type UStoreSubmapKeyT k = 'TPair (ToT MText) k

-- Extra attributes of fields
----------------------------------------------------------------------------

-- | Allows to specify format of key under which fields of this type are stored.
-- Useful to avoid collisions.
class KnownUStoreMarker (marker :: UStoreMarkerType) where
  -- | By field name derive key under which field should be stored.
  mkFieldMarkerUKey :: MText -> ByteString
  default mkFieldMarkerUKey :: MText -> ByteString
  mkFieldMarkerUKey = lPackValue

  -- | Display type-level information about UStore field with given marker and
  -- field value type.
  -- Used for error messages.
  type ShowUStoreField marker v :: ErrorMessage
  type ShowUStoreField marker v = 'Text "field of type " ':<>: 'ShowType v

-- | Version of 'mkFieldMarkerUKey' which accepts label.
mkFieldMarkerUKeyL
  :: forall marker field.
     (KnownUStoreMarker marker, KnownSymbol field)
  => Label field -> ByteString
mkFieldMarkerUKeyL _ =
  mkFieldMarkerUKey @marker (fieldNameToMText @field)

-- | Shortcut for 'mkFieldMarkerUKey' which accepts not marker but store template
-- and name of entry.
mkFieldUKey
  :: forall (store :: Kind.Type) field.
     (KnownSymbol field, KnownUStoreMarker (GetUStoreFieldMarker store field))
  => Label field -> ByteString
mkFieldUKey = mkFieldMarkerUKeyL @(GetUStoreFieldMarker store field)

instance KnownUStoreMarker UMarkerPlainField where

----------------------------------------------------------------------------
-- Type-safe lookup magic
----------------------------------------------------------------------------

{- Again we use generic magic to implement methods for 'Store'
(and thus 'Store' type constructor accepts a datatype, not a type-level list).

There are two reasons for this:

1. This gives us expected balanced tree of 'Or's for free.

2. This allows us selecting a map by field name, not by
e.g. type of map value. This is subjective, but looks like a good thing
for me (@martoon). On the other hand, it prevents us from sharing the
same interface between maps and 'Store'.

-}

-- | What was found on lookup by constructor name.
--
-- This keeps either type arguments of '|~>' or 'UStoreField'.
data ElemSignature
  = MapSignature Kind.Type Kind.Type
  | FieldSignature UStoreMarkerType Kind.Type

-- Again, we will use these getters instead of binding types within
-- 'MapSignature' using type equality because getters does not produce extra
-- compile errors on "field not found" cases.
type family MSKey (ms :: ElemSignature) :: Kind.Type where
  MSKey ('MapSignature k _) = k
  MSKey ('FieldSignature _ _) =
    TypeError ('Text "Expected UStore submap, but field was referred")
type family MSValue (ms :: ElemSignature) :: Kind.Type where
  MSValue ('MapSignature _ v) = v
  MSValue ('FieldSignature _ _) =
    TypeError ('Text "Expected UStore submap, but field was referred")
type family FSValue (ms :: ElemSignature) :: Kind.Type where
  FSValue ('FieldSignature _ v) = v
  FSValue ('MapSignature _ _) =
    TypeError ('Text "Expected UStore field, but submap was referred")
type family FSMarker (ms :: ElemSignature) :: UStoreMarkerType where
  FSValue ('FieldSignature m _) = m
  FSValue ('MapSignature _ _) =
    TypeError ('Text "Expected UStore field, but submap was referred")

-- | Get map signature from the constructor with a given name.
type GetUStore name a = MERequireFound name a (GLookupStore name (G.Rep a))

type family MERequireFound
  (name :: Symbol)
  (a :: Kind.Type)
  (mlr :: Maybe ElemSignature)
    :: ElemSignature where
  MERequireFound _ _ ('Just ms) = ms
  MERequireFound name a 'Nothing = TypeError
    ('Text "Failed to find plain field or submap in store template: datatype `"
     ':<>: 'ShowType a ':<>: 'Text "` has no field " ':<>: 'ShowType name)

type family GLookupStore (name :: Symbol) (x :: Kind.Type -> Kind.Type)
              :: Maybe ElemSignature where
  GLookupStore name (G.D1 _ x) = GLookupStore name x
  GLookupStore _ (_ :+: _) =
    TypeError ('Text "Templates used in UStore should have only one constructor")
  GLookupStore _ G.V1 =
    TypeError ('Text "No constructors in UStore template")

  GLookupStore name (G.C1 _ x) = GLookupStore name x

  GLookupStore name (x :*: y) = LSMergeFound name (GLookupStore name x)
                                                  (GLookupStore name y)

  -- When we encounter a field there are three cases we are interested in:
  -- 1. This field has type '|~>'. Then we check its name and return 'Just'
  -- with all required info on match, and 'Nothing' otherwise.
  -- 2. This field has type 'UStoreField'. We act in the same way
  -- as for '|~>', attaching 'ThePlainFieldKey' as key.
  -- 3. This field type is a different one. Then we expect this field to store
  -- '|~>' or 'UStoreField' somewhere deeper and try to find it there.
  GLookupStore name (G.S1 ('G.MetaSel mFieldName _ _ _) (G.Rec0 (k |~> v))) =
    Guard ('Just name == mFieldName) ('MapSignature k v)
  GLookupStore name (G.S1 ('G.MetaSel mFieldName _ _ _) (G.Rec0 (UStoreFieldExt m v))) =
    Guard ('Just name == mFieldName) ('FieldSignature m v)

  GLookupStore name (G.S1 _ (G.Rec0 a)) =
    GLookupStore name (G.Rep a)

  GLookupStore _ G.U1 = 'Nothing

type family LSMergeFound (name :: Symbol)
  (f1 :: Maybe ElemSignature) (f2 :: Maybe ElemSignature)
  :: Maybe ElemSignature where
  LSMergeFound _ 'Nothing 'Nothing = 'Nothing
  LSMergeFound _ ('Just ms) 'Nothing = 'Just ms
  LSMergeFound _ 'Nothing ('Just ms) = 'Just ms
  -- It's possible that there are two constructors with the same name,
  -- because main template pattern may be a sum of smaller template
  -- patterns with same constructor names.
  LSMergeFound ctor ('Just _) ('Just _) = TypeError
    ('Text "Found more than one constructor matching " ':<>: 'ShowType ctor)


-- | Get type of submap key.
type GetUStoreKey store name = MSKey (GetUStore name store)

-- | Get type of submap value.
type GetUStoreValue store name = MSValue (GetUStore name store)

-- | Get type of plain field.
-- This ignores marker with field type.
type GetUStoreField store name = FSValue (GetUStore name store)

-- | Get kind of field.
type GetUStoreFieldMarker store name = FSMarker (GetUStore name store)

-- One more magic
----------------------------------------------------------------------------

-- | Collect all fields with the given marker.
type PickMarkedFields marker template = GPickMarkedFields marker (G.Rep template)

type family GPickMarkedFields (marker :: UStoreMarkerType) (x :: Kind.Type -> Kind.Type)
             :: [(Symbol, Kind.Type)] where
  GPickMarkedFields m (G.D1 _ x) = GPickMarkedFields m x
  GPickMarkedFields m (G.C1 _ x) = GPickMarkedFields m x
  GPickMarkedFields m (x :*: y) = GPickMarkedFields m x ++ GPickMarkedFields m y
  GPickMarkedFields _ G.U1 = '[]

  GPickMarkedFields m (G.S1 ('G.MetaSel ('Just fieldName) _ _ _) (G.Rec0 (UStoreFieldExt m v))) =
    '[ '(fieldName, v) ]
  GPickMarkedFields _ (G.S1 _ (G.Rec0 (UStoreFieldExt _ _))) =
    '[]
  GPickMarkedFields _ (G.S1 _ (G.Rec0 (_ |~> _))) =
    '[]
  GPickMarkedFields m (G.S1 _ (G.Rec0 a)) =
    PickMarkedFields m a