{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE RecordWildCards            #-}
{-# LANGUAGE RecursiveDo                #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE Trustworthy                #-}

-- |
-- Copyright: © Herbert Valerio Riedel 2015-2018
-- SPDX-License-Identifier: GPL-2.0-or-later
--
module Data.YAML.Loader
    ( decodeLoader
    , Loader(..)
    , LoaderT
    , NodeId
    ) where

import           Control.Monad.State  (MonadState(..), gets, modify,
                                       StateT, evalStateT, state)
import           Control.Monad.Trans  (MonadTrans(..))
import qualified Data.ByteString.Lazy as BS.L
import qualified Data.Map             as Map
import qualified Data.Set             as Set

import           Data.YAML.Event      (Tag)
import qualified Data.YAML.Event      as YE
import           Util

-- | Unique identifier for identifying nodes
--
-- This is allows to observe the alias/anchor-reference structure
type NodeId = Word

-- | Structure defining how to construct a document tree/graph
--
-- @since 0.2.0
--
data Loader m n = Loader
  { yScalar   :: Tag -> YE.ScalarStyle -> Text -> LoaderT m n
  , ySequence :: Tag -> [n]                    -> LoaderT m n
  , yMapping  :: Tag -> [(n,n)]                -> LoaderT m n
  , yAlias    :: NodeId -> Bool -> n           -> LoaderT m n
  , yAnchor   :: NodeId -> n                   -> LoaderT m n
  }

-- | Helper type for 'Loader'
--
-- @since 0.2.0
type LoaderT m n = YE.Pos -> m (Either (YE.Pos,String) n)

-- TODO: newtype LoaderT m n = LoaderT { runLoaderT :: YE.Pos -> m (Either String n) }

-- | Generalised document tree/graph construction
--
-- This doesn't yet perform any tag resolution (thus all scalars are
-- represented as 'Text' values). See also 'Data.YAML.decodeNode' for a more
-- convenient interface.
--
-- @since 0.2.0
{-# INLINEABLE decodeLoader #-}
decodeLoader :: forall n m . MonadFix m => Loader m n -> BS.L.ByteString -> m (Either (YE.Pos, String) [n])
decodeLoader Loader{..} bs0 = do
    case sequence $ filter (not. isComment) (YE.parseEvents bs0) of
      Left (pos,err) -> return $ Left (pos,err)
      Right evs      -> runParserT goStream evs
  where
    isComment evPos = case evPos of
      Right (YE.EvPos {eEvent = (YE.Comment _), ePos = _}) -> True
      _                                                    -> False

    goStream :: PT n m [n]
    goStream = do
      _ <- satisfy (== YE.StreamStart)
      ds <- manyUnless (== YE.StreamEnd) goDoc
      eof
      return ds

    goDoc :: PT n m n
    goDoc = do
      _ <- satisfy isDocStart
      modify $ \s0 -> s0 { sDict = mempty, sCycle = mempty }
      n <- goNode
      _ <- satisfy isDocEnd
      return n

    getNewNid :: PT n m Word
    getNewNid = state $ \s0 -> let i0 = sIdCnt s0
                               in (i0, s0 { sIdCnt = i0+1 })

    returnNode :: YE.Pos -> Maybe YE.Anchor -> Either (YE.Pos, String) n -> PT n m n
    returnNode _ _ (Left err) = throwError err
    returnNode _ Nothing (Right node) = return node
    returnNode pos (Just a) (Right node) = do
      nid <- getNewNid
      node' <- liftEither' =<< lift (yAnchor nid node pos)
      modify $ \s0 -> s0 { sDict = Map.insert a (nid,node') (sDict s0) }
      return node'

    registerAnchor :: YE.Pos -> Maybe YE.Anchor -> PT n m n -> PT n m n
    registerAnchor _ Nothing  pn = pn
    registerAnchor pos (Just a) pn = do
      modify $ \s0 -> s0 { sCycle = Set.insert a (sCycle s0) }
      nid <- getNewNid

      mdo
        modify $ \s0 -> s0 { sDict = Map.insert a (nid,n) (sDict s0) }
        n0 <- pn
        n  <- liftEither' =<< lift (yAnchor nid n0 pos)
        return n

    exitAnchor :: Maybe YE.Anchor -> PT n m ()
    exitAnchor Nothing  = return ()
    exitAnchor (Just a) = modify $ \s0 -> s0 { sCycle = Set.delete a (sCycle s0) }

    goNode :: PT n m n
    goNode = do
      n <- anyEv
      let pos = YE.ePos n
      case YE.eEvent n of
        YE.Scalar manc tag sty val -> do
          exitAnchor manc
          n' <- lift (yScalar tag sty val pos)
          returnNode pos manc $! n'

        YE.SequenceStart manc tag _ -> registerAnchor pos manc $ do
          ns <- manyUnless (== YE.SequenceEnd) goNode
          exitAnchor manc
          liftEither' =<< lift (ySequence tag ns pos)

        YE.MappingStart manc tag _ -> registerAnchor pos manc $ do
          kvs <- manyUnless (== YE.MappingEnd) (liftM2 (,) goNode goNode)
          exitAnchor manc
          liftEither' =<< lift (yMapping tag kvs pos)

        YE.Alias a -> do
          d <- gets sDict
          cy <- gets sCycle
          case Map.lookup a d of
            Nothing       -> throwError (pos, ("anchor not found: " ++ show a))
            Just (nid,n') -> liftEither' =<< lift (yAlias nid (Set.member a cy) n' pos)

        _ -> throwError (pos, "goNode: unexpected event")

----------------------------------------------------------------------------
-- small parser framework


data S n = S { sEvs   :: [YE.EvPos]
             , sDict  :: Map YE.Anchor (Word,n)
             , sCycle :: Set YE.Anchor
             , sIdCnt :: !Word
             }

newtype PT n m a = PT (StateT (S n) (ExceptT (YE.Pos, String) m) a)
                 deriving ( Functor
                          , Applicative
                          , Monad
                          , MonadState (S n)
                          , MonadError (YE.Pos, String)
                          , MonadFix
                          )

instance MonadTrans (PT n) where
  lift = PT . lift . lift

runParserT :: Monad m => PT n m a -> [YE.EvPos] -> m (Either (YE.Pos, String) a)
runParserT (PT act) s0 = runExceptT $ evalStateT act (S s0 mempty mempty 0)

satisfy :: Monad m => (YE.Event -> Bool) -> PT n m YE.EvPos
satisfy p = do
  s0 <- get
  case sEvs s0 of
    [] -> throwError (fakePos, "satisfy: premature eof")
    (ev:rest)
       | p (YE.eEvent ev) -> do put (s0 { sEvs = rest})
                                return ev
       | otherwise        -> throwError (YE.ePos ev, ("satisfy: predicate failed " ++ show ev))

peek :: Monad m => PT n m (Maybe YE.EvPos)
peek = do
  s0 <- get
  case sEvs s0 of
    []     -> return Nothing
    (ev:_) -> return (Just ev)

peek1 :: Monad m => PT n m YE.EvPos
peek1 = maybe (throwError (fakePos,"peek1: premature eof")) return =<< peek

anyEv :: Monad m => PT n m YE.EvPos
anyEv = satisfy (const True)

eof :: Monad m => PT n m ()
eof = do
  s0 <- get
  case sEvs s0 of
    []     -> return ()
    (ev:_) -> throwError (YE.ePos ev, "eof expected")

-- NB: consumes the end-event
manyUnless :: Monad m => (YE.Event -> Bool) -> PT n m a -> PT n m [a]
manyUnless p act = do
  t0 <- peek1
  if p (YE.eEvent t0)
    then anyEv >> return []
    else liftM2 (:) act (manyUnless p act)

{-
tryError :: MonadError e m => m a -> m (Either e a)
tryError act = catchError (Right <$> act) (pure . Left)
-}

isDocStart :: YE.Event -> Bool
isDocStart (YE.DocumentStart _) = True
isDocStart _                    = False

isDocEnd :: YE.Event -> Bool
isDocEnd (YE.DocumentEnd _) = True
isDocEnd _                  = False

fakePos :: YE.Pos
fakePos = YE.Pos { posByteOffset = -1 , posCharOffset = -1  , posLine = 1 , posColumn = 0 }