{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DerivingStrategies #-}

module Language.Haskell.Brittany.Internal.Types where

import qualified Control.Monad.Trans.MultiRWS.Strict as MultiRWSS
import qualified Data.Data
import qualified Data.Kind as Kind
import qualified GHC.OldList as List
import qualified Data.Text.Lazy.Builder as Text.Builder
import qualified GHC
import           GHC                            ( Anno
                                                , DeltaPos
                                                  ( DifferentLine
                                                  , SameLine
                                                  )
                                                , EpaCommentTok
                                                , LHsDecl
                                                , ParsedSource
                                                , XRec
                                                , LImportDecl
                                                , LEpaComment
                                                )
import GHC.Utils.Outputable(Outputable)
import Language.Haskell.Brittany.Internal.Config.Types
import Language.Haskell.Brittany.Internal.Prelude
import Language.Haskell.Brittany.Internal.Components.BriDoc



-- brittany-internal error type, part of the brittany library public interface.
data BrittanyError
  = ErrorInput String
    -- ^ parsing failed
  | ErrorUnusedComment String
    -- ^ internal error: some comment went missing. Since ghc-9.2 the one below
    --   is used for missing comments while layouting top-level decls.
    --   You might still get thse for import/export statement comments (right?)
    --   (TODO this needs updating once the ghc-9.2 implementation is complete)
  | ErrorUnusedComments (LHsDecl GhcPs) Int Int
    -- ^ internal error: some comments went missing while layouting the
    --   specified top-level declaration (module child).
  | ErrorMacroConfig String String
    -- ^ in-source config string parsing error; first argument is the parser
    --   output and second the corresponding, ill-formed input.
  | LayoutWarning String
    -- ^ some warning
  | forall ast . (Data.Data.Data (XRec GhcPs ast), Outputable (Anno ast)) => ErrorUnknownNode String (XRec GhcPs ast)
    -- ^ internal error: pretty-printing is not implemented for type of node
    --   in the syntax-tree
  | ErrorOutputCheck String
    -- ^ checking the output for syntactic validity failed


-- General-purpose structure. Church-encoded version of a list that has a
-- parameterized final value.
-- That is:
-- 1) Start with a dumb old `List a = Cons a (List a) | Empty`
-- 2) Parameterize the empty `ListEnd a b = Cons a (List a b) | Final b`
-- 3) Church-encode the result for good measure (and better perf?)
--
-- This is similar to a Writer Monad, and maybe it already exists somewhere,
-- but it feels like a) This abstraction doesn't leak by itself b) you can
-- trivially turn it into something that streams into some "consuming" monadic
-- action c) it is _simple_, doesn't require a monoid or anything without
-- losing anything.
data FinalList a b = FinalList (forall r . (a -> r -> r) -> (b -> r) -> r)

finalPure :: b -> FinalList a b
finalPure b = FinalList $ \_ f -> f b
finalCons :: a -> FinalList a b -> FinalList a b
finalCons a (FinalList l) = FinalList (\f1 f2 -> f1 a $ l f1 f2)
finalBind :: FinalList a b -> (b -> FinalList a b) -> FinalList a b
finalBind (FinalList l) f =
  FinalList (\f1 f2 -> l f1 (\b -> let FinalList c = (f b) in c f1 f2))
