{-# LANGUAGE CPP #-}
module Graphics.UI.Gtk.ModelView.TreeStore (
  TreeStore,
  treeStoreNew,
  treeStoreNewDND,
  treeStoreDefaultDragSourceIface,
  treeStoreDefaultDragDestIface,
  treeStoreGetValue,
  treeStoreGetTree,
  treeStoreLookup,
  treeStoreSetValue,
  treeStoreInsert,
  treeStoreInsertTree,
  treeStoreInsertForest,
  treeStoreRemove,
  treeStoreClear,
  treeStoreChange,
  treeStoreChangeM,
  ) where
import Data.Bits
import Data.Word (Word32)
import Data.Maybe ( fromMaybe, isJust )
import Data.Tree
import Control.Monad ( when )
import Control.Exception (assert)
import Data.IORef
import Graphics.UI.Gtk.ModelView.Types
import Graphics.UI.Gtk.Types (GObjectClass(..))
import Graphics.UI.Gtk.ModelView.CustomStore
import Graphics.UI.Gtk.ModelView.TreeModel
import Graphics.UI.Gtk.ModelView.TreeDrag
import Control.Monad.Trans ( liftIO )
newtype TreeStore a = TreeStore (CustomStore (IORef (Store a)) a)
instance TypedTreeModelClass TreeStore
instance TreeModelClass (TreeStore a)
instance GObjectClass (TreeStore a) where
  toGObject (TreeStore tm) = toGObject tm
  unsafeCastGObject = TreeStore . unsafeCastGObject
type Depth = [Int]
data Store a = Store {
  depth :: Depth,
  content :: Cache a
}
treeStoreNew :: Forest a -> IO (TreeStore a)
treeStoreNew forest = treeStoreNewDND forest
                        (Just treeStoreDefaultDragSourceIface)
                        (Just treeStoreDefaultDragDestIface)
treeStoreNewDND :: Forest a 
  -> Maybe (DragSourceIface TreeStore a) 
  -> Maybe (DragDestIface TreeStore a) 
  -> IO (TreeStore a)
treeStoreNewDND forest mDSource mDDest = do
  storeRef <- newIORef Store {
      depth = calcForestDepth forest,
      content = storeToCache forest
    }
  let withStore f = readIORef storeRef >>= return . f
      withStoreUpdateCache f = do
        store <- readIORef storeRef
        let (result, cache') = f store
        writeIORef storeRef store { content = cache' }
        return result
  customStoreNew storeRef TreeStore TreeModelIface {
    treeModelIfaceGetFlags = return [],
    treeModelIfaceGetIter = \path -> withStore $
      \Store { depth = d } -> fromPath d path,
    treeModelIfaceGetPath = \iter -> withStore $
      \Store { depth = d } -> toPath d iter,
    treeModelIfaceGetRow  = \iter -> withStoreUpdateCache $
      \Store { depth = d, content = cache } ->
        case checkSuccess d iter cache of
          (True, cache'@((_, (Node { rootLabel = val }:_)):_)) ->
            (val, cache')
          _ -> error "TreeStore.getRow: iter does not refer to a valid entry",
    treeModelIfaceIterNext = \iter -> withStoreUpdateCache $
      \Store { depth = d, content = cache } -> iterNext d iter cache,
    treeModelIfaceIterChildren = \mIter -> withStoreUpdateCache $
      \Store { depth = d, content = cache } ->
      let iter = fromMaybe invalidIter mIter
       in iterNthChild d 0 iter cache,
    treeModelIfaceIterHasChild = \iter -> withStoreUpdateCache $
      \Store { depth = d, content = cache } ->
       let (mIter, cache') = iterNthChild d 0 iter cache
        in (isJust mIter, cache'),
    treeModelIfaceIterNChildren = \mIter -> withStoreUpdateCache $
      \Store { depth = d, content = cache } ->
      let iter = fromMaybe invalidIter mIter
       in iterNChildren d iter cache,
    treeModelIfaceIterNthChild = \mIter idx  -> withStoreUpdateCache $
      \Store { depth = d, content = cache } ->
      let iter = fromMaybe invalidIter mIter
       in iterNthChild d idx iter cache,
    treeModelIfaceIterParent = \iter -> withStore $
      \Store { depth = d } -> iterParent d iter,
    treeModelIfaceRefNode = \_ -> return (),
    treeModelIfaceUnrefNode = \_ -> return ()
   } mDSource mDDest
treeStoreDefaultDragSourceIface :: DragSourceIface TreeStore row
treeStoreDefaultDragSourceIface = DragSourceIface {
    treeDragSourceRowDraggable = \_ _-> return True,
    treeDragSourceDragDataGet = treeSetRowDragData,
    treeDragSourceDragDataDelete = \model dest@(_:_) -> do
            liftIO $ treeStoreRemove model dest
            return True
  }
treeStoreDefaultDragDestIface :: DragDestIface TreeStore row
treeStoreDefaultDragDestIface = DragDestIface {
    treeDragDestRowDropPossible = \model dest -> do
      mModelPath <- treeGetRowDragData
      case mModelPath of
        Nothing -> return False
        Just (model', source) -> return (toTreeModel model==toTreeModel model'),
    treeDragDestDragDataReceived = \model dest@(_:_) -> do
      mModelPath <- treeGetRowDragData
      case mModelPath of
        Nothing -> return False
        Just (model', source@(_:_)) ->
          if toTreeModel model/=toTreeModel model' then return False
          else liftIO $ do
            row <- treeStoreGetTree model source
            treeStoreInsertTree model (init dest) (last dest) row
            return True
  }
bitsNeeded :: Word32 -> Int
bitsNeeded n = bitsNeeded' 0 n
  where bitsNeeded' b 0 = b
        bitsNeeded' b n = bitsNeeded' (b+1) (n `shiftR` 1)
getBitSlice :: TreeIter -> Int -> Int -> Word32
getBitSlice (TreeIter _ a b c) off count =
      getBitSliceWord a  off     count
  .|. getBitSliceWord b (off-32) count
  .|. getBitSliceWord c (off-64) count
  where getBitSliceWord :: Word32 -> Int -> Int -> Word32
        getBitSliceWord word off count =
          word `shift` (-off) .&. (1 `shiftL` count - 1)
setBitSlice :: TreeIter -> Int -> Int -> Word32 -> TreeIter
setBitSlice (TreeIter stamp a b c) off count value =
  assert (value < 1 `shiftL` count) $
  TreeIter stamp
           (setBitSliceWord a  off     count value)
           (setBitSliceWord b (off-32) count value)
           (setBitSliceWord c (off-64) count value)
  where setBitSliceWord :: Word32 -> Int -> Int -> Word32 -> Word32
        setBitSliceWord word off count value =
          let mask = (1 `shiftL` count - 1) `shift` off
           in (word .&. complement mask) .|. (value `shift` off)
invalidIter :: TreeIter
invalidIter = TreeIter 0 0 0 0
calcForestDepth :: Forest a -> Depth
calcForestDepth f = map bitsNeeded $
                    takeWhile (/=0) $
                    foldr calcTreeDepth (repeat 0) f
  where
  calcTreeDepth Node { subForest = f } (d:ds) =
      (d+1): zipWith max ds (foldr calcTreeDepth (repeat 0) f)
toPath :: Depth -> TreeIter -> TreePath
toPath d iter = gP 0 d
  where
  gP pos [] = []
  gP pos (d:ds) = let idx = getBitSlice iter pos d in
                  if idx==0 then [] else fromIntegral (idx-1) : gP (pos+d) ds
fromPath :: Depth -> TreePath -> Maybe TreeIter
fromPath = fP 0 invalidIter
  where
  fP pos ti _ [] = Just ti 
  fP pos ti [] _ = Nothing
  fP pos ti (d:ds) (p:ps) = let idx = fromIntegral (p+1) in
    if idx >= bit d then Nothing else
    fP (pos+d) (setBitSlice ti pos d idx) ds ps
type Cache a = [(TreeIter, Forest a)]
storeToCache :: Forest a -> Cache a
storeToCache [] = []
storeToCache forest = [(invalidIter, [Node root forest])]
  where
  root = error "TreeStore.storeToCache: accessed non-exitent root of tree"
cacheToStore :: Cache a -> Forest a
cacheToStore [] = []
cacheToStore cache = case last cache of (_, [Node _ forest]) -> forest
advanceCache :: Depth -> TreeIter -> Cache a -> Cache a
advanceCache depth goal [] = []
advanceCache depth goal cache@((rootIter,_):_) =
  moveToSameLevel 0 depth
  where
  moveToSameLevel pos [] = cache
  moveToSameLevel pos (d:ds) =
    let
      goalIdx = getBitSlice goal pos d
      curIdx = getBitSlice rootIter pos d
      isNonZero pos d (ti,_) = getBitSlice ti pos d/=0
    in
    if goalIdx==curIdx then moveToSameLevel (pos+d) ds else
    if goalIdx==0 then dropWhile (isNonZero pos d) cache else
    if curIdx==0 then moveToChild pos (d:ds) cache else
    if goalIdx<curIdx then
      moveToChild pos (d:ds) (dropWhile (isNonZero pos d) cache)
    else let
      
      
      moveWithinLevel pos d ((ti,forest):parents) = let
          diff = fromIntegral (goalIdx-curIdx)
          (dropped, remain) = splitAt diff forest
          advance = length dropped
          ti' = setBitSlice ti pos d (curIdx+fromIntegral advance)
        in
        if advance==diff then moveToChild (pos+d) ds ((ti',remain):parents)
        else (ti',remain):parents 
    in moveWithinLevel pos d $ case ds of
        [] -> cache
        (d':_) -> dropWhile (isNonZero (pos+d) d') cache
  
  
  
  moveToChild :: Int -> Depth -> Cache a -> Cache a
  moveToChild pos [] cache = cache 
  moveToChild pos (d:ds) cache@((ti,forest):parents)
    | getBitSlice goal pos d == 0 = cache
    | otherwise = case forest of
      [] -> cache 
      Node { subForest = children }:_ ->
        let
          childIdx :: Int
          childIdx = fromIntegral (getBitSlice goal pos d)-1
          (dropped, remain) = splitAt childIdx children
          advanced = length dropped
          ti' = setBitSlice ti pos d (fromIntegral advanced+1)
        in if advanced<childIdx then ((ti',remain):cache) else
           moveToChild (pos+d) ds ((ti',remain):cache)
checkSuccess :: Depth -> TreeIter -> Cache a -> (Bool, Cache a)
checkSuccess depth iter cache = case advanceCache depth iter cache of
    cache'@((cur,sibs):_) -> (cmp cur iter && not (null sibs), cache')
    [] -> (False, [])
  where
  cmp (TreeIter _ a1 b1 c1) (TreeIter _ a2 b2 c2) =
      a1==a2 && b1==b2 && c2==c2
getTreeIterLeaf :: Depth -> TreeIter -> (Int, Int, Int)
getTreeIterLeaf ds ti = gTIL 0 0 ds
  where
  gTIL pos dCur (dNext:ds)
    | getBitSlice ti (pos+dCur) dNext==0 = (pos,dCur,dNext)
    | otherwise = gTIL (pos+dCur) dNext ds
  gTIL pos d [] = (pos, d, 0)
iterNext :: Depth -> TreeIter -> Cache a -> (Maybe TreeIter, Cache a)
iterNext depth iter cache = let
    (pos,leaf,_child) = getTreeIterLeaf depth iter
    curIdx = getBitSlice iter pos leaf
    nextIdx = curIdx+1
    nextIter = setBitSlice iter pos leaf nextIdx
  in
  if nextIdx==bit leaf then (Nothing, cache) else
  case checkSuccess depth nextIter cache of
    (True, cache) -> (Just nextIter, cache)
    (False, cache) -> (Nothing, cache)
iterNthChild :: Depth -> Int -> TreeIter -> Cache a  ->
                (Maybe TreeIter, Cache a)
iterNthChild depth childIdx_ iter cache = let
    (pos,leaf,child) = getTreeIterLeaf depth iter
    childIdx = fromIntegral childIdx_+1
    nextIter = setBitSlice iter (pos+leaf) child childIdx
  in
  if childIdx>=bit child then (Nothing, cache) else
  case checkSuccess depth nextIter cache of
    (True, cache) -> (Just nextIter, cache)
    (False, cache) -> (Nothing, cache)
iterNChildren :: Depth -> TreeIter -> Cache a -> (Int, Cache a)
iterNChildren depth iter cache = case checkSuccess depth iter cache of
  (True, cache@((_,Node { subForest = forest}:_):_)) -> (length forest, cache)
  (_, cache) -> (0, cache)
iterParent :: Depth -> TreeIter -> Maybe TreeIter
iterParent depth iter = let
    (pos,leaf,_child) = getTreeIterLeaf depth iter
  in if pos==0 then Nothing else
     if getBitSlice iter pos leaf==0 then Nothing else
     Just (setBitSlice iter pos leaf 0)
treeStoreInsertForest ::
    TreeStore a 
 -> TreePath    
 -> Int         
 -> Forest a    
 -> IO ()
treeStoreInsertForest (TreeStore model) path pos nodes = do
  customStoreInvalidateIters model
  (idx, toggle) <- atomicModifyIORef (customStoreGetPrivate model) $
    \store@Store { depth = d, content = cache } ->
    case insertIntoForest (cacheToStore cache) nodes path pos of
      Nothing -> error ("treeStoreInsertForest: path does not exist " ++ show path)
      Just (newForest, idx, toggle) ->
       let depth = calcForestDepth newForest
        in (Store { depth = depth,
                    content = storeToCache newForest },
           (idx, toggle))
  Store { depth = depth } <- readIORef (customStoreGetPrivate model)
  let rpath = reverse path
  stamp <- customStoreGetStamp model
  sequence_ [ let p' = reverse p
                  Just iter = fromPath depth p'
               in treeModelRowInserted model p' (treeIterSetStamp iter stamp)
            | (i, node) <- zip [idx..] nodes
            , p <- paths (i : rpath) node ]
  let Just iter = fromPath depth path
  when toggle $ treeModelRowHasChildToggled model path
                (treeIterSetStamp iter stamp)
  where paths :: TreePath -> Tree a -> [TreePath]
        paths path Node { subForest = ts } =
          path : concat [ paths (n:path) t | (n, t) <- zip [0..] ts ]
treeStoreInsertTree ::
    TreeStore a 
 -> TreePath    
 -> Int         
 -> Tree a      
 -> IO ()
treeStoreInsertTree store path pos node =
  treeStoreInsertForest store path pos [node]
treeStoreInsert ::
    TreeStore a 
 -> TreePath    
 -> Int         
 -> a           
 -> IO ()
treeStoreInsert store path pos node =
  treeStoreInsertForest store path pos [Node node []]
insertIntoForest :: Forest a -> Forest a -> TreePath -> Int ->
                    Maybe (Forest a, Int, Bool)
insertIntoForest forest nodes [] pos
  | pos<0 = Just (forest++nodes, length forest, null forest)
  | otherwise = Just (prev++nodes++next, length prev, null forest)
    where (prev, next) = splitAt pos forest
insertIntoForest forest nodes (p:ps) pos = case splitAt p forest of
  (prev, []) -> Nothing
  (prev, Node { rootLabel = val,
                subForest = for}:next) ->
    case insertIntoForest for nodes ps pos of
      Nothing -> Nothing
      Just (for, pos, toggle) -> Just (prev++Node { rootLabel = val,
                                                    subForest = for }:next,
                                       pos, toggle)
treeStoreRemove :: TreeStore a -> TreePath -> IO Bool
  
treeStoreRemove (TreeStore model) [] = return False
treeStoreRemove (TreeStore model) path = do
  customStoreInvalidateIters model
  (found, toggle) <- atomicModifyIORef (customStoreGetPrivate model) $
    \store@Store { depth = d, content = cache } ->
    if null cache then (store, (False, False)) else
    case deleteFromForest (cacheToStore cache) path of
      Nothing -> (store, (False, False))
      Just (newForest, toggle) ->
        (Store { depth = d, 
                 content = storeToCache newForest }, (True, toggle))
  when found $ do
    when (toggle && not (null path)) $ do
      Store { depth = depth } <- readIORef (customStoreGetPrivate model)
      let parent = init path
          Just iter = fromPath depth parent
      treeModelRowHasChildToggled model parent iter
    treeModelRowDeleted model path
  return found
treeStoreClear :: TreeStore a -> IO ()
treeStoreClear (TreeStore model) = do
  customStoreInvalidateIters model
  Store { content = cache } <- readIORef (customStoreGetPrivate model)
  let forest = cacheToStore cache
  writeIORef (customStoreGetPrivate model) Store {
      depth = calcForestDepth [],
      content = storeToCache []
    }
  let loop (-1) = return ()
      loop   n  = treeModelRowDeleted model [n] >> loop (n-1)
  loop (length forest - 1)
deleteFromForest :: Forest a -> TreePath -> Maybe (Forest a, Bool)
deleteFromForest forest [] = Just ([], False)
deleteFromForest forest (p:ps) =
  case splitAt p forest of
    (prev, kill@Node { rootLabel = val,
                       subForest = for}:next) ->
      if null ps then Just (prev++next, null prev && null next) else
      case deleteFromForest for ps of
        Nothing -> Nothing
        Just (for,toggle) -> Just (prev++Node {rootLabel = val,
                                               subForest = for }:next, toggle)
    (prev, []) -> Nothing
treeStoreSetValue :: TreeStore a -> TreePath -> a -> IO ()
treeStoreSetValue store path value = treeStoreChangeM store path (\_ -> return value)
                                  >> return ()
treeStoreChange :: TreeStore a -> TreePath -> (a -> a) -> IO Bool
treeStoreChange store path func = treeStoreChangeM store path (return . func)
treeStoreChangeM :: TreeStore a -> TreePath -> (a -> IO a) -> IO Bool
treeStoreChangeM (TreeStore model) path act = do
  customStoreInvalidateIters model
  store@Store { depth = d, content = cache } <-
      readIORef (customStoreGetPrivate model)
  (store'@Store { depth = d, content = cache }, found) <- do
    mRes <- changeForest (cacheToStore cache) act path
    return $ case mRes of
      Nothing -> (store, False)
      Just newForest -> (Store { depth = d,
                                 content = storeToCache newForest }, True)
  writeIORef (customStoreGetPrivate model) store'
  let Just iter = fromPath d path
  stamp <- customStoreGetStamp model
  when found $ treeModelRowChanged model path (treeIterSetStamp iter stamp)
  return found
changeForest :: Forest a -> (a -> IO a) -> TreePath -> IO (Maybe (Forest a))
changeForest forest act [] = return Nothing
changeForest forest act (p:ps) = case splitAt p forest of
  (prev, []) -> return Nothing
  (prev, Node { rootLabel = val,
                subForest = for}:next) ->
    if null ps then do
      val' <- act val
      return (Just (prev++Node { rootLabel = val',
                                 subForest = for }:next))
    else do
      mFor <- changeForest for act ps
      case mFor of
        Nothing -> return Nothing
        Just for -> return $ Just (prev++Node { rootLabel = val,
                                                subForest = for }:next)
treeStoreGetValue :: TreeStore a -> TreePath -> IO a
treeStoreGetValue model path = fmap rootLabel (treeStoreGetTree model path)
treeStoreGetTree :: TreeStore a -> TreePath -> IO (Tree a)
treeStoreGetTree (TreeStore model) path = do
  store@Store { depth = d, content = cache } <-
      readIORef (customStoreGetPrivate model)
  case fromPath d path of
    (Just iter) -> do
      let (res, cache') = checkSuccess d iter cache
      writeIORef (customStoreGetPrivate model) store { content = cache' }
      case cache' of
        ((_,node:_):_) | res -> return node
        _ -> fail ("treeStoreGetTree: path does not exist " ++ show path)
    _ -> fail ("treeStoreGetTree: path does not exist " ++ show path)
treeStoreLookup :: TreeStore a -> TreePath -> IO (Maybe (Tree a))
treeStoreLookup (TreeStore model) path = do
  store@Store { depth = d, content = cache } <-
      readIORef (customStoreGetPrivate model)
  case fromPath d path of
    (Just iter) -> do
      let (res, cache') = checkSuccess d iter cache
      writeIORef (customStoreGetPrivate model) store { content = cache' }
      case cache' of
        ((_,node:_):_) | res -> return (Just node)
        _ -> return Nothing
    _ -> return Nothing