{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

{-|
Description: An advanced 'Patch' on 'Map'

Patches of this type can can insert, delete, and move values from one key to
another, and move patches may also additionally patch the value being moved.
-}
module Data.Patch.MapWithPatchingMove
  ( PatchMapWithPatchingMove (..)
  , patchMapWithPatchingMove
  , patchMapWithPatchingMoveInsertAll
  , insertMapKey
  , moveMapKey
  , patchMapKey
  , swapMapKey
  , deleteMapKey
  , unsafePatchMapWithPatchingMove
  , patchMapWithPatchingMoveNewElements
  , patchMapWithPatchingMoveNewElementsMap
  , patchThatSortsMapWith
  , patchThatChangesAndSortsMapWith
  , patchThatChangesMap

  -- * Node Info
  , NodeInfo (..)
  , bitraverseNodeInfo
  , nodeInfoMapFrom
  , nodeInfoMapMFrom
  , nodeInfoSetTo

  -- * From
  , From(..)
  , bitraverseFrom

  -- * To
  , To

  -- TODO internals module
  , Fixup (..)
  ) where

import Data.Patch.Class

import Control.Lens ((<&>))
import Control.Lens.TH (makeWrapped)
import Data.Align (align)
import Data.Foldable (toList)
import Data.Function
import Data.List
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe
#if !MIN_VERSION_base(4,11,0)
import Data.Semigroup (Semigroup (..))
#endif
import Data.Monoid.DecidablyEmpty
import Data.Set (Set)
import qualified Data.Set as Set
import Data.These (These (..))

-- | Patch a Map with additions, deletions, and moves.  Invariant: If key @k1@
-- is coming from @From_Move k2@, then key @k2@ should be going to @Just k1@,
-- and vice versa.  There should never be any unpaired From/To keys.
newtype PatchMapWithPatchingMove k p = PatchMapWithPatchingMove
  { -- | Extract the internal representation of the 'PatchMapWithPatchingMove'
    unPatchMapWithPatchingMove :: Map k (NodeInfo k p)
  }

deriving instance (Show k, Show p, Show (PatchTarget p)) => Show (PatchMapWithPatchingMove k p)
deriving instance (Ord k, Read k, Read p, Read (PatchTarget p)) => Read (PatchMapWithPatchingMove k p)
deriving instance (Eq k, Eq p, Eq (PatchTarget p)) => Eq (PatchMapWithPatchingMove k p)
deriving instance (Ord k, Ord p, Ord (PatchTarget p)) => Ord (PatchMapWithPatchingMove k p)

deriving instance ( Ord k
#if !MIN_VERSION_base(4,11,0)
                  , Semigroup p
#endif
                  , DecidablyEmpty p
                  , Patch p
                  ) => DecidablyEmpty (PatchMapWithPatchingMove k p)

-- | Create a 'PatchMapWithPatchingMove', validating it
patchMapWithPatchingMove
  :: Ord k => Map k (NodeInfo k p) -> Maybe (PatchMapWithPatchingMove k p)
patchMapWithPatchingMove m = if valid then Just $ PatchMapWithPatchingMove m else Nothing
  where valid = forwardLinks == backwardLinks
        forwardLinks = Map.mapMaybe _nodeInfo_to m
        backwardLinks = Map.fromList $ catMaybes $ flip fmap (Map.toList m) $ \(to, p) ->
          case _nodeInfo_from p of
            From_Move from _ -> Just (from, to)
            _ -> Nothing

-- | Create a 'PatchMapWithPatchingMove' that inserts everything in the given 'Map'
patchMapWithPatchingMoveInsertAll
  :: Map k (PatchTarget p) -> PatchMapWithPatchingMove k p
patchMapWithPatchingMoveInsertAll m = PatchMapWithPatchingMove $ flip fmap m $ \v -> NodeInfo
  { _nodeInfo_from = From_Insert v
  , _nodeInfo_to = Nothing
  }

-- | Make a @'PatchMapWithPatchingMove' k p@ which has the effect of inserting or replacing a value @v@ at the given key @k@, like 'Map.insert'.
insertMapKey
  :: k -> PatchTarget p -> PatchMapWithPatchingMove k p
insertMapKey k v = PatchMapWithPatchingMove . Map.singleton k $ NodeInfo (From_Insert v) Nothing

-- |Make a @'PatchMapWithPatchingMove' k p@ which has the effect of moving the value from the first key @k@ to the second key @k@, equivalent to:
--
-- @
--     'Map.delete' src (maybe map ('Map.insert' dst) (Map.lookup src map))
-- @
moveMapKey
  :: ( DecidablyEmpty p
#if !MIN_VERSION_base(4,11,0)
     , Semigroup p
#endif
     , Patch p
     )
  => Ord k => k -> k -> PatchMapWithPatchingMove k p
moveMapKey src dst
  | src == dst = mempty
  | otherwise =
      PatchMapWithPatchingMove $ Map.fromList
        [ (dst, NodeInfo (From_Move src mempty) Nothing)
        , (src, NodeInfo From_Delete (Just dst))
        ]

patchMapKey
  :: ( DecidablyEmpty p
#if !MIN_VERSION_base(4,11,0)
     , Semigroup p
#endif
     )
  => k -> p -> PatchMapWithPatchingMove k p
patchMapKey k p
  | isEmpty p = PatchMapWithPatchingMove Map.empty
  | otherwise =
      PatchMapWithPatchingMove $ Map.singleton k $ NodeInfo (From_Move k p) (Just k)

-- |Make a @'PatchMapWithPatchingMove' k p@ which has the effect of swapping two keys in the mapping, equivalent to:
--
-- @
--     let aMay = Map.lookup a map
--         bMay = Map.lookup b map
--     in maybe id (Map.insert a) (bMay <> aMay)
--      . maybe id (Map.insert b) (aMay <> bMay)
--      . Map.delete a . Map.delete b $ map
-- @
swapMapKey
  :: ( DecidablyEmpty p
#if !MIN_VERSION_base(4,11,0)
     , Semigroup p
#endif
     , Patch p
     )
  => Ord k => k -> k -> PatchMapWithPatchingMove k p
swapMapKey src dst
  | src == dst = mempty
  | otherwise =
    PatchMapWithPatchingMove $ Map.fromList
      [ (dst, NodeInfo (From_Move src mempty) (Just src))
      , (src, NodeInfo (From_Move dst mempty) (Just dst))
      ]

-- | Make a @'PatchMapWithPatchingMove' k v@ which has the effect of deleting a key in
-- the mapping, equivalent to 'Map.delete'.
deleteMapKey
  :: k -> PatchMapWithPatchingMove k v
deleteMapKey k = PatchMapWithPatchingMove . Map.singleton k $ NodeInfo From_Delete Nothing

-- | Wrap a @'Map' k (NodeInfo k v)@ representing patch changes into a @'PatchMapWithPatchingMove' k v@, without checking any invariants.
--
-- __Warning:__ when using this function, you must ensure that the invariants of 'PatchMapWithPatchingMove' are preserved; they will not be checked.
unsafePatchMapWithPatchingMove
  :: Map k (NodeInfo k p) -> PatchMapWithPatchingMove k p
unsafePatchMapWithPatchingMove = PatchMapWithPatchingMove

-- | Apply the insertions, deletions, and moves to a given 'Map'
instance (Ord k, Patch p) => Patch (PatchMapWithPatchingMove k p) where
  type PatchTarget (PatchMapWithPatchingMove k p) = Map k (PatchTarget p)
  -- TODO: return Nothing sometimes
  -- Note: the strict application here is critical to ensuring that incremental
  -- merges don't hold onto all their prerequisite events forever; can we make
  -- this more robust?
  apply (PatchMapWithPatchingMove m) old = Just $! insertions `Map.union` (old `Map.difference` deletions)
    where insertions = flip Map.mapMaybeWithKey m $ \_ ni -> case _nodeInfo_from ni of
            From_Insert v -> Just v
            From_Move k p -> applyAlways p <$> Map.lookup k old
            From_Delete -> Nothing
          deletions = flip Map.mapMaybeWithKey m $ \_ ni -> case _nodeInfo_from ni of
            From_Delete -> Just ()
            _ -> Nothing

-- | Returns all the new elements that will be added to the 'Map'
patchMapWithPatchingMoveNewElements
  :: PatchMapWithPatchingMove k p -> [PatchTarget p]
patchMapWithPatchingMoveNewElements = Map.elems . patchMapWithPatchingMoveNewElementsMap

-- | Return a @'Map' k v@ with all the inserts/updates from the given @'PatchMapWithPatchingMove' k v@.
patchMapWithPatchingMoveNewElementsMap
  :: PatchMapWithPatchingMove k p -> Map k (PatchTarget p)
patchMapWithPatchingMoveNewElementsMap (PatchMapWithPatchingMove p) = Map.mapMaybe f p
  where f ni = case _nodeInfo_from ni of
          From_Insert v -> Just v
          From_Move _ _ -> Nothing
          From_Delete -> Nothing

-- | Create a 'PatchMapWithPatchingMove' that, if applied to the given 'Map', will sort
-- its values using the given ordering function.  The set keys of the 'Map' is
-- not changed.
patchThatSortsMapWith
  :: (Ord k, Monoid p)
  => (PatchTarget p -> PatchTarget p -> Ordering)
  -> Map k (PatchTarget p) -> PatchMapWithPatchingMove k p
patchThatSortsMapWith cmp m = PatchMapWithPatchingMove $ Map.fromList $ catMaybes $ zipWith g unsorted sorted
  where unsorted = Map.toList m
        sorted = sortBy (cmp `on` snd) unsorted
        f (to, _) (from, _) = if to == from then Nothing else
          Just (from, to)
        reverseMapping = Map.fromList $ catMaybes $ zipWith f unsorted sorted
        g (to, _) (from, _) = if to == from then Nothing else
          let Just movingTo = Map.lookup from reverseMapping
          in Just (to, NodeInfo (From_Move from mempty) $ Just movingTo)

-- | Create a 'PatchMapWithPatchingMove' that, if applied to the first 'Map' provided,
-- will produce a 'Map' with the same values as the second 'Map' but with the
-- values sorted with the given ordering function.
patchThatChangesAndSortsMapWith
  :: forall k p. (Ord k, Ord (PatchTarget p), Monoid p)
  => (PatchTarget p -> PatchTarget p -> Ordering)
  -> Map k (PatchTarget p) -> Map k (PatchTarget p) -> PatchMapWithPatchingMove k p
patchThatChangesAndSortsMapWith cmp oldByIndex newByIndexUnsorted = patchThatChangesMap oldByIndex newByIndex
  where newList = Map.toList newByIndexUnsorted
        newByIndex = Map.fromList $ zip (fst <$> newList) $ sortBy cmp $ snd <$> newList

-- | Create a 'PatchMapWithPatchingMove' that, if applied to the first 'Map' provided,
-- will produce the second 'Map'.
-- Note: this will never produce a patch on a value.
patchThatChangesMap
  :: forall k p
  .  (Ord k, Ord (PatchTarget p), Monoid p)
  => Map k (PatchTarget p) -> Map k (PatchTarget p) -> PatchMapWithPatchingMove k p
patchThatChangesMap oldByIndex newByIndex = patch
  where invert :: Map k (PatchTarget p) -> Map (PatchTarget p) (Set k)
        invert = Map.fromListWith (<>) . fmap (\(k, v) -> (v, Set.singleton k)) . Map.toList
        -- In the places where we use unionDistinct, a non-distinct key indicates a bug in this function
        unionDistinct :: forall k' v'. Ord k' => Map k' v' -> Map k' v' -> Map k' v'
        unionDistinct = Map.unionWith (error "patchThatChangesMap: non-distinct keys")
        unionPairDistinct :: (Map k (From k v), Map k (To k)) -> (Map k (From k v), Map k (To k)) -> (Map k (From k v), Map k (To k))
        unionPairDistinct (oldFroms, oldTos) (newFroms, newTos) = (unionDistinct oldFroms newFroms, unionDistinct oldTos newTos)
        -- Generate patch info for a single value
        -- Keys that are found in both the old and new sets will not be patched
        -- Keys that are found in only the old set will be moved to a new position if any are available; otherwise they will be deleted
        -- Keys that are found in only the new set will be populated by moving an old key if any are available; otherwise they will be inserted
        patchSingleValue :: PatchTarget p -> Set k -> Set k -> (Map k (From k p), Map k (To k))
        patchSingleValue v oldKeys newKeys = foldl' unionPairDistinct mempty $ align (toList $ oldKeys `Set.difference` newKeys) (toList $ newKeys `Set.difference` oldKeys) <&> \case
          This oldK -> (mempty, Map.singleton oldK Nothing) -- There's nowhere for this value to go, so we know we are deleting it
          That newK -> (Map.singleton newK $ From_Insert v, mempty) -- There's nowhere fo this value to come from, so we know we are inserting it
          These oldK newK -> (Map.singleton newK $ From_Move oldK mempty, Map.singleton oldK $ Just newK)
        -- Run patchSingleValue on a These.  Missing old or new sets are considered empty
        patchSingleValueThese :: PatchTarget p -> These (Set k) (Set k) -> (Map k (From k p), Map k (To k))
        patchSingleValueThese v = \case
          This oldKeys -> patchSingleValue v oldKeys mempty
          That newKeys -> patchSingleValue v mempty newKeys
          These oldKeys newKeys -> patchSingleValue v oldKeys newKeys
        -- Generate froms and tos for all values, then merge them together
        (froms, tos) = foldl' unionPairDistinct mempty $ Map.mapWithKey patchSingleValueThese $ align (invert oldByIndex) (invert newByIndex)
        patch = unsafePatchMapWithPatchingMove $ align froms tos <&> \case
          This from -> NodeInfo from Nothing -- Since we don't have a 'to' record for this key, that must mean it isn't being moved anywhere, so it should be deleted.
          That to -> NodeInfo From_Delete to -- Since we don't have a 'from' record for this key, it must be getting deleted
          These from to -> NodeInfo from to

--
-- NodeInfo
--

-- | Holds the information about each key: where its new value should come from,
-- and where its old value should go to
data NodeInfo k p = NodeInfo
  { _nodeInfo_from :: !(From k p)
    -- ^ Where do we get the new value for this key?
  , _nodeInfo_to :: !(To k)
    -- ^ If the old value is being kept (i.e. moved rather than deleted or
    -- replaced), where is it going?
  }
deriving instance (Show k, Show p, Show (PatchTarget p)) => Show (NodeInfo k p)
deriving instance (Read k, Read p, Read (PatchTarget p)) => Read (NodeInfo k p)
deriving instance (Eq k, Eq p, Eq (PatchTarget p)) => Eq (NodeInfo k p)
deriving instance (Ord k, Ord p, Ord (PatchTarget p)) => Ord (NodeInfo k p)

-- | Traverse the 'NodeInfo' over the key, patch, and patch target. Because of
-- the type families here, this doesn't it any bi- or tri-traversal class.
bitraverseNodeInfo
  :: Applicative f
  => (k0 -> f k1)
  -> (p0 -> f p1)
  -> (PatchTarget p0 -> f (PatchTarget p1))
  -> NodeInfo k0 p0 -> f (NodeInfo k1 p1)
bitraverseNodeInfo fk fp fpt (NodeInfo from to) = NodeInfo
  <$> bitraverseFrom fk fp fpt from
  <*> traverse fk to

-- | Change the 'From' value of a 'NodeInfo'
nodeInfoMapFrom
  :: (From k v -> From k v) -> NodeInfo k v -> NodeInfo k v
nodeInfoMapFrom f ni = ni { _nodeInfo_from = f $ _nodeInfo_from ni }

-- | Change the 'From' value of a 'NodeInfo', using a 'Functor' (or
-- 'Applicative', 'Monad', etc.) action to get the new value
nodeInfoMapMFrom
  :: Functor f => (From k v -> f (From k v)) -> NodeInfo k v -> f (NodeInfo k v)
nodeInfoMapMFrom f ni = fmap (\result -> ni { _nodeInfo_from = result }) $ f $ _nodeInfo_from ni

-- | Set the 'To' field of a 'NodeInfo'
nodeInfoSetTo
  :: To k -> NodeInfo k v -> NodeInfo k v
nodeInfoSetTo to ni = ni { _nodeInfo_to = to }

--
-- From
--

-- | Describe how a key's new value should be produced
data From k p
   = From_Insert (PatchTarget p) -- ^ Insert the given value here
   | From_Delete -- ^ Delete the existing value, if any, from here
   | From_Move !k !p -- ^ Move the value here from the given key, and apply the given patch

deriving instance (Show k, Show p, Show (PatchTarget p)) => Show (From k p)
deriving instance (Read k, Read p, Read (PatchTarget p)) => Read (From k p)
deriving instance (Eq k, Eq p, Eq (PatchTarget p)) => Eq (From k p)
deriving instance (Ord k, Ord p, Ord (PatchTarget p)) => Ord (From k p)

-- | Traverse the 'From' over the key, patch, and patch target. Because of
-- the type families here, this doesn't it any bi- or tri-traversal class.
bitraverseFrom
  :: Applicative f
  => (k0 -> f k1)
  -> (p0 -> f p1)
  -> (PatchTarget p0 -> f (PatchTarget p1))
  -> From k0 p0 -> f (From k1 p1)
bitraverseFrom fk fp fpt = \case
  From_Insert pt -> From_Insert <$> fpt pt
  From_Delete -> pure From_Delete
  From_Move k p -> From_Move <$> fk k <*> fp p

--
-- To
--

-- | Describe where a key's old value will go.  If this is 'Just', that means
-- the key's old value will be moved to the given other key; if it is 'Nothing',
-- that means it will be deleted.
type To = Maybe

--
-- Fixup
--

-- | Helper data structure used for composing patches using the monoid instance.
data Fixup k v
   = Fixup_Delete
   | Fixup_Update (These (From k v) (To k))

-- | Compose patches having the same effect as applying the patches in turn:
-- @'applyAlways' (p <> q) == 'applyAlways' p . 'applyAlways' q@
instance ( Ord k
#if !MIN_VERSION_base(4,11,0)
         , Semigroup p
#endif
         , DecidablyEmpty p
         , Patch p
         ) => Semigroup (PatchMapWithPatchingMove k p) where
  PatchMapWithPatchingMove mNew <> PatchMapWithPatchingMove mOld = PatchMapWithPatchingMove m
    where
      connections = Map.elems $ Map.intersectionWithKey (\_ new old -> (_nodeInfo_to new, _nodeInfo_from old)) mNew mOld
      h :: (Maybe k, From k p) -> [(k, Fixup k p)]
      h = \case
        (Just toAfter, From_Move fromBefore p)
          | fromBefore == toAfter && isEmpty p
            -> [ (toAfter, Fixup_Delete)
               ]
          | otherwise
            -> [ (toAfter, Fixup_Update (This (From_Move fromBefore p)))
               , (fromBefore, Fixup_Update (That (Just toAfter)))
               ]
        (Nothing, From_Move fromBefore _) -> [(fromBefore, Fixup_Update (That Nothing))] -- The item is destroyed in the second patch, so indicate that it is destroyed in the source map
        (Just toAfter, editBefore) -> [(toAfter, Fixup_Update (This editBefore))]
        (Nothing, _) -> []
      mergeFixups Fixup_Delete Fixup_Delete = Fixup_Delete
      mergeFixups (Fixup_Update a) (Fixup_Update b)
        | This x <- a, That y <- b
        = Fixup_Update $ These x y
        | That y <- a, This x <- b
        = Fixup_Update $ These x y
      mergeFixups _ _ = error "PatchMapWithPatchingMove: incompatible fixups"
      fixups = Map.fromListWithKey (\_ -> mergeFixups) $ concatMap h connections
      combineNodeInfos niNew niOld = NodeInfo
        { _nodeInfo_from = _nodeInfo_from niNew
        , _nodeInfo_to = _nodeInfo_to niOld
        }
      applyFixup ni = \case
        Fixup_Delete -> Nothing
        Fixup_Update u -> Just $ NodeInfo
          { _nodeInfo_from = case _nodeInfo_from ni of
              -- The new patch has a Move, so it could be affected by the
              -- corresponding From in the old patch.  If that From exists, then
              -- it is in the fixup here.
              f@(From_Move _ p') -> case getHere u of
                -- If there's no `From` fixup, just use the "new" `From`
                Nothing -> f
                -- If there's a `From` fixup which is an Insert, we can just apply
                -- our patch to that and turn ourselves into an insert.
                Just (From_Insert v) -> From_Insert $ applyAlways p' v
                -- If there's a `From` fixup which is a Delete, then we can throw
                -- our patch away because there's nothing to apply it to and
                -- become a Delete ourselves.
                Just From_Delete -> From_Delete
                -- If there's a `From` fixup which is a Move, we need to apply
                -- both the old patch and the new patch (in that order) to the
                -- value, so we append the patches here.
                Just (From_Move oldKey p) -> From_Move oldKey $ p' <> p
              -- If the new patch has an Insert, it doesn't care what the fixup
              -- value is, because it will overwrite it anyway.
              f@(From_Insert _) -> f
              -- If the new patch has an Delete, it doesn't care what the fixup
              -- value is, because it will overwrite it anyway.
              f@From_Delete -> f
          , _nodeInfo_to = case _nodeInfo_to ni of
              -- The old patch deletes this data, so we must delete it as well.
              -- According to the code above, any time we have this situation we
              -- should also have `getThere u == Nothing` because a fixup
              -- shouldn't be generated.
              Nothing -> Nothing
              -- The old patch sends the value to oldToAfter
              Just oldToAfter -> case getThere u of
                -- If there is no fixup, that should mean that the new patch
                -- doesn't do anything with the value in oldToAfter, so we still
                -- send it to oldToAfter
                Nothing -> Just oldToAfter
                -- If there is a fixup, it should tell us where the new patch
                -- sends the value at key oldToAfter.  We send our value there.
                Just mNewToAfter -> mNewToAfter
          }
      m = Map.differenceWithKey (\_ -> applyFixup) (Map.unionWith combineNodeInfos mNew mOld) fixups
      getHere :: These a b -> Maybe a
      getHere = \case
        This a -> Just a
        These a _ -> Just a
        That _ -> Nothing
      getThere :: These a b -> Maybe b
      getThere = \case
        This _ -> Nothing
        These _ b -> Just b
        That b -> Just b

--TODO: Figure out how to implement this in terms of PatchDMapWithPatchingMove rather than duplicating it here
-- | Compose patches having the same effect as applying the patches in turn:
-- @'applyAlways' (p <> q) == 'applyAlways' p . 'applyAlways' q@
instance ( Ord k
#if !MIN_VERSION_base(4,11,0)
         , Semigroup p
#endif
         , DecidablyEmpty p
         , Patch p
         ) => Monoid (PatchMapWithPatchingMove k p) where
  mempty = PatchMapWithPatchingMove mempty
  mappend = (<>)

makeWrapped ''PatchMapWithPatchingMove