-----------------------------------------------------------------------------
-- |
-- Module: Text.JSON.YAJL.Enumerator
-- Copyright: 2010 John Millikin
-- License: GPL-3
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-----------------------------------------------------------------------------
module Text.JSON.YAJL.Enumerator
	( Event (..)
	
	-- * Parsing
	, parseBytesIO
	, parseTextIO
	, parseBytesST
	, parseTextST
	
	-- * Generating
	, Y.GeneratorConfig (..)
	, Y.GeneratorError (..)
	, generateBytesIO
	, generateTextIO
	, generateBytesST
	, generateTextST
	) where
import Prelude hiding (null)
import qualified Prelude as Prelude
import qualified Data.ByteString as B
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Enumerator as E
import Data.Enumerator ((>>==))
import qualified Text.JSON.YAJL as Y

import Control.Exception as Exc
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.ST (ST, stToIO, unsafeIOToST, unsafeSTToIO, RealWorld)
import qualified Data.STRef as ST
import qualified Data.IORef as IO

data Event
	= EventNull
	| EventBoolean Bool
	| EventNumber B.ByteString
	| EventText T.Text
	| EventBeginArray
	| EventEndArray
	| EventBeginObject
	| EventAttributeName T.Text
	| EventEndObject
	deriving (Show)

-- Parser {{{

parseBytesIO :: MonadIO m => E.Enumeratee B.ByteString Event m b
parseBytesIO = parseIO Y.parseBytes

parseTextIO :: MonadIO m => E.Enumeratee T.Text Event m b
parseTextIO = parseIO Y.parseText

parseBytesST :: E.Enumeratee B.ByteString Event (ST s) b
parseBytesST = parseST Y.parseBytes

parseTextST :: E.Enumeratee T.Text Event (ST s) b
parseTextST = parseST Y.parseText

newParserIO :: IO (Y.Parser IO, IO.IORef [Event])
newParserIO = do
	p <- Y.newParserIO
	eventRef <- IO.newIORef []
	let addEvent e = IO.modifyIORef eventRef (e:) >> return True
	setCallbacks p addEvent
	return (p, eventRef)

newParserST :: ST s (Y.Parser (ST s), ST.STRef s [Event])
newParserST = do
	p <- Y.newParserST
	eventRef <- ST.newSTRef []
	let addEvent e = ST.modifySTRef eventRef (e:) >> return True
	setCallbacks p addEvent
	return (p, eventRef)

setCallbacks :: Monad m => Y.Parser m -> (Event -> m Bool) -> m ()
setCallbacks p addEvent = do
	let set = Y.setCallback p
	
	set Y.parsedBeginArray (addEvent EventBeginArray)
	set Y.parsedEndArray (addEvent EventEndArray)
	set Y.parsedBeginObject (addEvent EventBeginObject)
	set Y.parsedEndObject (addEvent EventEndObject)
	set Y.parsedNull (addEvent EventNull)
	set Y.parsedBoolean (addEvent . EventBoolean)
	set Y.parsedNumber (addEvent . EventNumber)
	set Y.parsedAttributeText (addEvent . EventAttributeName)
	set Y.parsedText (addEvent . EventText)

parseIO :: MonadIO m
        => (Y.Parser IO -> a -> IO Y.ParseStatus)
        -> E.Enumeratee a Event m b
parseIO parseFn s = do
	(p, eventRef) <- liftIO newParserIO
	
	let withEvents io = liftIO $ do
		IO.writeIORef eventRef []
		status <- io
		events <- IO.readIORef eventRef
		return (reverse events, status)
	
	let parseChunk bytes = withEvents (parseFn p bytes)
	let complete = withEvents (Y.parseComplete p)
	eneeParser parseChunk complete s

parseST :: (Y.Parser (ST s) -> a -> ST s Y.ParseStatus)
        -> E.Enumeratee a Event (ST s) b
parseST parseFn s = do
	(p, eventRef) <- lift newParserST
	
	let withEvents st = do
		ST.writeSTRef eventRef []
		status <- st
		events <- ST.readSTRef eventRef
		return (reverse events, status)
	
	let parseChunk bytes = withEvents (parseFn p bytes)
	let complete = withEvents (Y.parseComplete p)
	eneeParser parseChunk complete s

eneeParser :: Monad m
           => (a -> m ([Event], Y.ParseStatus))
           -> m ([Event], Y.ParseStatus)
           -> E.Enumeratee a Event m b
eneeParser parseChunk parseComplete = E.checkDone (E.continue . step) where
	step k (E.Chunks xs) = parseLoop k xs
	step k E.EOF = checkEvents k parseComplete
		(\k' -> E.yield (E.Continue k') E.EOF)
		(\_ -> throwError (T.pack "Unexpected EOF"))
	
	parseLoop k [] = E.continue (step k)
	parseLoop k (x:xs) = checkEvents k (parseChunk x)
		(\k' -> E.yield (E.Continue k') (E.Chunks xs))
		(\k' -> parseLoop k' xs)
	
	checkEvents k getEvents onFinished onContinue = do
		(events, status) <- lift getEvents
		let checkError k' = case status of
			Y.ParseError err -> throwError err
			Y.ParseFinished -> onFinished k'
			Y.ParseContinue -> onContinue k'
			Y.ParseCancelled -> throwError (T.pack "Parse cancelled")
		if null events
			then checkError k
			else k (E.Chunks events) >>== E.checkDone checkError
	
	throwError = E.throwError . Exc.ErrorCall . T.unpack

-- }}}

