{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

module Cursor.Tree.Base
  ( singletonTreeCursor,
    makeTreeCursor,
    makeNodeTreeCursor,
    makeTreeCursorWithSelection,
    rebuildTreeCursor,
    mapTreeCursor,
    currentTree,
    makeTreeCursorWithAbove,
    traverseTreeCursor,
    foldTreeCursor,
  )
where

import Control.Monad
import Cursor.Tree.Types

singletonTreeCursor :: a -> TreeCursor a b
singletonTreeCursor v = TreeCursor {treeAbove = Nothing, treeCurrent = v, treeBelow = emptyCForest}

makeTreeCursor :: (b -> a) -> CTree b -> TreeCursor a b
makeTreeCursor g (CNode v fs) = TreeCursor {treeAbove = Nothing, treeCurrent = g v, treeBelow = fs}

makeNodeTreeCursor :: a -> CForest b -> TreeCursor a b
makeNodeTreeCursor v fs = TreeCursor {treeAbove = Nothing, treeCurrent = v, treeBelow = fs}

makeTreeCursorWithSelection ::
  (a -> b) -> (b -> a) -> TreeCursorSelection -> CTree b -> Maybe (TreeCursor a b)
makeTreeCursorWithSelection f g sel = walkDown sel . makeTreeCursor g
  where
    walkDown SelectNode tc = pure tc
    walkDown (SelectChild i s) TreeCursor {..} =
      (walkDown s =<<) $
        case splitAt i $ unpackCForest treeBelow of
          (_, []) -> Nothing
          (lefts, current : rights) ->
            Just $
              makeTreeCursorWithAbove g current $
                Just $
                  TreeAbove
                    { treeAboveLefts = reverse lefts,
                      treeAboveAbove = treeAbove,
                      treeAboveNode = f treeCurrent,
                      treeAboveRights = rights
                    }

rebuildTreeCursor :: (a -> b) -> TreeCursor a b -> CTree b
rebuildTreeCursor f TreeCursor {..} = wrapAbove treeAbove $ CNode (f treeCurrent) treeBelow
  where
    wrapAbove Nothing t = t
    wrapAbove (Just TreeAbove {..}) t =
      wrapAbove treeAboveAbove $
        CNode treeAboveNode $
          openForest $
            concat [reverse treeAboveLefts, [t], treeAboveRights]

mapTreeCursor :: (a -> c) -> (b -> d) -> TreeCursor a b -> TreeCursor c d
mapTreeCursor f g TreeCursor {..} =
  TreeCursor
    { treeAbove = fmap g <$> treeAbove,
      treeCurrent = f treeCurrent,
      treeBelow = fmap g treeBelow
    }

currentTree :: (a -> b) -> TreeCursor a b -> CTree b
currentTree f TreeCursor {..} = CNode (f treeCurrent) treeBelow

makeTreeCursorWithAbove :: (b -> a) -> CTree b -> Maybe (TreeAbove b) -> TreeCursor a b
makeTreeCursorWithAbove g (CNode a forest) mta =
  TreeCursor {treeAbove = mta, treeCurrent = g a, treeBelow = forest}

traverseTreeCursor ::
  forall a b m c.
  Monad m =>
  ([CTree b] -> b -> [CTree b] -> c -> m c) ->
  (a -> CForest b -> m c) ->
  TreeCursor a b ->
  m c
traverseTreeCursor wrapFunc currentFunc TreeCursor {..} =
  currentFunc treeCurrent treeBelow >>= wrapAbove treeAbove
  where
    wrapAbove :: Maybe (TreeAbove b) -> c -> m c
    wrapAbove Nothing = pure
    wrapAbove (Just ta) = goAbove ta
    goAbove :: TreeAbove b -> c -> m c
    goAbove TreeAbove {..} =
      wrapFunc (reverse treeAboveLefts) treeAboveNode treeAboveRights >=> wrapAbove treeAboveAbove

foldTreeCursor ::
  forall a b c.
  ([CTree b] -> b -> [CTree b] -> c -> c) ->
  (a -> CForest b -> c) ->
  TreeCursor a b ->
  c
foldTreeCursor wrapFunc currentFunc TreeCursor {..} =
  wrapAbove treeAbove $ currentFunc treeCurrent treeBelow
  where
    wrapAbove :: Maybe (TreeAbove b) -> c -> c
    wrapAbove Nothing = id
    wrapAbove (Just ta) = goAbove ta
    goAbove :: TreeAbove b -> c -> c
    goAbove TreeAbove {..} =
      wrapAbove treeAboveAbove . wrapFunc (reverse treeAboveLefts) treeAboveNode treeAboveRights