module RedBlackTree where

open import stack
open import Level

record TreeMethods {n m : Level } {a : Set n } {t : Set m } (treeImpl : Set n ) : Set (m Level.⊔ n) where
  field
    putImpl : treeImpl -> a -> (treeImpl -> t) -> t
    getImpl  : treeImpl -> (treeImpl -> Maybe a -> t) -> t
open TreeMethods

record Tree  {n m : Level } {a : Set n } {t : Set m } (treeImpl : Set n ) : Set (m Level.⊔ n) where
  field
    tree : treeImpl
    treeMethods : TreeMethods {n} {m} {a} {t} treeImpl
  putTree : a -> (Tree treeImpl -> t) -> t
  putTree d next = putImpl (treeMethods ) tree d (\t1 -> next (record {tree = t1 ; treeMethods = treeMethods} ))
  getTree : (Tree treeImpl -> Maybe a -> t) -> t
  getTree next = getImpl (treeMethods ) tree (\t1 d -> next (record {tree = t1 ; treeMethods = treeMethods} ) d )

open Tree

data Color {n : Level } : Set n where
  Red   : Color
  Black : Color

data CompareResult {n : Level } : Set n where
  LT : CompareResult
  GT : CompareResult
  EQ : CompareResult

record Node {n : Level } (a k : Set n) : Set n where
  inductive
  field
    key   : k
    value : a
    right : Maybe (Node a k)
    left  : Maybe (Node a k)
    color : Color {n}
open Node

record RedBlackTree {n m : Level } {t : Set m} (a k si : Set n) : Set (m Level.⊔ n) where
  field
    root : Maybe (Node a k)
    nodeStack : Stack {n} {m} (Node a k) {t} si
    compare : k -> k -> CompareResult {n}

open RedBlackTree

open Stack


insertNode : {n m : Level } {t : Set m } {a k si : Set n} -> RedBlackTree {n} {m} {t} a k si -> Stack (Node a k) si -> Node a k -> (RedBlackTree {n} {m} {t} a k si -> t) -> t
insertNode tree s datum next = get2Stack s (\ s d1 d2 -> {!!} tree s datum d1 d2 next)

findNode : {n m : Level } {a k si : Set n} {t : Set m} -> RedBlackTree {n} {m} {t} a k si -> Stack (Node a k) si -> (Node a k) -> (Node a k) -> (RedBlackTree {n} {m} {t} a k si -> Stack (Node a k) si -> Node a k -> t) -> t
findNode {n} {m} {a} {k} {si} {t} tree s n0 n1 next = pushStack s n1 (\ s -> findNode1 s n1)
  where
    findNode2 : Stack (Node a k) si -> (Maybe (Node a k)) -> t
    findNode2 s Nothing = next tree s n0
    findNode2 s (Just n) = findNode tree s n0 n next
    findNode1 : Stack (Node a k) si -> (Node a k)  -> t
    findNode1 s n1 with (compare tree (key n0) (key n1))
    ...                                | EQ = next tree s n0 
    ...                                | GT = findNode2 s (right n1)
    ...                                | LT = findNode2 s (left n1)
      where
        -- findNode3 : Stack (Node a k) si -> (Maybe (Node a k)) -> t
        -- findNode3 s nothing = next tree s n0
        -- findNode3 s (Just n) = 
        --           popStack (nodeStack tree) (\s d -> findNode3 s d)


leafNode : {n m : Level } {a k si : Set n} {t : Set m} -> k -> a -> Node a k
leafNode k1 value = record {
    key   = k1 ;
    value = value ;
    right = Nothing ;
    left  = Nothing ;
    color = Black 
    }

putRedBlackTree : {n m : Level } {a k si : Set n} {t : Set m} -> RedBlackTree {n} {m} {t} a k si -> k -> a -> (RedBlackTree {n} {m} {t} a k si -> t) -> t
putRedBlackTree {n} {m} {a} {k} {si} {t} tree k1 value next with (root tree)
...                                | Nothing = next (record tree {root = Just (leafNode k1 value) })
...                                | Just n2  = findNode tree (nodeStack tree) (leafNode {n} {m} {a} {k} {si} {t} k1 value) n2 (\ tree1 s n1 -> insertNode tree1 s n1 next)

getRedBlackTree : {n m : Level } {a k si : Set n} {t : Set m} -> RedBlackTree {n} {m} {t} a k si -> k -> (RedBlackTree {n} {m} {t} a k si -> (Maybe (Node a k)) -> t) -> t
getRedBlackTree {_} {_} {a} {k} {_} {t} tree k1 cs = checkNode (root tree)
  where
    checkNode : Maybe (Node a k) -> t
    checkNode Nothing = cs tree Nothing
    checkNode (Just n) = search n
      where
        search : Node a k -> t
        search n with compare tree k1 (key n)
        search n | LT = checkNode (left n)
        search n | GT = checkNode (right n)
        search n | EQ = cs tree (Just n)