{-# LANGUAGE OverloadedStrings #-}
module NestingXML where

import Data.Conduit
import Data.XML.Types
import Control.Monad.State
import qualified Data.Text as S
import Data.List (find)
-- import qualified Text.XML.Stream.Parse as Parse (content)

type Lang = S.Text

data StrictList a = a :! !(StrictList a) | StrictNil

top ( a :! as ) = Just a
top _           = Nothing

data XMLState = XMLState { 
        nestingLevel :: Int,
        langStack :: StrictList (Int,Lang)
}

type NestingXML o m a = StateT XMLState (ConduitM Event o m) a


runNestingXML :: NestingXML o m a -> Int -> ConduitM Event o m (a,XMLState)
runNestingXML m lvl = runStateT m (XMLState lvl StrictNil)

doNestingXML :: Monad m => NestingXML o m a -> ConduitM Event o m a
doNestingXML = fmap fst . flip runNestingXML 0

nesting :: Monad m => NestingXML o m Int
nesting = fmap nestingLevel get

xmlLang :: Monad m => NestingXML o m (Maybe Lang)
xmlLang = fmap (fmap snd . top . langStack) get

lookupLang attrs =
    case find ( (=="xml:lang") . fst) attrs of
        Just (_,[ContentText x])   -> Just x
        Just (_,[ContentEntity x]) -> Just x
        _                          -> Nothing

awaitXML :: Monad m => NestingXML o m (Maybe Event)
awaitXML = do
    XMLState lvl langs <- get
    xml <- lift await
    put $ case xml of
              Just (EventBeginElement _ attrs) -> 
                case lookupLang attrs of
                  Nothing   -> XMLState (lvl+1) langs
                  Just lang -> XMLState (lvl+1) ( (lvl+1,lang) :! langs)
              Just (EventEndElement _)     -> 
                case langs of 
                  (llvl,_) :! ls | llvl==lvl -> XMLState (lvl-1) ls
                  _              | otherwise -> XMLState (lvl-1) langs
              _  -> XMLState lvl langs
    return xml

withXML
  :: Monad m => (Event -> NestingXML o m ()) -> NestingXML o m ()
withXML f = do
    xml <- awaitXML
    maybe (return ()) f xml


maybeXML
  :: Monad m => NestingXML o m a -> (Event -> NestingXML o m a) -> NestingXML o m a
maybeXML whenNothing withJust = do
    xml <- awaitXML
    maybe whenNothing withJust xml


awaitCloser :: Monad m => Int -> NestingXML o m ()
awaitCloser lvl = do
    fix $ \loop -> do
        lvl' <- nesting
        when (lvl' >= lvl) $ do
        withXML $ \xml -> do
            loop

nextElement :: Monad m => NestingXML o m (Maybe Event)
nextElement = do
    lvl <- nesting
    fix $ \loop  -> do
        xml <- awaitXML
        case xml of
            Nothing                      -> return Nothing
            Just (EventBeginElement _ _) -> do
                return xml
            Just _ -> do
                lvl' <- nesting
                if (lvl'>=lvl) then loop 
                               else return Nothing