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

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

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

import qualified Unsafe

import Control.Monad.Except (runExcept, throwError)
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 Lorentz.Base
import Lorentz.Coercions
import Lorentz.Constraints
import qualified Lorentz.Instr as L
import Lorentz.Pack
import Lorentz.UStore.Migration
import Lorentz.UStore.Migration.Diff
import Lorentz.UStore.Traversal
import Lorentz.UStore.Types
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
  :: (UStoreTraversable MkUStoreTW template)
  => template -> UStore template
mkUStore = UStore . BigMap . mkUStoreInternal

-- | 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.
     (UStoreTraversable DecomposeUStoreTW template)
  => UStore template -> Either Text (UStoreContent, template)
ustoreDecompose = storeDecomposeInternal . Map.toList . unBigMap . unUStore

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

-- | Make migration script which initializes 'UStore' from scratch.
fillUStore
  :: (UStoreTraversable FillUStoreTW template)
  => template -> UStoreMigration () template
fillUStore v = UStoreMigration $ fillUStoreInternal v

-- | Version of 'migrateFillUStore' for batched migrations.
--
-- Each field write will be placed to a separate batch.
fillUStoreMigrationBlock
  :: ( UStoreTraversable FillUStoreTW 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 $ fillUStoreInternal 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
  :: ( UStoreTraversable FillUStoreTW 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 = fillUStoreInternal 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
----------------------------------------------------------------------------

-- | Internal helper for 'mkUStore'.
mkUStoreInternal
  :: (UStoreTraversable MkUStoreTW template)
  => template -> Map ByteString ByteString
mkUStoreInternal = foldUStore MkUStoreTW

-- | Internal helper for 'ustoreDecompose'.
storeDecomposeInternal
  :: forall template.
     (UStoreTraversable DecomposeUStoreTW template)
  => UStoreContent -> Either Text (UStoreContent, template)
storeDecomposeInternal =
  runExcept . fmap swap . runStateT (genUStore DecomposeUStoreTW)

-- | Internal helper for 'fillUStore'.
fillUStoreInternal
  :: (UStoreTraversable FillUStoreTW template)
  => template
  -> [MigrationAtom]
fillUStoreInternal a = appEndo (foldUStore FillUStoreTW a) []

-- | Declares handlers for UStore creation from template.
data MkUStoreTW = MkUStoreTW

instance UStoreTraversalWay MkUStoreTW where
  type UStoreTraversalArgumentWrapper MkUStoreTW = Identity
  type UStoreTraversalMonad MkUStoreTW = Const (Map ByteString ByteString)

instance (NicePackedValue val) =>
         UStoreTraversalFieldHandler MkUStoreTW marker val where
  ustoreTraversalFieldHandler MkUStoreTW fieldName (Identity val) =
    Const $
    one ( mkFieldMarkerUKeyL @marker fieldName
        , lPackValue val
        )

instance (NicePackedValue k, NicePackedValue v) =>
         UStoreTraversalSubmapHandler MkUStoreTW k v where
  ustoreTraversalSubmapHandler MkUStoreTW fieldName (Identity m) =
    Const $
    mconcat
      [ one ( lPackValue (labelToMText fieldName, k)
            , lPackValue v
            )
      | (k, v) <- Map.toList m
      ]

-- | Declares handlers for UStore conversion to template.
data DecomposeUStoreTW = DecomposeUStoreTW

instance UStoreTraversalWay DecomposeUStoreTW where
  type UStoreTraversalArgumentWrapper DecomposeUStoreTW = Const ()
  type UStoreTraversalMonad DecomposeUStoreTW =
    StateT UStoreContent (ExceptT Text Identity)

instance (NiceUnpackedValue val) =>
         UStoreTraversalFieldHandler DecomposeUStoreTW marker val where
  ustoreTraversalFieldHandler DecomposeUStoreTW fieldName (Const ()) = do
    let expectedKey = mkFieldMarkerUKey @marker (labelToMText fieldName)
    allMatched <- mapMaybesState $ \(key, value) -> do
      unless (key == expectedKey) mzero
      case lUnpackValue value of
        Left err -> throwError $
          "Failed to parse UStore value for field " +|
          demote @(ToT val) |+ ": " +| err |+ ""
        Right valValue ->
          pure valValue
    case allMatched of
        [] -> throwError $
          "Failed to find field in UStore: " +| labelToMText fieldName |+ ""
        [matched] ->
          pure matched
        (_ : _ : _) ->
          error "UStore content contained multiple entries with the same key"

instance (Ord k, NiceUnpackedValue k, NiceUnpackedValue v) =>
         UStoreTraversalSubmapHandler DecomposeUStoreTW k v where
  ustoreTraversalSubmapHandler _ fieldName (Const ()) =
    fmap Map.fromList $
    mapMaybesState $ \(key, value) ->
      case lUnpackValue @(UStoreSubmapKey k) key of
        Left _ -> mzero
        Right (name :: MText, keyValue :: k)
          | name /= labelToMText fieldName ->
              mzero
          | otherwise ->
            case lUnpackValue value of
              Left err -> throwError $
                "Failed to parse UStore value for " +|
                demote @(ToT k) |+ " |~> " +| demote @(ToT v) |+
                ": " +| err |+ ""
              Right valValue ->
                pure (keyValue, valValue)

-- | Declares handlers for UStore filling via lambda.
data FillUStoreTW = FillUStoreTW

instance UStoreTraversalWay FillUStoreTW where
  type UStoreTraversalArgumentWrapper FillUStoreTW = Identity
  type UStoreTraversalMonad FillUStoreTW = Const (Endo [MigrationAtom])

instance (NiceConstant v) =>
         UStoreTraversalFieldHandler FillUStoreTW marker v where
  ustoreTraversalFieldHandler FillUStoreTW fieldName (Identity val) =
    Const $
    Endo . (:) . formMigrationAtom Nothing $
      attachMigrationActionName (DAddAction "init field") 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 @marker fieldName) #
      L.update

instance (NiceConstant k, NiceConstant v) =>
         UStoreTraversalSubmapHandler FillUStoreTW k v where
  ustoreTraversalSubmapHandler _ fieldName (Identity m) =
    Const $
    Endo . (<>) $
    Map.toList m <&> \(k, v) ->
    formMigrationAtom Nothing $
      attachMigrationActionName (DAddAction "init submap") 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 (labelToMText fieldName) # L.pair #
      L.pack @(UStoreSubmapKey _) #
      L.update

-- | Tries to map all items in the state and returns those which were mapped
-- successfully; others are retained in the state.
mapMaybesState :: forall a b m. MonadState [a] m => (a -> MaybeT m b) -> m [b]
mapMaybesState mapper =
  get >>= \st -> do
    mapped <- mapM (\a -> (a, ) <$> runMaybeT (mapper a)) st
    let
      (passed, failed) =
        bimap (map (Unsafe.fromJust . snd)) (map fst) $
        L.partition @(a, Maybe b) (isJust . snd) $
        mapped
    put failed
    return passed

-- 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"))