{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Conversion between 'UStore' in Haskell and Michelson representation.
module Lorentz.UStore.Haskell
  ( UStoreContent
  , UStoreConversible
  , mkUStore
  , ustoreDecompose
  , ustoreDecomposeFull
  , fillUStore
  , migrateFillUStore
  , fillUStoreMigrationBlock
  ) where

import qualified Data.Kind as Kind
import qualified Data.List as L
import qualified Data.Map as Map
import Data.Singletons (demote)
import Fcf (type (=<<), Eval, Pure2)
import qualified Fcf
import Fmt ((+|), (+||), (|+), (||+))
import GHC.Generics ((:*:)(..), (:+:)(..))
import qualified GHC.Generics as G
import GHC.TypeLits (ErrorMessage(..), KnownSymbol, TypeError, symbolVal)

import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Constraints
import qualified Lorentz.Instr as L
import Lorentz.Pack
import Lorentz.UStore.Common
import Lorentz.UStore.Migration
import Lorentz.UStore.Migration.Diff
import Lorentz.UStore.Types
import Michelson.Interpret.Pack
import Michelson.Text
import Michelson.Typed.Haskell.Value
import Util.Type

-- | 'UStore' content represented as key-value pairs.
type UStoreContent = [(ByteString, ByteString)]

-- | Make 'UStore' from separate @big_map@s and fields.
mkUStore
  :: (UStoreConversible template)
  => template -> UStore template
mkUStore = UStore . BigMap . mkUStoreRec

-- | Decompose 'UStore' into separate @big_map@s and fields.
--
-- Since this function needs to @UNPACK@ content of @UStore@ to actual
-- keys and values, you have to provide 'UnpackEnv'.
--
-- Along with resulting value, you get a list of @UStore@ entries which
-- were not recognized as belonging to any submap or field according to
-- @UStore@'s template - this should be empty unless @UStore@ invariants
-- were violated.
ustoreDecompose
  :: forall template.
     (UStoreConversible template)
  => UStore template -> Either Text (UStoreContent, template)
ustoreDecompose = storeDecomposeRec . Map.toList . unBigMap . unUStore

-- | Make migration script which initializes 'UStore' from scratch.
fillUStore
  :: (UStoreConversible template)
  => template -> UStoreMigration () template
fillUStore v = UStoreMigration $ appEndo (fillUStoreRec v) []

-- | Version of 'migrateFillUStore' for batched migrations.
--
-- Each field write will be placed to a separate batch.
fillUStoreMigrationBlock
  :: ( UStoreConversible template
     , allFieldsExp ~ AllUStoreFieldsF template
     , newDiff ~ FillingNewDiff template diff
     , newTouched ~ FillingNewTouched template touched
     , PatternMatchL newDiff, PatternMatchL newTouched
     )
  => template
  -> MigrationBlocks oldTempl newTempl diff touched newDiff newTouched
fillUStoreMigrationBlock v = MigrationBlocks $ appEndo (fillUStoreRec v) []

-- | Fill 'UStore' with entries from the given template as part of simple
-- migration.
--
-- Sometimes you already have some fields initialized and 'fillUStore' does not
-- suit, then in case if your UStore template is a nested structure you can use
-- sub-templates to initialize the corresponding parts of UStore.
--
-- For batched migrations see 'fillUStoreMigrationBlock'.
migrateFillUStore
  :: ( UStoreConversible template
     , allFieldsExp ~ AllUStoreFieldsF template
     , newDiff ~ FillingNewDiff template diff
     , newTouched ~ FillingNewTouched template touched
     , PatternMatchL newDiff, PatternMatchL newTouched
     )
  => template
  -> Lambda
       (MUStore oldTempl newTempl diff touched)
       (MUStore oldTempl newTempl newDiff newTouched)
migrateFillUStore v =
  let atoms = appEndo (fillUStoreRec v) []
      script = foldMap (unMigrationScript . maScript) atoms
  in forcedCoerce_ # script # forcedCoerce_

type FillingNewDiff template diff =
  CoverDiffMany diff
    (Eval (Fcf.Map (Pure2 '(,) 'DcAdd) =<< LinearizeUStoreF template))

type FillingNewTouched template touched =
  Eval (AllUStoreFieldsF template) ++ touched

-- Implementation
----------------------------------------------------------------------------

-- | Like 'ustoreDecompose', but requires all entries from @UStore@ to be
-- recognized.
ustoreDecomposeFull
  :: forall template.
     (UStoreConversible template)
  => UStore template -> Either Text template
ustoreDecomposeFull ustore = do
  (remained, res) <- ustoreDecompose ustore
  unless (null remained) $
    Left $ "Unrecognized entries in UStore: " +|| remained ||+ ""
  return res

-- | Recursive template traversal for 'mkUStore'.
mkUStoreRec
  :: (UStoreConversible template)
  => template -> Map ByteString ByteString
mkUStoreRec = gUstoreToVal . G.from

-- | Recursive template traversal for 'ustoreDecompose'.
storeDecomposeRec
  :: forall template.
     (UStoreConversible template)
  => UStoreContent -> Either Text (UStoreContent, template)
storeDecomposeRec = fmap (second G.to) ... gUstoreFromVal

-- | Recursive template traversal for 'fillUStore'.
fillUStoreRec
  :: (UStoreConversible template)
  => template
  -> Endo [MigrationAtom]
fillUStoreRec = gUstoreToScript . G.from

-- | Given template can be converted to 'UStore' value.
class (Generic template, GUStoreConversible (G.Rep template)) =>
       UStoreConversible template
instance (Generic template, GUStoreConversible (G.Rep template)) =>
         UStoreConversible template

-- | Generic traversal for 'mkUStore' and 'ustoreDecompose'.
class GUStoreConversible (template :: Kind.Type -> Kind.Type) where
  -- | Convert generic value to internal 'UStore' representation.
  gUstoreToVal :: template p -> Map ByteString ByteString

  -- | Parse internal 'UStore' representation into generic Haskell value of
  -- 'UStore', also returning unparsed entries.
  gUstoreFromVal
    :: UStoreContent
    -> Either Text (UStoreContent, template p)

  -- | Convert generic value to a script filling the corresponding 'UStore'.
  gUstoreToScript
    :: template p
    -> Endo [MigrationAtom]

instance GUStoreConversible x => GUStoreConversible (G.D1 i x) where
  gUstoreToVal = gUstoreToVal . G.unM1
  gUstoreFromVal = fmap (second G.M1) ... gUstoreFromVal
  gUstoreToScript = gUstoreToScript . G.unM1

instance GUStoreConversible x => GUStoreConversible (G.C1 i x) where
  gUstoreToVal = gUstoreToVal . G.unM1
  gUstoreFromVal = fmap (second G.M1) ... gUstoreFromVal
  gUstoreToScript = gUstoreToScript . G.unM1

instance TypeError ('Text "Unexpected sum type in UStore template") =>
         GUStoreConversible (x :+: y) where
  gUstoreToVal = error "impossible"
  gUstoreFromVal = error "impossible"
  gUstoreToScript = error "impossible"

instance TypeError ('Text "UStore template should have one constructor") =>
         GUStoreConversible G.V1 where
  gUstoreToVal = error "impossible"
  gUstoreFromVal = error "impossible"
  gUstoreToScript = error "impossible"

instance (GUStoreConversible x, GUStoreConversible y) =>
         GUStoreConversible (x :*: y) where
  gUstoreToVal (x :*: y) = gUstoreToVal x <> gUstoreToVal y
  gUstoreFromVal entries = do
    (entries', res1) <- gUstoreFromVal entries
    (entries'', res2) <- gUstoreFromVal entries'
    return (entries'', res1 :*: res2)
  gUstoreToScript (x :*: y) = gUstoreToScript x <> gUstoreToScript y

instance GUStoreConversible G.U1 where
  gUstoreToVal G.U1 = mempty
  gUstoreFromVal entries = pure (entries, G.U1)
  gUstoreToScript G.U1 = mempty

-- | Case of nested template.
instance {-# OVERLAPPABLE #-}
         (UStoreConversible template) =>
         GUStoreConversible (G.S1 i (G.Rec0 template)) where
  gUstoreToVal = mkUStoreRec . G.unK1 . G.unM1
  gUstoreFromVal = fmap (second $ G.M1 . G.K1) ... storeDecomposeRec
  gUstoreToScript = fillUStoreRec . G.unK1 . G.unM1

-- | Case of '|~>'.
instance ( NiceFullPackedValue k, NiceFullPackedValue v
         , KnownSymbol fieldName, Ord k
         ) =>
         GUStoreConversible (G.S1 ('G.MetaSel ('Just fieldName) _1 _2 _3)
                                  (G.Rec0 (k |~> v))) where
  gUstoreToVal (G.M1 (G.K1 (UStoreSubMap m))) =
    mconcat
      [ one ( lPackValue (fieldNameToMText @fieldName, k)
            , lPackValue v
            )
      | (k, v) <- Map.toList m
      ]

  gUstoreFromVal allEntries = do
    (unrecognized, res) <- foldM parseEntry (mempty, mempty) allEntries
    return (unrecognized, G.M1 . G.K1 $ UStoreSubMap res)
    where
    parseEntry
      :: (UStoreContent, Map k v)
      -> (ByteString, ByteString)
      -> Either Text (UStoreContent, Map k v)
    parseEntry (entries, !acc) entry@(key, val) =
      case lUnpackValue @(UStoreSubmapKey _) key of
        Left _ -> Right (entry : entries, acc)
        Right (name :: MText, keyValue :: k)
          | toText name /= toText (symbolVal $ Proxy @fieldName) ->
              Right (entry : entries, acc)
          | otherwise ->
            case lUnpackValue val of
              Left err ->
                Left $ "Failed to parse UStore value for " +|
                      demote @(ToT k) |+ " |~> " +| demote @(ToT v) |+
                      ": " +| err |+ ""
              Right valValue ->
                Right (entries, Map.insert keyValue valValue acc)

  gUstoreToScript (G.M1 (G.K1 (UStoreSubMap m))) = Endo . (<>) $
    Map.toList m <&> \(k, v) ->
    formMigrationAtom Nothing $
      attachMigrationActionName (DAddAction "init submap") (fromLabel @fieldName) (Proxy @v) #
      -- @PUSH + PACK@ will be merged by optimizer, but there is still place
      -- for further improvement both or value and key pushing.
      -- We cannot push already packed value because that would break code
      -- analyzers and transformers, consider adding necessary rules to
      -- optimizer.
      -- TODO [TM-379]: consider improving this case
      -- or
      -- TODO: add necessary rules to optimizer
      L.push v # L.pack # L.some #
      L.push k # L.push (fieldNameToMText @fieldName) # L.pair #
      L.pack @(UStoreSubmapKey _) #
      L.update

-- | Case of 'UStoreField'.
instance (NiceFullPackedValue v, KnownUStoreMarker m, KnownSymbol fieldName) =>
         GUStoreConversible (G.S1 ('G.MetaSel ('Just fieldName) _1 _2 _3)
                                  (G.Rec0 (UStoreFieldExt m v))) where
  gUstoreToVal (G.M1 (G.K1 (UStoreField val))) =
    one ( mkFieldMarkerUKeyL @m (fromLabel @fieldName)
        , lPackValue val
        )

  gUstoreFromVal entries =
    let key = packValue' $ toVal (fieldNameToMText @fieldName)
    in case L.partition ((== key) . fst) entries of
        ([], _) ->
          Left $ "Failed to find field in UStore: " +|
                fieldNameToMText @fieldName |+ ""
        ([(_, val)], otherEntries) ->
          case lUnpackValue val of
            Left err ->
              Left $ "Failed to parse UStore value for field " +|
                    demote @(ToT v) |+ ": " +| err |+ ""
            Right valValue ->
              Right (otherEntries, G.M1 . G.K1 $ UStoreField valValue)
        (_ : _ : _, _) ->
          error "UStore content contained multiple entries with the same key"

  gUstoreToScript (G.M1 (G.K1 (UStoreField val))) =
    Endo . (:) . formMigrationAtom Nothing $
      attachMigrationActionName (DAddAction "init field") (fromLabel @fieldName) (Proxy @v) #
      -- Not pushing already packed value (which would be more efficient) because
      -- analyzers cannot work with packed values.
      -- TODO: make optimizer compress this to @push (Just $ lPackValue val)@
      L.push val # L.pack # L.some #

      L.push (mkFieldMarkerUKeyL @m (fromLabel @fieldName)) #
      L.update

-- Examples
----------------------------------------------------------------------------

data MyStoreTemplate = MyStoreTemplate
  { ints :: Integer |~> ()
  , flag :: UStoreField Bool
  }
  deriving stock (Generic)

data MyStoreTemplateBig = MyStoreTemplateBig
  { templ :: MyStoreTemplate
  , bytes :: ByteString |~> ByteString
  }
  deriving stock (Generic)

_storeSample :: UStore MyStoreTemplate
_storeSample = mkUStore
  MyStoreTemplate
  { ints = UStoreSubMap $ one (1, ())
  , flag = UStoreField False
  }

_storeSampleBig :: UStore MyStoreTemplateBig
_storeSampleBig = mkUStore $
  MyStoreTemplateBig
    MyStoreTemplate
      { ints = UStoreSubMap $ one (1, ())
      , flag = UStoreField False
      }
    (UStoreSubMap $ one ("a", "b"))