module Jungle
( Jungle
, Node
, Path
, Position
, createJungle
, createTree
, getRootNode
, updateRootNode
, updateRootNodeWith
, addNewChildAt
, deleteChildAt
, putAttribute
, deleteAttribute
, getAttributes
, drawNode
, printAttributes
, size
) 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 :: (M.Map String Tree) } 

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

data Node = Node
          { children   :: Children
          , attributes :: Attributes
          } deriving (Show)

-- Mapのkeyやvalueの型は固定しているが、
-- jungle作成時、もしくは木作成時に与えるように変更することも容易
newtype Children   = Children   { getChildrenMap :: (M.Map Int Node) } deriving (Show)
newtype Attributes = Attributes { getAttributesMap :: (M.Map String B.ByteString) } deriving (Show)

type Path = [Int]
type Position = Int

createJungle :: Jungle
createJungle = Jungle M.empty

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

-- getRootNodeやupdateRootNodeをJungleとTree名を取るようにしたため不要
getTreeByName :: Jungle -> String -> Maybe Tree
getTreeByName (Jungle map) tree_name = M.lookup tree_name map

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

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

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

-- 新しい木構造を作成し、最新のルートノードとなるNodeを返す
-- Pathの位置にNodeが存在しない場合どうするか？
addNewChildAt :: Node -> Path -> Position -> Node
addNewChildAt parent []     pos = addChild parent pos emptyNode
addNewChildAt parent (x:xs) pos = addChild parent x $ addNewChildAt x_node xs pos
  where
    map = getChildrenMap $ 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 = getChildrenMap $ children node
    new_child = Children (M.insert pos child map)
    attr = attributes node

-- 子を削除した新しいNodeを追加するのに等しい
-- addNewChildAtのコピペ、一般化して関数に抽出したい
-- 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 = getChildrenMap $ 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 = getChildrenMap $ children node
    new_child = Children (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 = getChildrenMap $ 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 = getAttributesMap $ attributes node
    attr = Attributes (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 = getChildrenMap $ 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 = getAttributesMap $ attributes node
    attr = Attributes (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 = getChildrenMap $ 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 = getAttributesMap $ attributes target

-- デバッグ用表示関数

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

draw :: String -> Node -> [String]
draw string node = string : drawSubTrees keys
  where
    map = getChildrenMap $ 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  = getAttributesMap $ 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 = getChildrenMap $ 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 = getChildrenMap $ children node
    keys = M.keys map
    subTreesSize [] = 0
    subTreesSize (x:xs) = size (getNode x) + subTreesSize xs
    getNode x = fromJust $ M.lookup x map