_finalSingleton :: a -> b -> FinalList a b
_finalSingleton a b = FinalList (\f1 f2 -> f1 a (f2 b))
_finalLMap :: (a -> a') -> FinalList a b -> FinalList a' b
_finalLMap f (FinalList l) = FinalList (\f1 f2 -> l (f1 . f) f2)
_finalRMap :: (b -> b') -> FinalList a b -> FinalList a b'
_finalRMap f (FinalList l) = FinalList (\f1 f2 -> l f1 (f2 . f))
finalYield :: a -> FinalList a ()
finalYield x = FinalList $ \f1 f2 -> f1 x (f2 ())
finalToList_ :: FinalList a () -> [a]
finalToList_ (FinalList l) = l (:) (\() -> [])
finalToList :: FinalList a b -> ([a], b)
finalToList (FinalList l) = l (\x (a, b) -> (x:a, b)) (\b -> ([], b))
concatMapFinal :: FinalList a () -> (a -> [b]) -> [b]
concatMapFinal (FinalList l) f = l (\x rest -> f x ++ rest) (\() -> [])

instance Functor (FinalList a) where
  fmap = _finalRMap

instance Applicative (FinalList a) where
  pure = finalPure
  FinalList ff <*> FinalList fx =
    FinalList $ \f1 f2 -> ff f1 (\g -> fx f1 (f2 . g))

instance Monad (FinalList a) where
  FinalList fx >>= f =
    FinalList $ \f1 f2 -> fx f1 (\x -> let FinalList c = f x in c f1 f2)

-- After parsing, we split a module into a series of these so that
-- we can process them independently.
data ModuleElement
  = MEExactModuleHead ParsedSource
    -- ^ When not pretty-printing the module head:
    --   module name, imports, exports, but without decls.
    --   Could split out the relevant fields, but we need the full (with decls
    --   set to []) for exactprinting anyway.
  | MEPrettyModuleHead ParsedSource
    -- ^ Similar, but when pretty-printing the head:
    --   This encompasses just the module name, the exports and
    --   the where keyword. Everything else (including comments in
    --   various places) get 
  | MEImportDecl (LImportDecl GhcPs) [(Int, EpaCommentTok)]
    -- ^ an import decl, only occurs if pretty-printing the module head.
  | MEDecl (LHsDecl GhcPs) [(Int, EpaCommentTok)]
    -- ^ a top-level declaration
  | MEComment (Int, LEpaComment)
    -- ^ a top-level comment, i.e. a comment located between top-level elements
    --   (and not associated to some nested node, which might in theory happen).
    --   The Int carries the indentation of the comment.
  | MEWhitespace DeltaPos
  -- ^ Empty lines around decls.


newtype TraceFunc = TraceFunc { useTraceFunc :: String -> IO () }

type PPMLocal = MultiRWSS.MultiRWS
  '[Config, TraceFunc]
  '[Text.Builder.Builder, [BrittanyError], Seq String]
  '[CommentCounter]

-- type LayoutConstraints m
--   = ( MonadMultiReader Config m
--     , MonadMultiWriter TextL.Builder.Builder m
--     , MonadMultiWriter (Seq String) m
--     , MonadMultiState LayoutState m
--     , MonadMultiState CommentCounter m
--     , MonadMultiState [GHC.LEpaComment] m
--     )

ppmMoveToExactLoc
  :: MonadMultiWriter Text.Builder.Builder m => GHC.DeltaPos -> m ()
ppmMoveToExactLoc = \case
  SameLine c ->
    mTell $ Text.Builder.fromString (List.replicate c ' ')
  DifferentLine l c -> mTell $ Text.Builder.fromString
    (List.replicate l '\n' ++ List.replicate (c - 1) ' ')


type ToBriDocM = MultiRWSS.MultiRWS
  '[Config, TraceFunc, Layouters] -- reader
  '[[BrittanyError], Seq String] -- writer
  '[NodeAllocIndex, CommentCounter] -- state

data OpParenMode
  = NoParen
  | ParenNoSpace
  | ParenWithSpace
  deriving (Eq, Show)

data OpTree
  = OpUnknown OpParenMode -- Z paren?
              (Maybe GHC.RealSrcLoc) -- paren open loc
              (Maybe GHC.RealSrcLoc) -- paren close loc
              OpTree -- left operand
              [(BriDocNumbered, OpTree)] -- list of (next operand, symbol)
  | OpKnown   OpParenMode     -- with paren?
              Bool            -- allow non-whitespace transforms
              (Maybe GHC.RealSrcLoc) -- paren open loc
              (Maybe GHC.RealSrcLoc) -- paren close loc
              GHC.Fixity   -- only Just after (successful!) lookup phase
              OpTree
              [(BriDocNumbered, OpTree)]
  | OpLeaf Bool BriDocNumbered -- bool determines if parspacing is preferred

data SortItemsFlag = ShouldSortItems | KeepItemsUnsorted

data Layouters = Layouters
  { layout_expr :: ToBriDoc GHC.HsExpr
  , layout_lit :: GHC.HsLit GhcPs -> BriDocWrapped
  , layout_overLit :: GHC.OverLitVal -> BriDocWrapped
  , layout_type :: Bool -> ToBriDoc GHC.HsType
  , layout_sigType :: ToBriDoc GHC.HsSigType
  , layout_stmt :: GHC.GuardLStmt GhcPs -> ToBriDocM BriDocNumbered
  , layout_gatherOpTreeE
    :: OpParenMode
    -> Bool
    -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
    -> Maybe GHC.RealSrcLoc
    -> Maybe GHC.RealSrcLoc
    -> [(ToBriDocM BriDocNumbered, OpTree)]
    -> GHC.LHsExpr GhcPs
    -> ToBriDocM (OpTree, Bool)
  , layout_gatherOpTreeT
    :: OpParenMode
    -> Bool
    -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
    -> Maybe GHC.RealSrcLoc
    -> Maybe GHC.RealSrcLoc
    -> [(ToBriDocM BriDocNumbered, OpTree)]
    -> GHC.LHsType GhcPs
    -> ToBriDocM (OpTree, Bool)
  , layout_opTree :: (OpTree, Bool) -> ToBriDocM BriDocNumbered
  , layout_grhs
    :: GHC.LGRHS GhcPs (GHC.LHsExpr GhcPs)
    -> ToBriDocM
         ( Maybe (GHC.EpAnn GHC.GrhsAnn)
         , [BriDocNumbered]
         , BriDocNumbered
         )
  , layout_pat :: GHC.LPat GhcPs -> ToBriDocM (Seq BriDocNumbered)
  , layout_colsWrapPat :: Seq BriDocNumbered -> ToBriDocM BriDocNumbered
  , layout_patternBind
    :: Maybe Text
    -> BriDocNumbered
    -> GHC.LMatch GhcPs (GHC.LHsExpr GhcPs)
    -> ToBriDocM BriDocNumbered
  , layout_patternBindFinal
    :: Maybe Text
    -> BriDocNumbered
    -> Maybe BriDocNumbered
    -> Either (GHC.LHsExpr GhcPs) [GHC.LGRHS GhcPs (GHC.LHsExpr GhcPs)]
    -> ( ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
       , Maybe
           ( ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
           , [BriDocNumbered]
           )
       )
    -> Bool
    -> ToBriDocM BriDocNumbered
  , layout_localBinds
    :: GHC.HsLocalBindsLR GhcPs GhcPs
    -> ToBriDocM
         ( ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
         , Maybe
             ( ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
             , [BriDocNumbered]
             )
         )
  , layout_dataDecl
    :: Maybe (GHC.LTyClDecl GhcPs)
    -> GHC.EpAnn [GHC.AddEpAnn]
    -> GHC.LIdP GhcPs
    -> GHC.LHsQTyVars GhcPs
    -> [GHC.LHsTypeArg GhcPs]
    -> GHC.HsDataDefn GhcPs
    -> ToBriDocM BriDocNumbered
  , layout_tyVarBndrs
    :: [GHC.LHsTyVarBndr () GhcPs]
    -> ToBriDocM [(Text, Maybe (ToBriDocM BriDocNumbered))]
  , layout_tyVarBndrsSingleline
    :: [(Text, Maybe (ToBriDocM BriDocNumbered))] -> [ToBriDocM BriDocNumbered]
  , layout_hsTyPats
    :: [GHC.LHsTypeArg GhcPs] -> [ToBriDocM BriDocNumbered]
  , layout_LLIEs
    :: HasCallStack
    => Bool
    -> SortItemsFlag
    -> GHC.LocatedL [GHC.LIE GhcPs]
    -> ToBriDocM BriDocNumbered
  , layout_annAndSepLLIEs
    :: forall a
    .  (Data.Data.Data a, HasCallStack)
    => SortItemsFlag
    -> a
    -> [GHC.LIE GhcPs]
    -> ToBriDocM [ToBriDocM BriDocNumbered]
  }

callLayouter
  :: (Layouters -> a -> ToBriDocM r)
  -> a
  -> ToBriDocM r
callLayouter lens x = do
  layouters <- mAsk
  lens layouters x

callLayouter2
  :: (Layouters -> a -> b -> ToBriDocM r)
  -> a
  -> b
  -> ToBriDocM r
callLayouter2 lens x y = do
  layouters <- mAsk
  lens layouters x y


type ToBriDoc (sym :: Kind.Type -> Kind.Type) = XRec GhcPs (sym GhcPs) -> ToBriDocM BriDocNumbered
type ToBriDoc' sym            = XRec GhcPs (sym GhcPs)                 -> ToBriDocM BriDocNumbered
type ToBriDocC sym c          = XRec GhcPs (sym GhcPs)                 -> ToBriDocM c
type ToBriDocP sym            = sym GhcPs                              -> ToBriDocM BriDocNumbered


newtype CommentCounter = CommentCounter { unCommentCounter :: Int }
  deriving (Eq, Ord, Num, Data.Data.Data)
  deriving newtype (Show)

-- Why does this live in types? As long as it does, you _have_
-- to remember to call `docFlushRemaining` in combination with this!
briDocMToPPM :: Layouters -> ToBriDocM a -> PPMLocal (a, Int)
briDocMToPPM layouters m = do
  readers   <- MultiRWSS.mGetRawR
  initCount <- MultiRWSS.mGet @CommentCounter
  let (((x, errs), debugs), commentCount) =
        runIdentity
          $ MultiRWSS.runMultiRWSTNil
          $ MultiRWSS.withMultiStateAS initCount
          $ MultiRWSS.withMultiStateA (NodeAllocIndex 1)
          $ MultiRWSS.withMultiReader layouters
          $ MultiRWSS.withMultiReaders readers
          $ MultiRWSS.withMultiWriterAW
          $ MultiRWSS.withMultiWriterAW
          $ m
  mTell debugs
  mTell errs
  mSet commentCount
  pure
    ( x
    , sum
        [ case e of
            ErrorInput{}          -> 1
            ErrorUnusedComment{}  -> 1
            ErrorUnusedComments{} -> 1
            ErrorMacroConfig{}    -> 0
            LayoutWarning{}       -> 0
            ErrorUnknownNode{}    -> 1
            ErrorOutputCheck{}    -> 0
        | e <- errs
        ]
    )



newtype NodeAllocIndex = NodeAllocIndex Int



-- data NewLineState = NewLineStateInit -- initial state. we do not know if in a
--                                      -- newline, really. by special-casing
--                                      -- this we can appropriately handle it
--                                      -- differently at use-site.
--                   | NewLineStateYes
--                   | NewLineStateNo
--   deriving Eq

-- data LayoutSettings = LayoutSettings
--   { _lsettings_cols :: Int -- the thing that has default 80.
--   , _lsettings_indentPolicy :: IndentPolicy
--   , _lsettings_indentAmount :: Int
--   , _lsettings_indentWhereSpecial :: Bool -- indent where only 1 sometimes (TODO).
--   , _lsettings_indentListSpecial  :: Bool -- use some special indentation for ","
--                                           -- when creating zero-indentation
--                                           -- multi-line list literals.
--   , _lsettings_importColumn :: Int
--   , _lsettings_initialAnns :: ExactPrint.Anns
--   }


-- TODO92 is this old leftover, or useful future idea?
-- data BriSpacing = BriSpacing
--   { _bs_spacePastLineIndent :: Int -- space in the current,
--                                    -- potentially somewhat filled
--                                    -- line.
--   , _bs_spacePastIndent :: Int     -- space required in properly
--                                    -- indented blocks below the
--                                    -- current line.
--   }
