content: module Futhark.Fmt.Monad
    ( Fmt,
      -- functions for building fmt
      nil,
      nest,
      stdNest,
      text,
      space,
      hardline,
      line,
      sep,
      brackets,
      braces,
      parens,
      (<|>),
      (<+>),
      (</>),
      (<:/>),
      hardIndent,
      indent,
      hardStdIndent,
      stdIndent,
      FmtM,
      popComments,
      runFormat,
      align,
      fmtCopyLoc,
      comment,
      sepArgs,
      localLayout,
      localLayoutList,
      sepDecs,
      fmtByLayout,
      addComments,
      sepComments,
      sepLineComments,
      sepLine,
  
      -- * Formatting styles
      commentStyle,
      constantStyle,
      keywordStyle,
      bindingStyle,
      infixStyle,
    )
  where
  
  import Control.Monad (liftM2)
  import Control.Monad.Reader
    ( MonadReader (..),
      ReaderT (..),
    )
  import Control.Monad.State
    ( MonadState (..),
      State,
      evalState,
      gets,
      modify,
    )
  import Data.ByteString qualified as BS
  import Data.List.NonEmpty qualified as NE
  import Data.Loc (Loc (..), Located (..), locStart, posCoff, posLine)
  import Data.Maybe (fromMaybe)
  import Data.String
  import Data.Text qualified as T
  import Data.Text.Encoding qualified as T
  import Language.Futhark.Parser.Monad (Comment (..))
  import Prettyprinter qualified as P
  import Prettyprinter.Render.Terminal
    ( AnsiStyle,
      Color (..),
      bold,
      color,
      colorDull,
      italicized,
    )
  
  -- These are right associative since we want to evaluate the monadic
  -- computation from left to right. Since the left most expression is
  -- printed first and our monad is checking if a comment should be
  -- printed.
  
  infixr 6 <:/>
  
  infixr 6 <+>
  
  infixr 6 </>
  
  infixr 4 <|>
  
  type Fmt = FmtM (P.Doc AnsiStyle)
  
  instance Semigroup Fmt where
    (<>) = liftM2 (<>)
  
  instance Monoid Fmt where
    mempty = nil
  
  instance IsString Fmt where
    fromString s = text style s'
      where
        s' = fromString s
        style =
          if s' `elem` keywords
            then keywordStyle
            else mempty
        keywords =
          [ "true",
            "false",
            "if",
            "then",
            "else",
            "def",
            "let",
            "loop",
            "in",
            "val",
            "for",
            "do",
            "with",
            "local",
            "open",
            "include",
            "import",
            "type",
            "entry",
            "module",
            "while",
            "assert",
            "match",
            "case"
          ]
  
  commentStyle, keywordStyle, constantStyle, bindingStyle, infixStyle :: AnsiStyle
  commentStyle = italicized
  keywordStyle = color Magenta <> bold
  constantStyle = color Green
  bindingStyle = colorDull Blue
  infixStyle = colorDull Cyan
  
  -- | This function allows to inspect the layout of an expression @a@ and if it
  -- is singleline line then use format @s@ and if it is multiline format @m@.
  fmtByLayout ::
    (Located a) => a -> Fmt -> Fmt -> Fmt
  fmtByLayout a s m =
    s
      <|> ( case lineLayout a of
              Just SingleLine -> s
              _any -> m
          )
  
  -- | This function determines the Layout of @a@ and updates the monads
  -- environment to format in the appropriate style. It determines this
  -- by checking if the location of @a@ spans over two or more lines.
  localLayout :: (Located a) => a -> FmtM b -> FmtM b
  localLayout a = local (\lo -> fromMaybe lo $ lineLayout a)
  
  -- | This function determines the Layout of @[a]@ and if it is singleline then it
  -- updates the monads enviroment to format singleline style otherwise format using
  -- multiline style. It determines this by checking if the locations of @[a]@
  -- start and end at any different line number.
  localLayoutList :: (Located a) => [a] -> FmtM b -> FmtM b
  localLayoutList a m = do
    lo <- ask
    case lo of
      MultiLine -> local (const $ fromMaybe lo $ lineLayoutList a) m
      SingleLine -> m
  
  -- | This function uses the location of @a@ and prepends comments if
  -- the comments location is less than the location of @a@. It format
  -- @b@ in accordance with if @a@ is singleline or multiline using
  -- 'localLayout'. It currently does not handle trailing comment
  -- perfectly. See tests/fmt/traillingComments*.fut.
  addComments :: (Located a) => a -> Fmt -> Fmt
  addComments a b = localLayout a $ do
    c <- fmtComments a
    f <- b
    pure $ c <> f
  
  prependComments :: (a -> Loc) -> (a -> Fmt) -> a -> Fmt
  prependComments floc fmt a = do
    fmcs <- fcs
    f <- fmt a
    pure $ fromMaybe mempty fmcs <> f
    where
      fcs = do
        s <- get
        case comments s of
          c : cs | floc a /= NoLoc && floc a > locOf c -> do
            put $ s {comments = cs}
            mcs <- fcs
            pre' <- pre
            pure $ Just $ pre' <> fmtNoLine c <> maybe mempty (P.line <>) mcs
          _any -> pure Nothing
      fmtNoLine = P.pretty . commentText
      pre = do
        lastO <- gets lastOutput
        case lastO of
          Nothing -> nil
          Just Line -> nil
          Just _ -> modify (\s -> s {lastOutput = Just Line}) >> hardline
  
  -- | The internal state of the formatter monad 'FmtM'.
  data FmtState = FmtState
    { -- | The comments that will be inserted, ordered by increasing order in regards to location.
      comments :: [Comment],
      -- | The original source file that is being formatted.
      file :: BS.ByteString,
      -- | Keeps track of what type the last output was.
      lastOutput :: !(Maybe LastOutput)
    }
    deriving (Show, Eq, Ord)
  
  -- | A data type to describe the last output used during formatting.
  data LastOutput = Line | Space | Text | Comm deriving (Show, Eq, Ord)
  
  -- | A data type to describe the layout the formatter is using currently.
  data Layout = MultiLine | SingleLine deriving (Show, Eq)
  
  -- | The format monad used to keep track of comments and layout. It is a a
  -- combincation of a reader and state monad. The comments and reading from the
  -- input file are the state monads job to deal with. While the reader monad
  -- deals with the propagating the current layout.
  type FmtM a = ReaderT Layout (State FmtState) a
  
  fmtComment :: Comment -> Fmt
  fmtComment c = comment $ commentText c
  
  fmtCommentList :: [Comment] -> Fmt
  fmtCommentList [] = nil
  fmtCommentList (c : cs) =
    fst $ foldl f (fmtComment c, locOf c) cs
    where
      f (acc, loc) c' =
        if consecutive loc (locOf c')
          then (acc <> fmtComment c', locOf c')
          else (acc <> hardline <> fmtComment c', locOf c')
  
  hasComment :: (Located a) => a -> FmtM Bool
  hasComment a =
    gets $ not . null . takeWhile relevant . comments
    where
      relevant c = locOf a /= NoLoc && locOf a > locOf c
  
  -- | Prepends comments.
  fmtComments :: (Located a) => a -> Fmt
  fmtComments a = do
    (here, later) <- gets $ span relevant . comments
    if null here
      then pure mempty
      else do
        modify $ \s -> s {comments = later}
        fmtCommentList here
          <> if consecutive (locOf here) (locOf a) then nil else hardline
    where
      relevant c = locOf a /= NoLoc && locOf a > locOf c
  
  -- | Determines the layout of @a@ by checking if it spans a single line or two
  -- or more lines.
  lineLayout :: (Located a) => a -> Maybe Layout
  lineLayout a =
    case locOf a of
      Loc start end ->
        if posLine start == posLine end
          then Just SingleLine
          else Just MultiLine
      NoLoc -> Nothing -- error "Formatting term without location."
  
  -- | Determines the layout of @[a]@ by checking if it spans a single line or two
  -- or more lines.
  lineLayoutList :: (Located a) => [a] -> Maybe Layout
  lineLayoutList as =
    case concatMap auxiliary as of
      [] -> Nothing
      (t : ts) | any (/= t) ts -> Just MultiLine
      _ -> Just SingleLine
    where
      auxiliary a =
        case locOf a of
          Loc start end -> [posLine start, posLine end]
          NoLoc -> [] -- error "Formatting term without location"
  
  -- | Retrieves the last comments from the monad and concatenates them together.
  popComments :: Fmt
  popComments = do
    cs <- gets comments
    modify (\s -> s {comments = []})
    lastO <- gets lastOutput
    case lastO of
      Nothing ->
        fmtCommentList cs -- Happens when file has only comments.
      _
        | not $ null cs -> hardline <> fmtCommentList cs
        | otherwise -> nil
  
  -- | Using the location of @a@ get the segment of text in the original file to
  -- create a @Fmt@.
  fmtCopyLoc :: (Located a) => AnsiStyle -> a -> Fmt
  fmtCopyLoc style a = do
    f <- gets file
    case locOf a of
      Loc sPos ePos ->
        let sOff = posCoff sPos
            eOff = posCoff ePos
         in case T.decodeUtf8' $ BS.take (eOff - sOff) $ BS.drop sOff f of
              Left err -> error $ show err
              Right lit -> text style lit
      NoLoc -> error "Formatting term without location"
  
  -- | Given a formatter @FmtM a@, a sequence of comments ordered in increasing
  -- order by location, and the original text files content. Run the formatter and
  -- create @a@.
  runFormat :: FmtM a -> [Comment] -> T.Text -> a
  runFormat format cs file = evalState (runReaderT format e) s
    where
      s =
        FmtState
          { comments = cs,
            file = T.encodeUtf8 file,
            lastOutput = Nothing
          }
      e = MultiLine
  
  -- | An empty input.
  nil :: Fmt
  nil = pure mempty
  
  -- | Indents everything after a line occurs if in multiline and if in singleline
  -- then indent.
  nest :: Int -> Fmt -> Fmt
  nest i a = a <|> (P.nest i <$> a)
  
  -- | A space.
  space :: Fmt
  space = modify (\s -> s {lastOutput = Just Space}) >> pure P.space
  
  -- | Forces a line to be used regardless of layout, this should
  -- ideally not be used.
  hardline :: Fmt
  hardline = do
    modify $ \s -> s {lastOutput = Just Line}
    pure P.line
  
  -- | A line or a space depending on layout.
  line :: Fmt
  line = space <|> hardline
  
  -- | Seperates element by a @s@ followed by a space in singleline layout and
  -- seperates by a line followed by a @s@ in multine layout.
  sepLine :: Fmt -> [Fmt] -> Fmt
  sepLine s = sep (s <> space <|> hardline <> s)
  
  -- | A comment.
  comment :: T.Text -> Fmt
  comment c = do
    modify (\s -> s {lastOutput = Just Line})
    pure $ P.annotate commentStyle (P.pretty (T.stripEnd c)) <> P.line
  
  sep :: Fmt -> [Fmt] -> Fmt
  sep _ [] = nil
  sep s (a : as) = auxiliary a as
    where
      auxiliary acc [] = acc
      auxiliary acc (x : xs) = auxiliary (acc <> s <> x) xs
  
  sepComments :: (a -> Loc) -> (a -> Fmt) -> Fmt -> [a] -> Fmt
  sepComments _ _ _ [] = nil
  sepComments floc fmt s (a : as) = auxiliary (fmt a) as
    where
      auxiliary acc [] = acc
      auxiliary acc (x : xs) =
        auxiliary (acc <> prependComments floc (\y -> s <> fmt y) x) xs
  
  sepLineComments :: (a -> Loc) -> (a -> Fmt) -> Fmt -> [a] -> Fmt
  sepLineComments floc fmt s =
    sepComments floc fmt (s <> space <|> hardline <> s)
  
  -- | This is used for function arguments. It seperates multiline
  -- arguments by lines and singleline arguments by spaces. We specially
  -- handle the case where all the arguments are on a single line except
  -- for the last one, which may continue to the next line.
  sepArgs :: (Located a) => (a -> Fmt) -> NE.NonEmpty a -> Fmt
  sepArgs fmt ls =
    localLayout locs $ align' $ sep line $ map fmtArg ls'
    where
      locs = map (locStart . locOf) ls'
      align' = case lineLayout locs of
        Just SingleLine -> id
        _ -> align
      fmtArg x = localLayout x $ fmt x
      ls' = NE.toList ls
  
  -- | Nest but with the standard value of two spaces.
  stdNest :: Fmt -> Fmt
  stdNest = nest 2
  
  -- | Aligns line by line.
  align :: Fmt -> Fmt
  align a = do
    modify (\s -> s {lastOutput = Just Line}) -- XXX?
    P.align <$> a
  
  -- | Indents everything by @i@, should never be used.
  hardIndent :: Int -> Fmt -> Fmt
  hardIndent i a = P.indent i <$> a
  
  -- | Indents if in multiline by @i@ if in singleline it does not indent.
  indent :: Int -> Fmt -> Fmt
  indent i a = a <|> hardIndent i a
  
  -- | Hard indents with the standard size of two.
  hardStdIndent :: Fmt -> Fmt
  hardStdIndent = hardIndent 2
  
  -- | Idents with the standard size of two.
  stdIndent :: Fmt -> Fmt
  stdIndent = indent 2
  
  -- | Creates a piece of text, it should not contain any new lines.
  text :: AnsiStyle -> T.Text -> Fmt
  text style t = do
    modify (\s -> s {lastOutput = Just Text})
    pure $ P.annotate style $ P.pretty t
  
  -- | Adds brackets.
  brackets :: Fmt -> Fmt
  brackets a = "[" <> a <> "]"
  
  -- | Adds braces.
  braces :: Fmt -> Fmt
  braces a = "{" <> a <> "}"
  
  -- | Add parenthesis.
  parens :: Fmt -> Fmt
  parens a = "(" <> a <> ")"
  
  -- | If in a singleline layout then concatenate with 'nil' and in multiline
  -- concatenate by a line.
  (<:/>) :: Fmt -> Fmt -> Fmt
  a <:/> b = a <> (nil <|> hardline) <> b
  
  -- | Concatenate with a space between.
  (<+>) :: Fmt -> Fmt -> Fmt
  a <+> b = a <> space <> b
  
  -- | Concatenate with a space if in singleline layout and concatenate by a
  -- line in multiline.
  (</>) :: Fmt -> Fmt -> Fmt
  a </> b = a <> line <> b
  
  -- | If in a singleline layout then choose @a@, if in a multiline layout choose
  -- @b@.
  (<|>) :: Fmt -> Fmt -> Fmt
  a <|> b = do
    lo <- ask
    if lo == SingleLine
      then a
      else b
  
  -- | Are these locations on consecutive lines?
  consecutive :: Loc -> Loc -> Bool
  consecutive (Loc _ end) (Loc beg _) = posLine end + 1 == posLine beg
  consecutive _ _ = False
  
  -- | If in singleline layout seperate by spaces. In a multiline layout seperate
  -- by a single line if two neighbouring elements are singleline. Otherwise
  -- sepereate by two lines.
  sepDecs :: (Located a) => (a -> Fmt) -> [a] -> Fmt
  sepDecs _ [] = nil
  sepDecs fmt decs@(x : xs) =
    sep space (map fmt decs) <|> (fmt x <> auxiliary x xs)
    where
      auxiliary _ [] = nil
      auxiliary prev (y : ys) = p <> fmt y <> auxiliary y ys
        where
          p = do
            commented <- hasComment y
            case (commented, lineLayout y, lineLayout prev) of
              (False, Just SingleLine, Just SingleLine)
                | consecutive (locOf prev) (locOf y) -> hardline
              _any -> hardline <> hardline
  