-- Generator {{{

class Nullable a where
	null :: a -> Bool

instance Nullable [a] where
	null = Prelude.null

instance Nullable B.ByteString where
	null = B.null

instance Nullable T.Text where
	null = T.null

generateTextIO :: MonadIO m
               => Y.GeneratorConfig
               -> E.Enumeratee Event T.Text m b
generateTextIO = generateIO (fmap TE.decodeUtf8 . Y.getBuffer)

generateBytesIO :: MonadIO m
                => Y.GeneratorConfig
                -> E.Enumeratee Event B.ByteString m b
generateBytesIO = generateIO Y.getBuffer

generateTextST :: Y.GeneratorConfig
               -> E.Enumeratee Event T.Text (ST s) b
generateTextST = generateST (fmap TE.decodeUtf8 . Y.getBuffer)

generateBytesST :: Y.GeneratorConfig
                -> E.Enumeratee Event B.ByteString (ST s) b
generateBytesST = generateST Y.getBuffer

generateIO :: (Nullable a, MonadIO m)
           => (Y.Generator RealWorld -> ST RealWorld a)
           -> Y.GeneratorConfig
           -> E.Enumeratee Event a m b
generateIO getBuf config s = do
	g <- liftIO $ stToIO $ Y.newGenerator config
	let takeBuf = liftIO $ stToIO $ do
		buf <- getBuf g
		Y.clearBuffer g
		return buf
	let genEvent e = liftIO $ Exc.handle (return . Just) $ do
		stToIO $ genEventImpl g e
		return Nothing
	eneeGenerator genEvent takeBuf s

generateST :: Nullable a
           => (Y.Generator s -> ST s a)
           -> Y.GeneratorConfig
           -> E.Enumeratee Event a (ST s) b
generateST getBuf config s = do
	g <- lift $ Y.newGenerator config
	let takeBuf = do
		buf <- getBuf g
		Y.clearBuffer g
		return buf
	let genEvent e = unsafeIOToST $ Exc.handle (return . Just) $ do
		unsafeSTToIO $ genEventImpl g e
		return Nothing
	eneeGenerator genEvent takeBuf s

eneeGenerator :: (Nullable a, Monad m)
              => (Event -> m (Maybe Y.GeneratorError))
              -> m a
              -> E.Enumeratee Event a m b
eneeGenerator genEvent takeBuf = E.checkDone (E.continue . step) where
	step k (E.Chunks []) = E.continue (step k)
	step k (E.Chunks xs) = parseLoop k xs
	step k E.EOF = E.yield (E.Continue k) E.EOF
	
	parseLoop k [] = checkBuf k (E.continue . step)
	parseLoop k (x:xs) = do
		maybeError <- lift $ genEvent x
		case maybeError of
			Just Y.GenerationComplete -> checkBuf k
				(\k' -> E.yield (E.Continue k') (E.Chunks (x:xs)))
			Just err -> checkBuf k (\_ -> E.throwError err)
			Nothing -> parseLoop k xs
	
	checkBuf k next = do
		buf <- lift takeBuf
		if null buf
			then next k
			else k (E.Chunks [buf]) >>== E.checkDone next

genEventImpl :: Y.Generator s -> Event -> ST s ()
genEventImpl g e = case e of
	EventNull -> Y.generateNull g
	EventBoolean x -> Y.generateBoolean g x
	EventNumber num -> Y.generateNumber g num
	EventText text -> Y.generateText g text
	EventBeginArray -> Y.generateBeginArray g
	EventEndArray -> Y.generateEndArray g
	EventBeginObject -> Y.generateBeginObject g
	EventAttributeName name -> Y.generateText g name
	EventEndObject -> Y.generateEndObject g

-- }}}