{-# LANGUAGE ScopedTypeVariables #-}
module Text.XML.Expat.Chunked where

import Text.XML.Expat.Chunked.Iterator
import qualified Text.XML.Expat.IO as IO
import Text.XML.Expat.SAX
import Text.XML.Expat.Tree

import Control.Monad
import Control.Monad.Trans
import qualified Data.ByteString as B
import Data.IORef
import Data.Iteratee hiding (head)
import qualified Data.Iteratee as It
import Data.Iteratee.WrappedByteString
import Data.List.Class
import Data.Sequence (Seq, (|>), index)
import qualified Data.Sequence as Seq
import Data.Word


------ Queue ------------------------------------------------------------------

newtype Queue a = Queue (Seq (Maybe a))

emptyQueue :: Queue a
emptyQueue = Queue (Seq.empty)

queuePush :: a -> Queue a -> Queue a
queuePush value (Queue s) = Queue (s |> Just value)

queueEnd :: Queue a -> Queue a
queueEnd (Queue s) = Queue (s |> Nothing)

data QueueValue a
    = Value a
    | End
    | Pending

queueIndex :: Queue a -> Int -> QueueValue a
queueIndex q@(Queue s) i
    | i >= Seq.length s = Pending
    | otherwise = case s `index` i of
        Just value -> Value value
        Nothing    -> End


------ Stack ------------------------------------------------------------------

newtype Stack a = Stack [a]

emptyStack :: Stack a
emptyStack = Stack []

stackTop :: Stack a -> a
stackTop (Stack s) = head s

stackPop :: Stack a -> Stack a
stackPop (Stack s) = Stack (tail s)

stackPush :: a -> Stack a -> Stack a
stackPush it (Stack s) = Stack (it:s)


------ Handler ----------------------------------------------------------------

data Result m a = Yield (HandlerT m a) | Result a

data HandlerT m a = HandlerT {
        runHandlerT :: m (Result m a)
    }

instance Monad m => Functor (HandlerT m) where
    fmap = liftM

instance Monad m => Monad (HandlerT m) where
    return a = HandlerT $ return $ Result a
    f >>= g = HandlerT $ do
        res1 <- runHandlerT f
        case res1 of
            Yield c -> return $ Yield (c >>= g)
            Result a  -> runHandlerT (g a)
    fail err = HandlerT $ fail err

instance MonadTrans HandlerT where
    lift m = HandlerT $ do
        r <- m
        return $ Result r

instance MonadIO m => MonadIO (HandlerT m) where
    liftIO m = HandlerT $ do
        r <- liftIO m
        return $ Result r

yield :: Monad m => HandlerT m ()
yield = HandlerT $ return $ Yield (HandlerT $ return $ Result ())

-------------------------------------------------------------------------------

type CNode m tag text = NodeG (Iterator (HandlerT m)) tag text

parse :: forall m a tag text . (
             MonadIO m,
             GenericXMLString tag,
             GenericXMLString text
         ) =>
         ParserOptions tag text
      -> (Maybe (CNode m tag text) -> HandlerT m a)
      -> IterateeG WrappedByteString Word8 m (Either ErrMsg a)
parse opts code = IterateeG $ \str -> do
    let enc = parserEncoding opts

    parser <- liftIO $ IO.newParser enc

    rootRef <- liftIO $ newIORef emptyQueue
    stackRef <- liftIO $ newIORef $ rootRef `stackPush` emptyStack

    liftIO $ IO.setStartElementHandler parser $ \cName cAttrs -> do
        name <- mkText cName
        attrs <- forM cAttrs $ \(cAttrName,cAttrValue) -> do
            attrName <- mkText cAttrName
            attrValue <- mkText cAttrValue
            return (attrName, attrValue)
        --loc <- getParseLocation parser
        stack <- readIORef stackRef
        newRef <- newIORef emptyQueue
        let elt = Element name attrs (Iterator $ iter newRef 0)
        modifyIORef (stackTop stack) (elt `queuePush`)
        modifyIORef stackRef (newRef `stackPush`)
        return True
    liftIO $ IO.setEndElementHandler parser $ \_ -> do
        --name <- mkText cName
        --loc <- getParseLocation parser
        stack <- readIORef stackRef
        modifyIORef (stackTop stack) queueEnd
        writeIORef stackRef (stackPop stack)
        return True
    liftIO $ IO.setCharacterDataHandler parser $ \cText -> do
        txt <- gxFromCStringLen cText
        --loc <- getParseLocation parser
        stack <- readIORef stackRef
        modifyIORef (stackTop stack) (Text txt `queuePush`)
        return True

    let nextIter :: HandlerT m a
                 -> IterateeG WrappedByteString Word8 m (Either ErrMsg a)
        nextIter c = IterateeG $ \str -> do
            case str of
                EOF (Just err) -> return $ Done (Left err) str
                _ -> do
                    mErr <- case str of
                        Chunk (WrapBS blk) -> liftIO $ IO.parseChunk parser blk False
                        EOF _              -> liftIO $ IO.parseChunk parser B.empty True

                    res <- runHandlerT c

                    return $ case res of
                        Yield c' -> do
                            case (str, mErr) of
                                (Chunk _, Just err) -> Cont (nextIter c') (Just $ Err $ show err)
                                (Chunk _, Nothing)  -> Cont (nextIter c') Nothing
                                (EOF _,   Just err) -> Done (Left $ Err $ show err) str
                                (EOF _,   Nothing)  -> Done (Left $ Err "EOF not handled") str
                        Result a -> do
                            case mErr of
                                Just err -> Done (Left $ Err $ show err) str
                                Nothing  -> Done (Right a) str

    let process = do
            elt <- iter rootRef 0
            case elt of
                Nil         -> code Nothing
                Cons node _ -> code (Just node)
    runIter (nextIter process) str

  where
    iter :: IORef (Queue (CNode m tag text)) -> Int -> HandlerT m (ListItem (Iterator (HandlerT m)) (CNode m tag text))
    iter ref i = do
        queue <- liftIO (readIORef ref)
        case queue `queueIndex` i of  -- to do: make it throw away heads we don't have references to
            Pending -> yield >> iter ref i
            End     -> return $ Nil
            Value a -> return $
                let i' = i + 1
                in  Cons a (Iterator $ iter ref i')