{-# LANGUAGE TupleSections #-}

-- | Combining and composing trees of trees.
module Dodge.Tree.Compose (
    tToBTree,
    overwriteLabel,
    composeTree,
    decomposeTree,
    numMetaTree,
    numSelfTree,
    combineTree,
    decomposeSelfTree,
    shiftChildren,
    attachTree,
    toOnward,
    attachOnward,
    attachOnward',
    showIntsString,
    module Dodge.Tree.Compose.Data,
) where

import Control.Monad.State
import Data.List (intersperse)
import Data.Maybe
import Data.Tuple
import Dodge.Data.GenWorld
import Dodge.Tree.Compose.Data
import LensHelp
import TreeHelp

overwriteLabel :: Int -> Tree Room -> [Tree Room] -> Tree Room
overwriteLabel i t ts =
    safeUpdateSingleNode
        (\rm -> LabelCluster i `elem` (rm ^?! rmClusterStatus . csLinks))
        ((branches .~ ts) . (root . rmClusterStatus . csLinks . at (LabelCluster i) .~ Nothing))
        t

attachTree :: (a -> Maybe a) -> Tree a -> Tree ([Tree a], a) -> Tree ([Tree a], a)
attachTree upf t = safeUpdateSingleNode (isJust . upf . snd) ((root . _2 %~ g) . (root . _1 .:~ t))
  where
    g x = fromMaybe x (upf x)

shiftChildren :: Tree ([Tree a], a) -> Tree a
shiftChildren (Node (ts, x) ts') = Node x (ts ++ map shiftChildren ts')

composeTree :: MetaTree a b -> Tree a
composeTree mt = attachList' (composeNode $ _mtTree mt) (_mtBranches mt)

mtTopLabels :: MetaTree a b -> Tree b
mtTopLabels (MTree lab _ bs) = Node lab (map (mtTopLabels . _mbTree) bs)

mtUnderLabels :: MetaTree a b -> [Tree b]
mtUnderLabels mt =
    maybe [] decomposeTree (mt ^? mtTree . nodeMetaTree)
        ++ concatMap (mtUnderLabels . _mbTree) (_mtBranches mt)

decomposeTree :: MetaTree a b -> [Tree b]
decomposeTree mt = mtTopLabels mt : mtUnderLabels mt

decomposeSelfTree :: SelfTree a -> [Tree a]
decomposeSelfTree st = fmap fst (_unST st) : concatMap (f . snd) (_unST st)
  where
    f (Left t) = [t]
    f (Right st') = decomposeSelfTree st'

combineTree :: (a -> b) -> MetaTree a b -> SelfTree b
combineTree f mt = case _mtTree mt of
    NodeTree t -> ST (Node (_mtLabel mt, Left $ fmap f t) (_unST . combineTree f . _mbTree <$> _mtBranches mt))
    --NodeTree t -> ST (Node (Left $ fmap f t) _)
    NodeMTree mt' ->
        ST
            ( Node (_mtLabel mt, Right $ combineTree f mt') $
                map (_unST . combineTree f . _mbTree) $ _mtBranches mt
            )

composeNode :: MetaNode a b -> Tree a
composeNode (NodeTree t) = t
composeNode (NodeMTree mt) = composeTree mt

attachList' :: Tree a -> [MetaBranch a b] -> Tree a
attachList' t xs = shiftChildren $ foldr attachBranch (fmap ([],) t) xs

attachBranch :: MetaBranch a b -> Tree ([Tree a], a) -> Tree ([Tree a], a)
attachBranch mb =
    safeUpdateSingleNode
        (isJust . upf . snd)
        ((root . _2 %~ g) . (root . _1 .:~ composeTree (_mbTree mb)))
  where
    upf = _mbAttach mb
    g x = fromMaybe x (upf x)

tToBTree :: String -> Tree Room -> MetaTree Room String
tToBTree str t = MTree str (NodeTree t) []

attachOnward :: b -> MetaTree Room b -> MetaTree Room b -> MetaTree Room b
attachOnward s t1 t2 = MTree s (NodeMTree t1) [MBranch toOnward t2]

attachOnward' :: MetaTree Room b -> MetaTree Room b -> MetaTree Room b
attachOnward' t1 t2 = MTree (_mtLabel t1) (NodeMTree t1) [MBranch toOnward t2]

toOnward :: Room -> Maybe Room
toOnward rm
    | OnwardCluster `elem` rm ^?! rmClusterStatus . csLinks =
        Just (rm & rmClusterStatus . csLinks . at OnwardCluster .~ Nothing)
    | otherwise = Nothing

numSelfTree :: SelfTree a -> SelfTree ([Int], a)
numSelfTree st = evalState (numSelfTree' st) [0]

numSelfTree' :: SelfTree a -> State [Int] (SelfTree ([Int], a))
numSelfTree' (ST (Node lrt bs)) = do
    is <- get
    modify (ix 0 +~ 1)
    bs' <- mapM (fmap _unST . numSelfTree' . ST) bs
    case lrt of
        (y, Left t) -> return $ ST (Node ((is, y), Left $ (_1 %~ (: is)) . swap <$> numTraversable t) bs')
        (y, Right mt') -> return $ ST (Node ((is, y), Right (evalState (numSelfTree' mt') (0 : is))) bs')

numMetaTree :: MetaTree a String -> MetaTree a ([Int], String)
numMetaTree mt = evalState (numMetaTree' mt) [0]

numMetaTree' :: MetaTree a b -> State [Int] (MetaTree a ([Int], b))
numMetaTree' (MTree lab mn bs) = do
    is <- get
    modify (ix 0 +~ 1)
    bs' <- mapM (mbTree %%~ numMetaTree') bs
    return $ MTree (is, lab) (mn & nodeMetaTree %~ (\nmt -> evalState (numMetaTree' nmt) (0 : is))) bs'

showIntsString :: ([Int], String) -> String
showIntsString (is, s) = foldr1 (++) (intersperse ":" (map show $ reverse is)) ++ ':' : s