module Jungle
( Jungle
, Node
, Path
, Position
, createJungle
, createTree
, getRootNode
, updateRootNode
, updateRootNodeWith
, addNewChildAt
, deleteChildAt
, putAttribute
, deleteAttribute
, getAttributes
, getChildren
, getChildrenWithKey
, assocs
, numOfChild
, currentChild
, drawNode
, printAttributes
, size
, attrSize
) where

import qualified Data.Map as M
import qualified Data.ByteString.Lazy.Char8 as B
import Control.Concurrent.STM
import Data.Maybe (fromJust)

data Jungle = Jungle { getJungleMap :: (TVar (M.Map String Tree)) } 

data Tree = Tree
          { rootNode :: (TVar Node)
          , treeName :: String }

data Node = Node
          { children   :: (M.Map Int Node)
          , attributes :: (M.Map String B.ByteString) }

type Path = [Int]
type Position = Int

createJungle :: IO Jungle
createJungle = atomically $ do
    map <- newTVar M.empty
    return (Jungle map)

-- 同じ名前のTreeが存在する場合、上書きする
-- Eitherなどで失敗させるほうがいいかもしれない
createTree :: Jungle -> String -> IO ()
createTree (Jungle tmap) tree_name = atomically $ do
    map <- readTVar tmap
    tree <- emptyTree tree_name
    writeTVar tmap (M.insert tree_name tree map)

emptyTree :: String -> STM Tree
emptyTree tree_name = do
    node <- newTVar emptyNode
    return (Tree node tree_name)

emptyNode :: Node
emptyNode = Node (M.empty) (M.empty)

getRootNode :: Jungle -> String -> IO Node
getRootNode (Jungle tmap) tree_name = atomically $ do 
    map <- readTVar tmap
    readTVar (root_node map)
    where
      root_node map = case M.lookup tree_name map of
                        Just x -> rootNode x

-- ルートノードを更新する
updateRootNode :: Jungle -> String -> Node -> IO ()
updateRootNode (Jungle tmap) tree_name node = 
  atomically $ do 
    map <- readTVar tmap
    writeTVar (root_node map) node
  where
    root_node map = case M.lookup tree_name map of
                      Just x -> rootNode x

updateRootNodeWith :: (Node -> Node) -> Jungle -> String -> IO ()
updateRootNodeWith f (Jungle tmap) tree_name = 
  atomically $ do
    map <- readTVar tmap
    n <- readTVar (root_node map)
    writeTVar (root_node map) (f n)
  where
    root_node map = case M.lookup tree_name map of
                      Just x -> rootNode x

-- 新しい木構造を作成し、最新のルートノードとなるNodeを返す
-- Pathの位置にNodeが存在しない場合どうするか？
-- 1 -> 2 -> 3と勝手にPositionをインクリメントしながら追加する
addNewChildAt :: Node -> Path -> Node
addNewChildAt parent []     = addChild' parent emptyNode
addNewChildAt parent (x:xs) = addChild parent x $ addNewChildAt x_node xs
  where
    map = children parent
    x_node = case M.lookup x map of
               Just x -> x

-- 子を追加したNodeを新しく作成して返す
-- 同じ位置に既に子がある場合は？
-- 現在はinsertでそのまま上書き
addChild :: Node -> Position -> Node -> Node
addChild node pos child = Node new_child attr
  where
    map = children node
    new_child = M.insert pos child map
    attr = attributes node

-- 自動でPositionをincrementして追加してくれるaddChild
addChild' :: Node -> Node -> Node
addChild' node child = Node new_child attr
  where
    map = children node
    pos = (M.size map) + 1
    new_child = M.insert pos child map
    attr = attributes node

-- 子を削除した新しいNodeを追加するのに等しい
deleteChildAt :: Node -> Path -> Position -> Node
deleteChildAt parent []     pos = deleteChild parent pos
deleteChildAt parent (x:xs) pos = addChild parent x $ deleteChildAt x_node xs pos
  where
    map = children parent
    x_node = case M.lookup x map of
               Just x -> x

deleteChild :: Node -> Position -> Node
deleteChild node pos = Node new_child attr
  where
    map = children node
    new_child = M.delete pos map
    attr = attributes node

-- attribute関連はaddNewChildAtを利用する
-- 現在はコピペ、関数に抽出したい
putAttribute :: Node -> Path -> String -> B.ByteString -> Node
putAttribute parent []     key value = putAttr parent key value
putAttribute parent (x:xs) key value = addChild parent x $ putAttribute x_node xs key value
  where
    map = children parent
    x_node = case M.lookup x map of
               Just x -> x

