{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

module Cursor.Tree.Types
  ( TreeCursor (..),
    treeCursorAboveL,
    treeCursorCurrentL,
    treeCursorBelowL,
    treeCursorCurrentSubTreeL,
    TreeAbove (..),
    treeAboveLeftsL,
    treeAboveAboveL,
    treeAboveNodeL,
    treeAboveRightsL,
    TreeCursorSelection (..),

    -- * CTree
    CTree (..),
    makeCTree,
    cTree,
    rebuildCTree,
    CForest (..),
    makeCForest,
    cForest,
    rebuildCForest,
    emptyCForest,
    openForest,
    closedForest,
    lengthCForest,
    unpackCForest,
  )
where

import Control.DeepSeq
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.Tree
import Data.Validity
import Data.Validity.Tree ()
import GHC.Generics (Generic)
import Lens.Micro

data TreeCursor a b = TreeCursor
  { treeAbove :: !(Maybe (TreeAbove b)),
    treeCurrent :: !a,
    treeBelow :: !(CForest b)
  }
  deriving (Show, Eq, Generic)

instance (Validity a, Validity b) => Validity (TreeCursor a b)

instance (NFData a, NFData b) => NFData (TreeCursor a b)

treeCursorAboveL :: Lens' (TreeCursor a b) (Maybe (TreeAbove b))
treeCursorAboveL = lens treeAbove $ \tc ta -> tc {treeAbove = ta}

treeCursorCurrentL :: Lens (TreeCursor a b) (TreeCursor a' b) a a'
treeCursorCurrentL = lens treeCurrent $ \tc a -> tc {treeCurrent = a}

treeCursorBelowL :: Lens' (TreeCursor a b) (CForest b)
treeCursorBelowL = lens treeBelow $ \tc tb -> tc {treeBelow = tb}

treeCursorCurrentSubTreeL :: Lens (TreeCursor a b) (TreeCursor a' b) (a, CForest b) (a', CForest b)
treeCursorCurrentSubTreeL =
  lens (\tc -> (treeCurrent tc, treeBelow tc)) (\tc (a, cf) -> tc {treeCurrent = a, treeBelow = cf})

data TreeAbove b = TreeAbove
  { treeAboveLefts :: ![CTree b], -- In reverse order
    treeAboveAbove :: !(Maybe (TreeAbove b)),
    treeAboveNode :: !b,
    treeAboveRights :: ![CTree b]
  }
  deriving (Show, Eq, Generic, Functor)

instance Validity b => Validity (TreeAbove b)

instance NFData b => NFData (TreeAbove b)

treeAboveLeftsL :: Lens' (TreeAbove b) [CTree b]
treeAboveLeftsL = lens treeAboveLefts $ \ta tal -> ta {treeAboveLefts = tal}

treeAboveAboveL :: Lens' (TreeAbove b) (Maybe (TreeAbove b))
treeAboveAboveL = lens treeAboveAbove $ \ta taa -> ta {treeAboveAbove = taa}

treeAboveNodeL :: Lens' (TreeAbove b) b
treeAboveNodeL = lens treeAboveNode $ \ta a -> ta {treeAboveNode = a}

treeAboveRightsL :: Lens' (TreeAbove b) [CTree b]
treeAboveRightsL = lens treeAboveRights $ \ta tar -> ta {treeAboveRights = tar}

data TreeCursorSelection
  = SelectNode
  | SelectChild !Int !TreeCursorSelection
  deriving (Show, Eq, Generic)

instance Validity TreeCursorSelection

instance NFData TreeCursorSelection

data CTree a
  = CNode !a (CForest a)
  deriving (Show, Eq, Generic, Functor)

instance Validity a => Validity (CTree a)

instance NFData a => NFData (CTree a)

instance Foldable CTree where
  foldMap f (CNode a cf) = f a `mappend` foldMap f cf

instance Traversable CTree where
  traverse f (CNode a cf) = CNode <$> f a <*> traverse f cf

makeCTree :: Tree a -> CTree a
makeCTree = cTree False

cTree :: Bool -> Tree a -> CTree a
cTree b (Node v f) = CNode v $ cForest b f

rebuildCTree :: CTree a -> Tree a
rebuildCTree (CNode v cf) = Node v $ rebuildCForest cf

data CForest a
  = EmptyCForest
  | ClosedForest !(NonEmpty (Tree a))
  | OpenForest !(NonEmpty (CTree a))
  deriving (Show, Eq, Generic, Functor)

instance Validity a => Validity (CForest a)

instance NFData a => NFData (CForest a)

instance Foldable CForest where
  foldMap f = \case
    EmptyCForest -> mempty
    ClosedForest ne -> foldMap (foldMap f) ne
    OpenForest ne -> foldMap (foldMap f) ne

instance Traversable CForest where
  traverse f = \case
    EmptyCForest -> pure EmptyCForest
    ClosedForest ne -> ClosedForest <$> traverse (traverse f) ne
    OpenForest ne -> OpenForest <$> traverse (traverse f) ne

makeCForest :: Forest a -> CForest a
makeCForest = cForest True

cForest :: Bool -> Forest a -> CForest a
cForest b f =
  if b
    then openForest $ map (cTree b) f
    else closedForest f

rebuildCForest :: CForest a -> Forest a
rebuildCForest EmptyCForest = []
rebuildCForest (ClosedForest f) = NE.toList f
rebuildCForest (OpenForest ct) = NE.toList $ NE.map rebuildCTree ct

emptyCForest :: CForest a
emptyCForest = EmptyCForest

openForest :: [CTree a] -> CForest a
openForest ts = maybe emptyCForest OpenForest $ NE.nonEmpty ts

closedForest :: [Tree a] -> CForest a
closedForest ts = maybe emptyCForest ClosedForest $ NE.nonEmpty ts

lengthCForest :: CForest a -> Int
lengthCForest EmptyCForest = 0
lengthCForest (ClosedForest ts) = length ts
lengthCForest (OpenForest ts) = length ts

unpackCForest :: CForest a -> [CTree a]
unpackCForest EmptyCForest = []
unpackCForest (ClosedForest ts) = NE.toList $ NE.map makeCTree ts
unpackCForest (OpenForest ts) = NE.toList ts