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

module Cursor.Tree.Insert
  ( treeCursorInsert
  , treeCursorInsertAndSelect
  , treeCursorAppend
  , treeCursorAppendAndSelect
  , treeCursorAddChildAtPos
  , treeCursorAddChildAtStart
  , treeCursorAddChildAtEnd
  ) where

import qualified Data.List.NonEmpty as NE
import Data.List.NonEmpty ((<|))
import Data.Tree

import Cursor.Tree.Base
import Cursor.Tree.Types

treeCursorInsert :: Tree b -> TreeCursor a b -> Maybe (TreeCursor a b)
treeCursorInsert tree tc@TreeCursor {..} = do
  ta <- treeAbove
  let newTreeAbove = ta {treeAboveLefts = makeCTree tree : treeAboveLefts ta}
  pure tc {treeAbove = Just newTreeAbove}

treeCursorInsertAndSelect ::
     (a -> b) -> (b -> a) -> Tree b -> TreeCursor a b -> Maybe (TreeCursor a b)
treeCursorInsertAndSelect f g tree tc@TreeCursor {..} = do
  ta <- treeAbove
  let newTreeAbove = ta {treeAboveRights = currentTree f tc : treeAboveRights ta}
  pure $ makeTreeCursorWithAbove g (makeCTree tree) $ Just newTreeAbove

treeCursorAppend :: Tree b -> TreeCursor a b -> Maybe (TreeCursor a b)
treeCursorAppend tree tc@TreeCursor {..} = do
  ta <- treeAbove
  let newTreeAbove = ta {treeAboveRights = makeCTree tree : treeAboveRights ta}
  pure tc {treeAbove = Just newTreeAbove}

treeCursorAppendAndSelect ::
     (a -> b) -> (b -> a) -> Tree b -> TreeCursor a b -> Maybe (TreeCursor a b)
treeCursorAppendAndSelect f g tree tc@TreeCursor {..} = do
  ta <- treeAbove
  let newTreeAbove = ta {treeAboveLefts = currentTree f tc : treeAboveLefts ta}
  pure $ makeTreeCursorWithAbove g (makeCTree tree) $ Just newTreeAbove

-- TODO make this fail if the position doesn't make sense
treeCursorAddChildAtPos :: Int -> Tree b -> TreeCursor a b -> TreeCursor a b
treeCursorAddChildAtPos i t tc =
  case treeBelow tc of
    EmptyCForest -> tc {treeBelow = openForest [makeCTree t]}
    ClosedForest ts ->
      let (before, after) = splitAt i $ NE.toList ts
       in tc {treeBelow = openForest $ map makeCTree $ before ++ [t] ++ after}
    OpenForest ts ->
      let (before, after) = splitAt i $ NE.toList ts
       in tc {treeBelow = openForest $ before ++ [makeCTree t] ++ after}

treeCursorAddChildAtStart :: Tree b -> TreeCursor a b -> TreeCursor a b
treeCursorAddChildAtStart t tc =
  case treeBelow tc of
    EmptyCForest -> tc {treeBelow = openForest [makeCTree t]}
    ClosedForest ts -> tc {treeBelow = OpenForest $ NE.map makeCTree $ t <| ts}
    OpenForest ts -> tc {treeBelow = OpenForest $ makeCTree t <| ts}

treeCursorAddChildAtEnd :: Tree b -> TreeCursor a b -> TreeCursor a b
treeCursorAddChildAtEnd t tc =
  case treeBelow tc of
    EmptyCForest -> tc {treeBelow = openForest [makeCTree t]}
    ClosedForest ts -> tc {treeBelow = openForest $ map makeCTree $ NE.toList ts ++ [t]}
    OpenForest ts -> tc {treeBelow = openForest $ NE.toList ts ++ [makeCTree t]}