{-# LANGUAGE GeneralizedNewtypeDeriving, DuplicateRecordFields #-}
import Control.Monad
import Control.Monad.State
import Control.Monad.Reader
import Data.List

data Type = V | I | F | B deriving (Show, Eq)

data Symbol = Program | Variable { name :: String, vartype :: Type } | Function { name :: String, returntype :: Type }
  deriving (Show, Eq)

data Tree a = Node { value :: a, children :: [Tree a] } deriving (Show, Eq)

leaf a = Node a []

data Ctx a = Top | Child {
  parent :: Ctx a,
  val :: a,
  lefts :: [Tree a],
  rights :: [Tree a]
                         } deriving (Show, Eq)

data Loc c a = Loc { ctx :: c, struct :: a } deriving (Show, Eq)


newtype Travel loc a = Travel { unT :: State loc a } deriving (Functor, Applicative, Monad)

type TreeLoc a = Loc (Ctx a) (Tree a)
type TravelTree a = Travel (TreeLoc a) (Tree a)

topLoc :: Tree a -> TreeLoc a
topLoc = Loc Top

traverse :: Loc c a -> Travel (Loc c a) a -> a
traverse start tt = evalState (unT tt) start

modifyStruct :: (a -> a) -> Travel (Loc c a) a
modifyStruct f = Travel $ do
  loc <- get
  let (Loc ctx str) = loc
   in put $ Loc ctx (f str)
  liftM struct get

putStruct :: b -> Travel (Loc c b) b
putStruct = modifyStruct . const

getStruct :: Travel (Loc c a) a
getStruct = modifyStruct id

symTable = Node Program [leaf $ Function "foo" V, leaf $ Variable "kaas" I]

addSymbol :: Symbol -> Tree Symbol -> Tree Symbol
addSymbol sym (Node val children) = Node val (Node sym []:children)

down :: Int -> TravelTree a
down idx = Travel $ do
  Loc ctx (Node val children) <- get :: State (TreeLoc a) (TreeLoc a)
  let (left, node':right) = splitAt idx children
  put (Loc (Child ctx val left right) node')
  return node'

up :: TravelTree a
up = Travel $ do
  Loc ctx node <- get
  case ctx of
    Top -> return node
    Child ctx' val lefts rights -> let node' = Node val (lefts ++ node:rights)
                                    in put (Loc ctx' node') >> return node'


travelAddSymbol :: Symbol -> TravelTree Symbol
travelAddSymbol = modifyStruct . addSymbol

test :: TravelTree Symbol
test = do
  down 0
  travelAddSymbol (Variable "kroket" B)
  up

data AST = ASTNode

type Traversal a = State (TravelTree Symbol, a) a

traverseAddSymbol :: Symbol -> Traversal ()
traverseAddSymbol sym = do
  modify (\(st, ast) -> (st >> travelAddSymbol sym, ast))

test = do
  
