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

-- | Different approaches to batching.
--
-- For now we do not support perfect batching because operation size evaluation
-- (as well as gas consumption evaluation) is not implemented yet.
-- The only non-trivial batching implementation we provide is
-- 'mbSeparateLambdas'.
module Lorentz.UStore.Migration.Batching
  ( -- * Separate-lambdas batching
    SlBatchType (..)
  , SlBatchInfo (..)
  , mbSeparateLambdas
  ) where

import qualified Data.List as L
import Fmt (Buildable(..))
import System.Console.Pretty (Color(..), color)

import Lorentz.UStore.Migration.Base
import Michelson.Typed

----------------------------------------------------------------------------
-- Separating lambdas
----------------------------------------------------------------------------

-- | Type of batch.
data SlBatchType
  = SlbtData
    -- ^ Addition of any type of data.
  | SlbtLambda
    -- ^ Addition of code.
  | SlbtCustom
    -- ^ Several joined actions of different types.
  | SlbtUnknown
    -- ^ No information to chooseType about batching.
    -- This means that the given action does not contain 'DMigrationActionDesc'.
  deriving stock (Show, Eq)

slbtIsData :: SlBatchType -> Bool
slbtIsData = \case { SlbtData -> True; _ -> False }

data SlBatchInfo = SlBatchInfo
  { slbiType :: SlBatchType
  , slbiActions :: [Text]
  }

instance Buildable SlBatchInfo where
  build (SlBatchInfo ty actions) = mconcat
    [ build @Text $ case ty of
        SlbtData -> color Blue "[data]"
        SlbtLambda -> color Green "[code]"
        SlbtCustom -> color Yellow "[custom]"
        SlbtUnknown -> color Red "[unknown]"
    , " "
    , case actions of
        [] -> "-"
        [a] -> build a
        as -> foldMap (\a -> "\n  * " <> build a) as
    ]

-- | Puts all data updates in one batch, and all lambdas in separate batches,
-- one per batch.
--
-- The reason for such behaviour is that in production contracts amount of
-- changed data (be it in contract initialization or contract upgrade) is small,
-- while stored entrypoints are huge and addition of even one entrypoint often
-- barely fits into gas limit.
mbSeparateLambdas :: MigrationBatching [] SlBatchInfo
mbSeparateLambdas = MigrationBatching $ \atoms ->
  let
    atomsWithType = atoms <&> \a -> (atomType a, a)
    (dataAtoms, otherAtoms) = L.partition (slbtIsData . fst) atomsWithType
    dataMigration =
      ( SlBatchInfo SlbtData (nubCounting $ maName . snd <$> dataAtoms)
      , manualConcatMigrationScripts (maScript . snd <$> dataAtoms)
      )
    otherMigrations =
      [ (SlBatchInfo ty [maName atom], maScript atom)
      | (ty, atom) <- otherAtoms
      ]
  in dataMigration : otherMigrations
  where
    atomType :: MigrationAtom -> SlBatchType
    atomType = chooseType . maActionsDesc

    chooseType :: [DMigrationActionDesc] -> SlBatchType
    chooseType = \case
      [] -> SlbtUnknown
      xs | all isLambda xs -> SlbtLambda
      xs | (not . any isAddLambda) xs -> SlbtData
         | otherwise -> SlbtCustom

    isLambda :: DMigrationActionDesc -> Bool
    isLambda = \case { TLambda{} -> True; _ -> False } . manFieldType

    isAddLambda :: DMigrationActionDesc -> Bool
    isAddLambda a = and
      [ isLambda a
      , case manAction a of { DAddAction _ -> True; _ -> False }
      ]

-- | Similar to 'nub', counts number of invocations and attaches to text entry.
--
-- >>> nubCounting ["a", "b", "a"]
-- ["a (x2)", "b"]
nubCounting :: [Text] -> [Text]
nubCounting = \case
  [] -> []
  x : xs ->
    let ((length -> repetitions), others) = L.partition (== x) xs
        x' = if repetitions == 0
             then x
             else x <> " (x" <> show (repetitions + 1) <> ")"
    in x' : nubCounting others