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

-- | UStore templates generic traversals.
--
-- Normally you work with functionality of this module as follows:
-- 1. Pick the function fitting most for your traversal, one of
--    'traverseUStore', 'foldUStore' e.t.c.
-- 2. Create a custom datatype value of which will be put to that function.
-- 3. Implement a respective 'UStoreTemplateTraversable' instance for this
--    datatype.
module Lorentz.UStore.Traversal
  ( UStoreTraversalWay (..)
  , UStoreTraversalFieldHandler (..)
  , UStoreTraversalSubmapHandler (..)
  , UStoreTraversable
  , traverseUStore
  , modifyUStore
  , foldUStore
  , genUStore
  ) where

import qualified Data.Kind as Kind
import GHC.Generics ((:*:)(..), (:+:))
import qualified GHC.Generics as G

import Lorentz.UStore.Types
import Util.Label
import Util.TypeLits

----------------------------------------------------------------------------
-- Interface
----------------------------------------------------------------------------

-- | Defines general parameters of UStore template traversal.
-- You need a separate @way@ datatype with an instance of this typeclass for each
-- type of traversal.
class ( Applicative (UStoreTraversalArgumentWrapper way)
      , Applicative (UStoreTraversalMonad way)
      ) =>
      UStoreTraversalWay (way :: Kind.Type) where

  -- | Wrapper which will accompany the existing value of traversed template,
  -- aka argument.
  -- This is usually @'Identity'@ or @'Const' a@.
  type UStoreTraversalArgumentWrapper way :: Kind.Type -> Kind.Type

  -- | Additional constraints on monadic action used in traversal.
  -- Common ones include 'Identity', @'Const'@, @(,) a@
  type UStoreTraversalMonad way :: Kind.Type -> Kind.Type

-- | Declares a handler for UStore fields when given traversal way is applied.
class (UStoreTraversalWay way) =>
      UStoreTraversalFieldHandler
        (way :: Kind.Type) (marker :: UStoreMarkerType) (v :: Kind.Type) where
  -- | How to process each of UStore fields.
  ustoreTraversalFieldHandler
    :: (KnownUStoreMarker marker)
    => way
    -> Label name
    -> UStoreTraversalArgumentWrapper way v
    -> UStoreTraversalMonad way v

-- | Declares a handler for UStore submaps when given traversal way is applied.
class (UStoreTraversalWay way) =>
      UStoreTraversalSubmapHandler
        (way :: Kind.Type) (k :: Kind.Type) (v :: Kind.Type) where
  -- | How to process each of UStore submaps.
  ustoreTraversalSubmapHandler
    :: way
    -> Label name
    -> UStoreTraversalArgumentWrapper way (Map k v)
    -> UStoreTraversalMonad way (Map k v)

-- | Constraint for UStore traversal.
type UStoreTraversable way a =
  (Generic a, GUStoreTraversable way (G.Rep a), UStoreTraversalWay way)

-- | Perform UStore traversal. The most general way to perform a traversal.
traverseUStore
  :: forall way template.
     (UStoreTraversable way template)
  => way
  -> UStoreTraversalArgumentWrapper way template
  -> UStoreTraversalMonad way template
traverseUStore way =
  fmap G.to . gTraverseUStore way . fmap G.from

-- | Modify each UStore entry.
modifyUStore
  :: ( UStoreTraversable way template
     , UStoreTraversalArgumentWrapper way ~ Identity
     , UStoreTraversalMonad way ~ Identity
     )
  => way
  -> template
  -> template
modifyUStore way a =
  runIdentity $ traverseUStore way (Identity a)

-- | Collect information about UStore entries into monoid.
foldUStore
  :: ( UStoreTraversable way template
     , UStoreTraversalArgumentWrapper way ~ Identity
     , UStoreTraversalMonad way ~ Const res
     )
  => way
  -> template
  -> res
foldUStore way a =
  getConst $ traverseUStore way (Identity a)

-- | Fill UStore template with entries.
genUStore
  :: ( UStoreTraversable way template
     , UStoreTraversalArgumentWrapper way ~ Const ()
     )
  => way -> UStoreTraversalMonad way template
genUStore way =
  traverseUStore way (Const ())

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

-- | Generic traversal of UStore template.
class GUStoreTraversable (way :: Kind.Type) (x :: Kind.Type -> Kind.Type) where
  gTraverseUStore
    :: (UStoreTraversalWay way)
    => way
    -> UStoreTraversalArgumentWrapper way (x p)
    -> UStoreTraversalMonad way (x p)

instance GUStoreTraversable way x =>
         GUStoreTraversable way (G.D1 i x) where
  gTraverseUStore way x =
    G.M1 <$> gTraverseUStore way (G.unM1 <$> x)

instance GUStoreTraversable way x =>
         GUStoreTraversable way (G.C1 i x) where
  gTraverseUStore way x =
    G.M1 <$> gTraverseUStore way (G.unM1 <$> x)

instance TypeError ('Text "Unexpected sum type in UStore template") =>
         GUStoreTraversable way (x :+: y) where
  gTraverseUStore _ = error "imposible"

instance TypeError ('Text "Unexpected void-like type in UStore template") =>
         GUStoreTraversable way G.V1 where
  gTraverseUStore _ = error "impossible"

instance ( GUStoreTraversable way x
         , GUStoreTraversable way y
         ) =>
         GUStoreTraversable way (x :*: y) where
  gTraverseUStore way a =
    (:*:) <$> gTraverseUStore way (a <&> \(x :*: _) -> x)
          <*> gTraverseUStore way (a <&> \(_ :*: y) -> y)

instance GUStoreTraversable way G.U1 where
  gTraverseUStore _ _ = pure G.U1

instance {-# OVERLAPPABLE #-}
         UStoreTraversable way template =>
         GUStoreTraversable way (G.S1 i (G.Rec0 template)) where
  gTraverseUStore way sub =
    G.M1 . G.K1 <$> traverseUStore way (G.unK1 . G.unM1 <$> sub)

instance ( UStoreTraversalFieldHandler way marker v, KnownUStoreMarker marker
         , KnownSymbol ctor
         ) =>
         GUStoreTraversable
           way
           (G.S1 ('G.MetaSel ('Just ctor) _1 _2 _3) (G.Rec0 (UStoreFieldExt marker v))) where
  gTraverseUStore way entry =
    G.M1 . G.K1 . UStoreField <$>
      ustoreTraversalFieldHandler
        @_
        @marker
        way
        (Label @ctor)
        (entry <&> \(G.M1 (G.K1 (UStoreField v))) -> v)

instance (UStoreTraversalSubmapHandler way k v, KnownSymbol ctor) =>
         GUStoreTraversable
           way
           (G.S1 ('G.MetaSel ('Just ctor) _1 _2 _3) (G.Rec0 (k |~> v))) where
  gTraverseUStore way entry =
    G.M1 . G.K1 . UStoreSubMap <$>
      ustoreTraversalSubmapHandler
        way
        (Label @ctor)
        (entry <&> \(G.M1 (G.K1 (UStoreSubMap m))) -> m)