{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}

module Graphics.LambdaCube.Tree where

import Data.ByteString.Char8 (ByteString)
import Data.Trie (Trie)
import qualified Data.ByteString.Char8 as B
import qualified Data.Trie as T
import qualified Data.Trie.Convenience as T

-- Contains full paths mapped to data and node names mapped to full
-- paths, where the path separator is NUL.
data IndexedTree a = IT String (Trie a) (Trie ByteString) deriving Show

separator :: ByteString
separator = B.pack "\0"

sep :: Char
sep = B.head separator

root :: String -> a -> IndexedTree a
root name dat = IT "\1" (T.singleton nn dat) (T.singleton n nn)
  where
    n = B.pack name
    nn = B.pack "\1"

-- Returns Nothing if parent is not found.  Ignores parent if the node
-- already exists, and only replaces the data.
addNode :: String -> String -> a -> IndexedTree a -> Maybe (IndexedTree a)
addNode parent name dat it@(IT nn dm nm) = case T.member n nm of
    True -> Just (updateNode (const dat) name it)
    False -> fmap add (T.lookup p nm)
  where
    p = B.pack parent
    n = B.pack name
    nn' = nextName nn
    add path = IT nn' (T.insert path' dat dm) (T.insert n path' nm)
      where
        path' = B.intercalate separator [path, B.pack nn']

-- Deletes node with its whole subtree.
deleteNode :: String -> IndexedTree a -> IndexedTree a
deleteNode name it@(IT nn dm nm) = case T.lookup n nm of
    Nothing -> it
    Just path -> IT nn dm' nm'
      where
        subkeys = T.keys (T.submap path dm)
        subnames = map (snd . B.breakEnd (==sep)) subkeys
        dm' = T.disunion dm (T.fromListS (map (\k -> (k,undefined)) subkeys))
        nm' = T.disunion nm (T.fromListS (map (\k -> (k,B.pack "")) subnames))
  where
    n = B.pack name

updateNode :: (a -> a) -> String -> IndexedTree a -> IndexedTree a
updateNode f name it@(IT nn dm nm) = case T.lookup n nm of
    Nothing -> it
    Just path -> IT nn (T.adjust f path dm) nm
  where
    n = B.pack name

getDataAt :: String -> IndexedTree a -> Maybe a
getDataAt name (IT _ dm nm) = flip T.lookup dm =<< T.lookup n nm
  where
    n = B.pack name

flattenTree :: (a -> b -> a) -> (a -> b -> r) -> a -> IndexedTree b -> [r]
flattenTree f g acc (IT _ dm _) = go (T.toList dm) [] [acc]
  where
    go [] _ _ = []
    go ((path,dat):pds') [] as@(acc:_) = g acc dat : go pds' [path] (f acc dat:as)
    go pds@((path,dat):pds') ps@(p:ps') as@(acc:as')
        | prefix    = g acc dat : go pds' (path:ps) (f acc dat:as)
        | otherwise = go pds ps' as'
      where
        prefix = (p `B.isPrefixOf` path) && (B.index path (B.length p) == sep)

nextName :: String -> String
nextName "" = "\1"
nextName ('\255':s) = '\1' : nextName s
nextName (c:s) = succ c : s