{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- | Replicated Growable Array (RGA)
module RON.Data.RGA (
    RGA (..),
    RgaRep,
    RgaString,
    edit,
    editText,
    getAliveIndices,
    getList,
    getText,
    insert,
    insertAfter,
    insertAtBegin,
    insertText,
    insertTextAfter,
    insertTextAtBegin,
    newFromList,
    newFromText,
    remove,
    rgaType,
    RON.Data.RGA.toList,
    toText,
)
where

import Data.Algorithm.Diff (PolyDiff (Both, First, Second), getGroupedDiffBy)
import Data.Bifunctor (second)
import Data.HashMap.Strict qualified as HashMap
import Data.Map.Strict ((!?))
import Data.Map.Strict qualified as Map
import Data.Text qualified as Text

import RON.Data.Internal (
    MonadObjectState,
    ReducedChunk (ReducedChunk, rcBody, rcRef),
    Reducible,
    Rep,
    Replicated (encoding),
    ReplicatedAsObject,
    ReplicatedAsPayload,
    Unapplied,
    applyPatches,
    fromRon,
    getObjectStateChunk,
    modifyObjectStateChunk_,
    newObject,
    newRon,
    objectEncoding,
    readObject,
    reduceUnappliedPatches,
    reducibleOpType,
    stateFromChunk,
    stateToChunk,
    toPayload,
 )
import RON.Error (MonadE, errorContext, throwErrorText)
import RON.Event (ReplicaClock, getEventUuid, getEventUuids)
import RON.Prelude
import RON.Semilattice (Semilattice)
import RON.Types (
    ObjectRef (ObjectRef),
    Op (Op, opId, payload, refId),
    StateChunk (StateChunk),
    StateFrame,
    UUID,
    WireStateChunk (WireStateChunk, stateBody, stateType),
 )
import RON.UUID (pattern Zero)
import RON.UUID qualified as UUID
import RON.Util.Word (ls60, pattern B11)

{-# ANN module ("HLint: ignore Reduce duplication" :: String) #-}

{- | opId = vertex id
    refId:
        0 = value is alive,
        _ = tombstone event, value is backup for undo
    payload: the value
-}
type Vertex = Op

data VertexListItem
    = VertexListItem
    { itemValue :: Vertex
    , itemNext :: Maybe UUID
    }
    deriving (Eq, Show)

data VertexList
    = VertexList
    { listHead :: UUID
    , listItems :: HashMap UUID VertexListItem
    }
    deriving (Eq, Show)

instance Semigroup VertexList where
    (<>) = merge

instance Semilattice VertexList

vertexListToOps :: VertexList -> [Vertex]
vertexListToOps v@VertexList{..} = go listHead listItems
  where
    go root items =
        let VertexListItem{..} =
                HashMap.lookupDefault
                    ( error $
                        unlines $
                            ["Cannot find vertex id", show root, "in array"]
                                ++ map show (HashMap.toList items)
                                ++ ["Original array is", show v]
                    )
                    root
                    items
            rest = case itemNext of
                Just next -> go next (HashMap.delete root items)
                Nothing -> []
        in  itemValue : rest

vertexListFromOps :: [Vertex] -> Maybe VertexList
vertexListFromOps = foldr go mempty
  where
    go v@Op{opId} vlist = Just VertexList{listHead = opId, listItems = vlist'}
      where
        item itemNext = VertexListItem{itemValue = v, itemNext}
        vlist' = case vlist of
            Nothing -> HashMap.singleton opId (item Nothing)
            Just VertexList{listHead, listItems} ->
                HashMap.insert opId (item $ Just listHead) listItems

-- | Untyped RGA
newtype RgaRep = RgaRep (Maybe VertexList)
    deriving (Eq, Monoid, Semigroup, Semilattice, Show)

data PatchSet
    = PatchSet
    { psPatches :: Map UUID VertexList
    -- ^ the key is the parent event, the value is a non-empty VertexList
    , psRemovals :: Map UUID UUID
    -- ^ the key is the target event, the value is the tombstone event
    }
    deriving (Eq, Show)

instance Semigroup PatchSet where
    rga1 <> rga2 = reapplyPatchSet $ preMerge rga1 rga2

preMerge :: PatchSet -> PatchSet -> PatchSet
preMerge (PatchSet p1 r1) (PatchSet p2 r2) =
    PatchSet
        { psPatches = Map.unionWith (<>) p1 p2
        , psRemovals = Map.unionWith max r1 r2
        }

instance Monoid PatchSet where
    mempty = PatchSet{psPatches = mempty, psRemovals = mempty}

patchSetFromRawOp :: Op -> PatchSet
patchSetFromRawOp op@Op{opId, refId, payload} = case payload of
    [] ->
        -- remove op
        mempty{psRemovals = Map.singleton refId opId}
    _ : _ ->
        -- append op
        mempty
            { psPatches =
                Map.singleton
                    refId
                    VertexList
                        { listHead = opId
                        , listItems =
                            HashMap.singleton
                                opId
                                VertexListItem
                                    { itemValue = op{refId = Zero}
                                    , itemNext = Nothing
                                    }
                        }
            }

patchSetFromChunk :: ReducedChunk -> PatchSet
patchSetFromChunk ReducedChunk{rcRef, rcBody} =
    case (UUID.split rcRef).uuidVersion of
        B11 ->
            -- derived event -- rm-patch compatibility
            foldMap patchSetFromRawOp rcBody
        _ ->
            -- patch
            case vertexListFromOps rcBody of
                Just patch -> mempty{psPatches = Map.singleton rcRef patch}
                Nothing -> mempty

instance Reducible RgaRep where
    reducibleOpType = rgaType

    stateFromChunk = RgaRep . vertexListFromOps

    stateToChunk (RgaRep rga) = maybe [] vertexListToOps rga

    applyPatches rga (patches, ops) =
        second patchSetToChunks . reapplyPatchSetToState rga $
            foldMap patchSetFromChunk patches <> foldMap patchSetFromRawOp ops

    reduceUnappliedPatches (patches, ops) =
        patchSetToChunks . reapplyPatchSet $
            foldMap patchSetFromChunk patches <> foldMap patchSetFromRawOp ops

patchSetToChunks :: PatchSet -> Unapplied
patchSetToChunks PatchSet{psPatches, psRemovals} =
    ( [ ReducedChunk{..}
      | (rcRef, vertices) <- Map.assocs psPatches
      , let rcBody = vertexListToOps vertices
      ]
    , [ Op{opId = tombstone, refId, payload = []}
      | (refId, tombstone) <- Map.assocs psRemovals
      ]
    )

reapplyPatchSet :: PatchSet -> PatchSet
reapplyPatchSet ps =
    continue ps [reapplyPatchesToOtherPatches, reapplyRemovalsToPatches]

reapplyPatchSetToState :: RgaRep -> PatchSet -> (RgaRep, PatchSet)
reapplyPatchSetToState rga ps =
    continue (rga, ps) [reapplyPatchesToState, reapplyRemovalsToState]

continue :: x -> [x -> Maybe x] -> x
continue x fs = case asum $ map ($ x) fs of
    Nothing -> x
    Just x' -> continue x' fs

reapplyPatchesToState :: (RgaRep, PatchSet) -> Maybe (RgaRep, PatchSet)
reapplyPatchesToState (RgaRep rstate, ps@PatchSet{..}) = case rstate of
    Just VertexList{listHead = targetHead, listItems = targetItems} ->
        asum
            [ do
                targetItems' <- applyPatch parent patch targetItems
                pure
                    ( RgaRep . Just $ VertexList targetHead targetItems'
                    , ps{psPatches = Map.delete parent psPatches}
                    )
            | (parent, patch) <- Map.assocs psPatches
            ]
    Nothing -> do
        -- rstate is empty => only virtual 0 node exists
        -- => we can apply only 0 patch
        patch <- psPatches !? Zero
        pure (RgaRep $ Just patch, ps{psPatches = Map.delete Zero psPatches})

reapplyPatchesToOtherPatches :: PatchSet -> Maybe PatchSet
reapplyPatchesToOtherPatches ps@PatchSet{..} =
    asum
        [ do
            targetItems' <- applyPatch parent patch targetItems
            pure
                ps
                    { psPatches =
                        Map.insert
                            targetParent
                            (VertexList targetHead targetItems')
                            . Map.delete parent
                            $ psPatches
                    }
        | (parent, patch) <- Map.assocs psPatches
        , (targetParent, targetPatch) <- Map.assocs psPatches
        , parent /= targetParent
        , let VertexList targetHead targetItems = targetPatch
        ]

applyPatch ::
    UUID ->
    VertexList ->
    HashMap UUID VertexListItem ->
    Maybe (HashMap UUID VertexListItem)
applyPatch parent patch targetItems = case parent of
    Zero ->
        error "chunk with zero ref must be considered a state, not a patch"
    _ -> do
        item@VertexListItem{itemNext} <- HashMap.lookup parent targetItems
        let VertexList next' newItems = case itemNext of
                Nothing -> patch
                Just next -> VertexList next targetItems <> patch
        let item' = item{itemNext = Just next'}
        pure $ HashMap.insert parent item' targetItems <> newItems

reapplyRemovalsToState :: (RgaRep, PatchSet) -> Maybe (RgaRep, PatchSet)
reapplyRemovalsToState (RgaRep rstate, ps@PatchSet{..}) = do
    VertexList{listHead = targetHead, listItems = targetItems} <- rstate
    asum
        [ do
            targetItems' <- applyRemoval parent tombstone targetItems
            pure
                ( RgaRep . Just $ VertexList targetHead targetItems'
                , ps{psRemovals = Map.delete parent psRemovals}
                )
        | (parent, tombstone) <- Map.assocs psRemovals
        ]

reapplyRemovalsToPatches :: PatchSet -> Maybe PatchSet
reapplyRemovalsToPatches PatchSet{..} =
    asum
        [ do
            targetItems' <- applyRemoval parent tombstone targetItems
            pure
                PatchSet
                    { psRemovals = Map.delete parent psRemovals
                    , psPatches =
                        Map.insert
                            targetParent
                            (VertexList targetHead targetItems')
                            psPatches
                    }
        | (parent, tombstone) <- Map.assocs psRemovals
        , (targetParent, targetPatch) <- Map.assocs psPatches
        , let VertexList targetHead targetItems = targetPatch
        ]

applyRemoval ::
    UUID ->
    UUID ->
    HashMap UUID VertexListItem ->
    Maybe (HashMap UUID VertexListItem)
applyRemoval parent tombstone targetItems = do
    item@VertexListItem{itemValue = v@Op{refId}} <-
        HashMap.lookup parent targetItems
    let item' = item{itemValue = v{refId = max refId tombstone}}
    pure $ HashMap.insert parent item' targetItems

merge :: VertexList -> VertexList -> VertexList
merge v1 v2 =
    fromMaybe (error "merge of non-empty lists cannot be empty") $
        vertexListFromOps $
            (merge' `on` vertexListToOps) v1 v2

merge' :: [Vertex] -> [Vertex] -> [Vertex]
merge' [] vs2 = vs2
merge' vs1 [] = vs1
merge' w1@(v1 : vs1) w2@(v2 : vs2) =
    case compare e1 e2 of
        LT -> v2 : merge' w1 vs2
        GT -> v1 : merge' vs1 w2
        EQ -> mergeVertices : merge' vs1 vs2
  where
    Op{opId = e1, refId = tombstone1, payload = p1} = v1
    Op{opId = e2, refId = tombstone2, payload = p2} = v2
    -- priority of deletion
    mergeVertices =
        Op
            { opId = e1
            , refId = max tombstone1 tombstone2
            , payload = maxOn length p1 p2
            }

-- | Name-UUID to use as RGA type marker.
rgaType :: UUID
rgaType = $(UUID.liftName "rga")

-- | Typed RGA
newtype RGA a = RGA [a]
    deriving (Eq, Show)

instance (Replicated a) => Replicated (RGA a) where encoding = objectEncoding

instance (Replicated a) => ReplicatedAsObject (RGA a) where
    type Rep (RGA a) = RgaRep

    newObject (RGA items) = do
        oid <- getEventUuid
        vertexIds <- getEventUuids $ ls60 $ genericLength items
        ops <-
            for (zip items vertexIds) $ \(item, vertexId) -> do
                payload <- newRon item
                pure $ Op vertexId Zero payload
        modify' $ Map.insert oid $ wireStateChunk ops
        pure $ ObjectRef oid

    readObject = do
        StateChunk stateBody <- getObjectStateChunk
        mItems <-
            for stateBody $ \Op{refId, payload} -> case refId of
                Zero -> Just <$> fromRon payload
                _ -> pure Nothing
        pure . RGA $ catMaybes mItems

{- | Replace content of the RGA throug introducing changes detected by
'getGroupedDiffBy'.
-}
edit ::
    ( MonadE m
    , MonadObjectState (RGA a) m
    , ReplicaClock m
    , ReplicatedAsPayload a
    ) =>
    [a] ->
    m ()
edit newItems =
    modifyObjectStateChunk_ $ \chunk@(StateChunk stateBody) -> do
        let newItems' = [Op Zero Zero $ toPayload item | item <- newItems]
        -- TODO(2019-04-17, #59, cblp) replace 'toPayload' with 'newRon' and
        -- relax constraint on 'a' from 'ReplicatedAsPayload' to
        -- 'Replicated'
        let diff = getGroupedDiffBy eqAliveOnPayload stateBody newItems'
        (stateBody', Last lastEvent) <-
            runWriterT $
                fold <$> for diff \case
                    First removed -> do
                        tombstones <- getEventUuids $ genericLength removed
                        for (zip removed tombstones) \case
                            (op@Op{refId = Zero}, tombstone) -> do
                                -- not deleted yet
                                tell . Last $ Just tombstone
                                pure op{refId = tombstone}
                            (op, _) ->
                                -- deleted already
                                pure op
                    Both v _ -> pure v
                    Second added -> do
                        opIds <- getEventUuids $ genericLength added
                        for (zip added opIds) \(op, opId) -> do
                            tell . Last $ Just opId
                            pure op{opId}
        pure $ case lastEvent of
            Nothing -> chunk
            Just _ -> StateChunk stateBody'
  where
    eqAliveOnPayload
        Op{refId = Zero, payload = p1}
        Op{refId = Zero, payload = p2} =
            p1 == p2
    eqAliveOnPayload _ _ = False

-- | Speciaization of 'edit' for 'Text'
editText ::
    (MonadE m, MonadObjectState RgaString m, ReplicaClock m) => Text -> m ()
editText = edit . Text.unpack

{- | Speciaization of 'RGA' to 'Char'.
This is the recommended way to store a string.
-}
type RgaString = RGA Char

-- | Create an RGA from a list
newFromList ::
    (MonadState StateFrame m, ReplicaClock m, Replicated a) =>
    [a] ->
    m (ObjectRef (RGA a))
newFromList = newObject . RGA

-- | Create an 'RgaString' from a text
newFromText ::
    (MonadState StateFrame m, ReplicaClock m) =>
    Text ->
    m (ObjectRef RgaString)
newFromText = newFromList . Text.unpack

getAliveIndices :: (MonadE m, MonadObjectState (RGA a) m) => m [UUID]
getAliveIndices = do
    StateChunk stateBody <- getObjectStateChunk
    pure [opId | Op{opId, refId = Zero} <- stateBody]

-- | Read elements from RGA
getList :: (MonadE m, MonadObjectState (RGA a) m, Replicated a) => m [a]
getList = coerce <$> readObject

-- | Read characters from 'RgaString'
getText :: (MonadE m, MonadObjectState RgaString m) => m Text
getText = Text.pack <$> getList

{- | Insert a sequence of elements after the specified position.
Position is identified by 'UUID'. 'Nothing' means the beginning.
-}
insert ::
    (MonadE m, MonadObjectState (RGA a) m, ReplicaClock m, Replicated a) =>
    [a] ->
    -- | position
    Maybe UUID ->
    m ()
insert [] _ = pure ()
insert items mPosition =
    modifyObjectStateChunk_ $ \(StateChunk stateBody) -> do
        vertexIds <- getEventUuids $ ls60 $ genericLength items
        ops <-
            for (zip items vertexIds) $ \(item, vertexId) -> do
                payload <- newRon item
                pure $ Op vertexId Zero payload
        stateBody' <-
            case mPosition of
                Nothing -> pure $ ops <> stateBody
                Just position -> findAndInsertAfter position ops stateBody
        pure $ StateChunk stateBody'
  where
    findAndInsertAfter pos newOps = go
      where
        go = \case
            [] -> throwErrorText "Position not found"
            op@Op{opId} : ops
                | opId == pos -> pure $ op : newOps ++ ops
                | otherwise -> (op :) <$> go ops

insertAtBegin ::
    (MonadE m, MonadObjectState (RGA a) m, ReplicaClock m, Replicated a) =>
    [a] ->
    m ()
insertAtBegin items = insert items Nothing

insertAfter ::
    (MonadE m, MonadObjectState (RGA a) m, ReplicaClock m, Replicated a) =>
    [a] ->
    -- | position
    UUID ->
    m ()
insertAfter items = insert items . Just

{- | Insert a text after the specified position.
Position is identified by 'UUID'. 'Nothing' means the beginning.
-}
insertText ::
    (MonadE m, MonadObjectState RgaString m, ReplicaClock m) =>
    Text ->
    -- | position
    Maybe UUID ->
    m ()
insertText = insert . Text.unpack

insertTextAtBegin ::
    (MonadE m, MonadObjectState RgaString m, ReplicaClock m) => Text -> m ()
insertTextAtBegin text = insertText text Nothing

insertTextAfter ::
    (MonadE m, MonadObjectState RgaString m, ReplicaClock m) =>
    Text ->
    -- | position
    UUID ->
    m ()
insertTextAfter text = insertText text . Just

-- | Record a removal of a specific item
remove ::
    (MonadE m, MonadObjectState (RGA a) m, ReplicaClock m) =>
    -- | position
    UUID ->
    m ()
remove position =
    errorContext "RGA.remove" $
        errorContext ("position = " <> show position) $
            modifyObjectStateChunk_ $
                \(StateChunk stateBody) -> do
                    event <- getEventUuid
                    stateBody' <- findAndTombstone event stateBody
                    pure $ StateChunk stateBody'
  where
    findAndTombstone event = go
      where
        go = \case
            [] -> throwErrorText "Position not found"
            op@Op{opId} : ops
                | opId == position -> pure $ op{refId = event} : ops
                | otherwise -> (op :) <$> go ops

wireStateChunk :: [Op] -> WireStateChunk
wireStateChunk stateBody = WireStateChunk{stateType = rgaType, stateBody}

toList :: RGA a -> [a]
toList (RGA xs) = xs

toText :: RgaString -> Text
toText (RGA s) = Text.pack s