{-# Language DuplicateRecordFields,
    NoFieldSelectors,
    OverloadedRecordDot  #-}

module Parser where

  import Prelude hiding (lookup)
  import qualified Prelude as P
  import Control.Monad.State.Lazy
  import Data.Map.Lazy (Map, insert, lookup, empty, fromList, toList)
  import Data.Text.Encoding (encodeUtf8, decodeUtf8)
  import Data.Text.Lazy (toStrict)
  import Debug.Trace (trace)
  import Data.Maybe (fromJust, isNothing, isJust, fromMaybe)
  import Data.Text (pack, Text)
  import Data.YAML (decode1Strict, Node, Pos, Parser, parseEither)
  import Data.YAML.Event hiding (Scalar)
  import qualified Data.YAML.Event as Y
  import qualified Data.ByteString.Char8 as BS
  import qualified Data.ByteString.Lazy as BS.L



  activateTrace = False
  
  showTrace :: EvPos -> a -> a
  showTrace event action = if activateTrace
                           then trace (show (eEvent event) ++ " @" ++ show (posLine $ ePos event)) action
                           else action

  data ParseState = PState {
    rootNodes :: [YAMLNode],
    anchors :: Map Text YAMLNode,
    comments :: [Comment] -- YAML comment queue for the next node.
  }

  data AnchorData = NoAnchor | AnchorDef Text | AnchorAlias Text deriving (Show, Eq)


  data YAMLNode =
    Scalar {
      bytes :: BS.ByteString,
      {-tag :: Tag,-}
      {-style :: Style,-}
      comment :: [Comment], -- TODO every node preceded by a scalar preceded by a comment stores said comment  
      anchorData :: AnchorData,
      pos :: Pos
    } | Mapping {
      content :: [(YAMLNode, YAMLNode)],
      comment :: [Comment],
      anchorData :: AnchorData,
      pos :: Pos
    } | Sequence {
      children :: [YAMLNode],
      comment :: [Comment],
      anchorData :: AnchorData,
      pos :: Pos
    }


  type Comment = Text


  disp :: Int -> YAMLNode -> String
  disp indent (Scalar b c a p) = replicate indent ' ' ++ "Sca: " ++ show b ++ " #" ++ show c ++ " " ++ show a ++ " @" ++ show (posLine p) ++ "\n"
  disp indent (Mapping ct cm a p) = replicate indent ' ' ++ "Map:\n" ++ if null cm then disp' (indent + 2) ct else "" ++ " #" ++ show cm ++ " " ++ show a ++ " @" ++ show (posLine p) ++ "\n" where
    disp' :: Int -> [(YAMLNode, YAMLNode)] -> String
    disp' indent [] = ""
    disp' indent ((k,v):xs) = replicate indent ' ' ++ toStr k ++ ":\n" ++ disp (indent + 4) v ++ disp' indent xs
    toStr :: YAMLNode -> String
    toStr (Scalar b c _ _) = "| " ++ show b ++ if not (null c) then (" #" ++ show c) else ""
  disp indent (Sequence ch cm a p) = replicate indent ' ' ++ "Seq:\n" ++ disp'' (indent + 4) ch ++ " #" ++ show cm ++ " " ++ show a ++ " @" ++ show (posLine p) ++ "\n" where
    disp'' :: Int -> [YAMLNode] -> String
    disp'' indent [] = ""
    disp'' indent (x:xs) = replicate indent ' ' ++ toStr' (indent + 4) x ++ disp'' indent xs
    toStr' :: Int -> YAMLNode -> String
    toStr' _ (Scalar b c _ _) = "- " ++ show b ++ if not (null c) then (" #" ++ show c) else ""
    toStr' indent x = "- " ++ disp indent x

  instance Show YAMLNode where
    show = disp 0


  
  parse :: EvStream -> State ParseState YAMLNode
  parse ((Right (EvPos (DocumentEnd _) pos)):_) = get >>= \pState -> return $ if  length pState.rootNodes == 1
                                                                            then head pState.rootNodes
                                                                            else Sequence pState.rootNodes [] NoAnchor pos
  parse es = do
    (root, es') <- parseNode es
    pState <- get
    when (isJust root) . put $ pState {rootNodes = fromJust root : pState.rootNodes}
    parse es'

  parseComment :: EvStream -> State ParseState EvStream
  parseComment ((Right (EvPos (Y.Comment comment) _)):es) = do
    pState <- get
    put $ pState {comments = comment : pState.comments}
    parseComment es
  parseComment es = return es


  parseNode :: EvStream -> State ParseState (Maybe YAMLNode, EvStream)
  parseNode [] = error "Unexpected eof"
  parseNode ((Left _):es) = error "Failed to parse"
  parseNode es@((Right (EvPos event pos)):es') = do
    pState <- get
    showTrace (EvPos event pos) $ case event of
      Y.Comment _ -> parseComment es >>= parseNode
      Y.SequenceStart anchor _ _ -> parseSequence es' anchor [] >>= \(seq, es'') -> return (Just seq, es'')
      Y.MappingStart anchor _ _ -> parseMapping es' anchor [] >>= \(map, es'') -> return (Just map, es'')
      Y.Scalar anchor _ _ text -> parseScalar anchor text pos >>= \scal -> return (Just scal, es')
      Y.Alias anchor -> parseAlias anchor >>= \a -> return (Just a, es')
      _ -> return (Nothing, es) -- error $ "Unexpected event: " ++ show event ++ " @" ++ show (posLine pos)


  parseSequence :: EvStream -> Maybe Anchor -> [YAMLNode] -> State ParseState (YAMLNode, EvStream)
  parseSequence ((Right (EvPos SequenceEnd pos)):es) anchor children = showTrace (EvPos SequenceEnd pos) $ do
    pState <- get
    let anchorData = maybe NoAnchor AnchorDef anchor
    let seq = Sequence (reverse children) [] anchorData pos
    let anchors = if isNothing anchor then pState.anchors else insert (fromJust anchor) seq pState.anchors
    put $ pState {anchors = anchors}
    return (seq, es)
  parseSequence es anchor children = do
    (child, es') <- parseNode es
    case child of
      Nothing -> parseSequence es' anchor children
      Just c  -> parseSequence es' anchor (c : children)


  parseMapping :: EvStream -> Maybe Anchor -> [(YAMLNode, YAMLNode)] -> State ParseState (YAMLNode, EvStream)
  parseMapping ((Right (EvPos MappingEnd pos)):es) anchor content = showTrace (EvPos MappingEnd pos) $ do
    pState <- get
    let anchorData = maybe NoAnchor AnchorDef anchor
    let map = Mapping (reverse content) [] anchorData pos
    let anchors = if isNothing anchor then pState.anchors else insert (fromJust anchor) map pState.anchors
    put $ pState {anchors = anchors}
    return (map, es)
  parseMapping es anchor content = do
    (maybeKey, es') <- parseNode es
    case maybeKey of
      Nothing -> parseMapping es' anchor content
      Just key -> do
        unless (isScalar key) . error $ "Key not a scalar: " ++ show key
        (maybeVal, es'') <- parseNode es'
        let val = fromJust maybeVal
        let content' = (key {comment = []}, val {comment = key.comment}) : content -- migrate comment to val to preserve it for the workflow data structure
        parseMapping es'' anchor content' where
          isScalar :: YAMLNode -> Bool
          isScalar (Scalar {}) = True
          isScalar _ = False


  parseScalar :: Maybe Anchor -> Text -> Pos -> State ParseState YAMLNode
  parseScalar anchor text pos = do
    pState <- get
    let comments = pState.comments
    let anchorData = maybe NoAnchor AnchorDef anchor
    let scal = Scalar (encodeUtf8 text) comments anchorData pos
    let anchors = if isNothing anchor then pState.anchors else insert (fromJust anchor) scal pState.anchors
    put $ pState {anchors = anchors, comments = []}
    return scal

  
  parseAlias :: Anchor -> State ParseState YAMLNode
  parseAlias anchor = do
    pState <- get
    case lookup anchor pState.anchors of
      Nothing -> error $ "Anchor '" ++ show anchor ++ "' not defined"
      Just node -> return node {anchorData = AnchorAlias anchor}


  class FromYAML' a where
    fromYAML :: YAMLNode -> Parser a

  instance FromYAML' a => FromYAML' (Maybe a) where
    fromYAML y@(Scalar bs _ _ _)
      | decodeUtf8 bs == pack "null" = pure Nothing
      | otherwise = Just <$> fromYAML y
    fromYAML y = Just <$> fromYAML y

  instance (Ord k, FromYAML' k, FromYAML' v) => FromYAML' (Map k v) where
    fromYAML (Mapping c _ _ _) = fromList <$> mapM (\(a,b) -> (,) <$> fromYAML a <*> fromYAML b) c

  instance FromYAML' Text where
    fromYAML (Scalar bs _ _ _) = pure $ decodeUtf8 bs

  instance FromYAML' YAMLNode where
    fromYAML = pure

  instance FromYAML' v => FromYAML' [v] where
    fromYAML (Sequence c _ _ _) = mapM fromYAML c

  instance (FromYAML' a, FromYAML' b) => FromYAML' (a,b) where
    fromYAML (Sequence [n1, n2] _ _ _) = (,) <$> fromYAML n1
                                           <*> fromYAML n2

  decodeWithComments1 :: FromYAML' v => BS.L.ByteString -> Either (Pos, String) v
  decodeWithComments1 input = do
    let events = parseEvents input
    let mainEvents = validHead events
    unless (isJust mainEvents) . error $ "Missing DocumentStart event"
    let initState = PState [] empty []
    let content = evalState (parse $ fromJust mainEvents) initState
    parseEither . fromYAML $ content
    where
      validHead :: EvStream -> Maybe EvStream
      validHead ((Right (EvPos StreamStart _)):(Right (EvPos (DocumentStart _) _)):es) = Just es
      validHead _ = Nothing

  decodeWithComments1Strict :: FromYAML' v => BS.ByteString -> Either (Pos, String) v
  decodeWithComments1Strict = decodeWithComments1 . BS.L.fromChunks . (:[])

  
  (<|) :: FromYAML' a => [(YAMLNode, YAMLNode)] -> Text -> Parser a
  mapping <| key = maybe (fail $ "key " ++ show key ++ " not found") fromYAML (P.lookup key $ prep mapping) where
    prep :: [(YAMLNode, YAMLNode)] -> [(Text, YAMLNode)]
    prep mapping = [(decodeUtf8 scalar.bytes, val) | (scalar, val) <- mapping]

  (<|?) :: FromYAML' a => [(YAMLNode, YAMLNode)] -> Text -> Parser (Maybe a)
  mapping <|? key = maybe (pure Nothing) fromYAML (P.lookup key $ prep mapping) where
    prep :: [(YAMLNode, YAMLNode)] -> [(Text, YAMLNode)]
    prep mapping = [(decodeUtf8 scalar.bytes, val) | (scalar, val) <- mapping]