putAttr :: Node -> String -> B.ByteString -> Node
putAttr node key value = Node child attr
  where
    map = attributes node
    attr = M.insert key value map
    child = children node

deleteAttribute :: Node -> Path -> String -> Node
deleteAttribute parent []     key = deleteAttr parent key
deleteAttribute parent (x:xs) key = addChild parent x $ deleteAttribute x_node xs key
  where
    map = children parent
    x_node = case M.lookup x map of
               Just x -> x
    
deleteAttr :: Node -> String -> Node
deleteAttr node key = Node child attr
  where
    map = attributes node
    attr = M.delete key map
    child = children node

-- データを取り出す関連の関数
-- getNodeとgetchild, getattributeなど？

getNode :: Node -> Path -> Node
getNode node []     = node
getNode node (x:xs) = getNode child xs
  where
    map = children node
    child = case M.lookup x map of
              Just x -> x

getAttributes :: Node -> Path -> String -> Maybe B.ByteString
getAttributes node path key = M.lookup key map
  where
    target = getNode node path
    map = attributes target

-- 対象のNodeの全ての子を返す
getChildren :: Node -> Path -> [Node]
getChildren node path = M.elems map
  where
    target = getNode node path
    map = children target

getChildrenWithKey :: Node -> Path -> [(Int, Node)]
getChildrenWithKey node path = M.assocs map
  where
    target = getNode node path
    map = children target

-- Attributeの連想リストを返す
assocs :: Node -> Path ->  [(String, B.ByteString)]
assocs node path = M.assocs map
  where
    target = getNode node path
    map = attributes target

-- 対象のNodeの子供の数を教えてくれる
numOfChild :: Node -> Path -> Int
numOfChild node path = M.size map
  where
    target = getNode node path
    map = children target

-- foucus
-- 対象のpathのノードの最新の子を返す
-- 編集する際に使うのは無理
currentChild :: Node -> Path -> Maybe Node
currentChild node path = M.lookup pos map
  where
    target = getNode node path
    map = children target
    pos = M.size map

-- デバッグ用表示関数

-- 現在の木構造を整形して表示
drawNode :: Node -> String
drawNode node = unlines $ draw "root" node

draw :: String -> Node -> [String]
draw string node = string : drawSubTrees keys
  where
    map = children node
    keys = M.keys map
    drawSubTrees [] = []
    drawSubTrees [t] = 
      "|" : shift "`- " "   " (draw (show t) (fromJust $ M.lookup t map))
    drawSubTrees (t:ts) = 
      "|" : shift "+- " "|  " (draw (show t) (fromJust $ M.lookup t map )) ++ drawSubTrees ts
    shift first other = zipWith (++) (first : repeat other)

-- Attributesを全て表示
printAttributes :: Node -> String
printAttributes node = unlines $ printAttr "root" node

printAttr :: String -> Node -> [String]
printAttr string node =
    if not $ M.null attr_map
      then ("Node: " ++ string) : ("  " ++ attr) : printSubTrees keys
      else printSubTrees keys
  where
    attr_map  = attributes node
    show_attr [] = []
    show_attr [x] = fst x ++ ": " ++ (B.unpack $ snd x)
    show_attr (x:xs) = fst x ++ ": " ++ (B.unpack $ snd x) ++ "\n  " ++ show_attr xs
    attr = show_attr $ M.assocs attr_map
    map = children node
    keys = M.keys map
    printSubTrees []     = []
    printSubTrees (x:xs) = printAttr (string ++ "-" ++ (show x)) (fromJust $ M.lookup x map) ++ printSubTrees xs

-- ルートノードの下にいくつの子があるか数える
size :: Node -> Int
size node = M.size map + subTreesSize keys
  where
    map = children node
    keys = M.keys map
    subTreesSize [] = 0
    subTreesSize (x:xs) = size (getNode x) + subTreesSize xs
    getNode x = fromJust $ M.lookup x map

-- Attributesの数を調べる
attrSize :: Node -> Int
attrSize node = M.size attr_map + subTreesSize keys
  where
    attr_map = attributes node
    map = children node
    keys = M.keys map
    subTreesSize [] = 0
    subTreesSize (x:xs) = attrSize (getNode x) + subTreesSize xs
    getNode x = fromJust $ M.lookup x map