{-# LANGUAGE NoImplicitPrelude #-}

module Language.Haskell.Brittany.Internal.ToBriDoc.DataDecl where

import qualified Data.Text as Text
import GHC (GenLocated(L))
import GHC.Hs
import qualified GHC.OldList as List
import Language.Haskell.Brittany.Internal.Config.Types
import Language.Haskell.Brittany.Internal.S3_ToBriDocTools
import Language.Haskell.Brittany.Internal.ToBriDoc.Type
import Language.Haskell.Brittany.Internal.Prelude
import Language.Haskell.Brittany.Internal.Types
import Language.Haskell.Brittany.Internal.Components.BriDoc



layoutDataDecl
  :: LTyClDecl GhcPs
  -> EpAnn [AddEpAnn]
  -> LIdP GhcPs
  -> LHsQTyVars GhcPs
  -> [LHsTypeArg GhcPs]
  -> HsDataDefn GhcPs
  -> ToBriDocM BriDocNumbered
layoutDataDecl ltycl declEpAnn name (HsQTvs _ bndrs) pats defn = case defn of
  -- newtype MyType a b = MyType ..
  HsDataDefn NoExtField NewType Nothing _ctype Nothing [cons] mDerivs ->
    case cons of
      (L _ (ConDeclH98 epAnn consName False _qvars ctxMay details _conDoc)) ->
        let isSimple = case ctxMay of
              Nothing -> True
              Just (L _ []) -> True
              _ -> False
        in if isSimple
          then do
            nameStr <- lrdrNameToTextAnn name
            consNameStr <- lrdrNameToTextAnn consName
            tyVarLine <- shareDoc $ createBndrDoc bndrs
            patDocs <- mapM shareDoc $ layoutHsTyPats pats
            -- headDoc     <- fmap return $ docSeq
            --   [ appSep $ docLitS "newtype")
            --   , appSep $ docLit nameStr
            --   , appSep tyVarLine
            --   ]
            rhsDoc <- return <$> createDetailsDoc consNameStr details
            docHandleComms declEpAnn $ createDerivingPar mDerivs $ docSeq
              [ appSep $ docLitS "newtype"
              , appSep $ docLit nameStr
              , appSep tyVarLine
              , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
              , docSeparator
              , docLitS "="
              , docSeparator
              , docHandleComms epAnn $ rhsDoc
              ]
          else briDocByExactNoComment ltycl
      _ -> briDocByExactNoComment ltycl


  -- data MyData a b
  -- (zero constructors)
  HsDataDefn NoExtField DataType ctxMay _ctype Nothing [] mDerivs -> do
    lhsContextDoc <- case ctxMay of
      Just (L _ lhsContext) -> shareDoc $ createContextDoc lhsContext
      Nothing -> pure docEmpty
    nameStr <- lrdrNameToTextAnn name
    tyVarLine <- return <$> createBndrDoc bndrs
    patDocs <- mapM shareDoc $ layoutHsTyPats pats
    docHandleComms declEpAnn $ createDerivingPar mDerivs $ docSeq
      [ appSep $ docLitS "data"
      , lhsContextDoc
      , appSep $ docLit nameStr
      , appSep tyVarLine
      , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
      ]

  -- data MyData = MyData ..
  -- data MyData = MyData { .. }
  HsDataDefn NoExtField DataType ctxMay _ctype Nothing [cons] mDerivs ->
    case cons of
      (L _ (ConDeclH98 epAnn consName _hasExt qvars mRhsContext details _conDoc))
        -> do
          lhsContextDoc <- case ctxMay of
            Just (L _ lhsContext) -> shareDoc $ createContextDoc lhsContext
            Nothing -> pure docEmpty
          nameStr <- lrdrNameToTextAnn name
          consNameStr <- lrdrNameToTextAnn consName
          tyVarLine <- return <$> createBndrDoc bndrs
          patDocs <- mapM shareDoc $ layoutHsTyPats pats
          forallDocMay <- case createForallDoc qvars of
            Nothing -> pure Nothing
            Just x -> Just . pure <$> x
          rhsContextDocMay <- case mRhsContext of
            Nothing -> pure Nothing
            Just (L _ ctxt) -> Just . pure <$> createContextDoc ctxt
          rhsDoc <- return <$> createDetailsDoc consNameStr details
          let posEqual = obtainAnnPos epAnn AnnEqual
          consDoc <-
            shareDoc
            $ docHandleComms epAnn
            $ docHandleComms posEqual
            $ docNonBottomSpacing
            $ case (forallDocMay, rhsContextDocMay) of
                (Just forallDoc, Just rhsContextDoc) -> docLines
                  [ docSeq
                    [docLitS "=", docSeparator, docForceSingleline forallDoc]
                  , docSeq
                    [ docLitS "."
                    , docSeparator
                    , docSetBaseY $ docLines [rhsContextDoc, docSetBaseY rhsDoc]
                    ]
                  ]
                (Just forallDoc, Nothing) -> docLines
                  [ docSeq
                    [docLitS "=", docSeparator, docForceSingleline forallDoc]
                  , docSeq [docLitS ".", docSeparator, rhsDoc]
                  ]
                (Nothing, Just rhsContextDoc) -> docSeq
                  [ docLitS "="
                  , docSeparator
                  , docSetBaseY $ docLines [rhsContextDoc, docSetBaseY rhsDoc]
                  ]
                (Nothing, Nothing) ->
                  docSeq [docLitS "=", docSeparator, rhsDoc]
          docHandleComms declEpAnn $ createDerivingPar mDerivs $ docAlt
            [ -- data D = forall a . Show a => D a
              docSeq
              [ -- TOOD92 docNodeAnnKW ltycl (Just GHC.AnnData) $ 
                docSeq
                [ appSep $ docLitS "data"
                , docForceSingleline $ lhsContextDoc
                , appSep $ docLit nameStr
                , appSep tyVarLine
                , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                ]
              , docHandleComms epAnn $ docHandleComms posEqual $ docLitS "="
              , docSeparator
              , docSetIndentLevel $ docSeq
                [ case forallDocMay of
                  Nothing -> docEmpty
                  Just forallDoc ->
                    docSeq
                      [ docForceSingleline forallDoc
                      , docSeparator
                      , docLitS "."
                      , docSeparator
                      ]
                , maybe docEmpty docForceSingleline rhsContextDocMay
                , rhsDoc
                ]
              ]
            , -- data D
              --   = forall a . Show a => D a
              docAddBaseY BrIndentRegular $ docPar
              ( -- TODO92 docNodeAnnKW ltycl (Just GHC.AnnData) $ 
                docSeq
                [ appSep $ docLitS "data"
                , docForceSingleline lhsContextDoc
                , appSep $ docLit nameStr
                , tyVarLine
                , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                ]
              )
              (docSeq
                [ docHandleComms epAnn $ docHandleComms posEqual $ docLitS "="
                , docSeparator
                , docSetIndentLevel $ docSeq
                  [ case forallDocMay of
                    Nothing -> docEmpty
                    Just forallDoc ->
                      docSeq
                        [ docForceSingleline forallDoc
                        , docSeparator
                        , docLitS "."
                        , docSeparator
                        ]
                  , maybe docEmpty docForceSingleline rhsContextDocMay
                  , rhsDoc
                  ]
                ]
              )
            , -- data D
              --   = forall a
              --   . Show a =>
              --     D a
              docAddBaseY BrIndentRegular $ docPar
              ( -- TODO92 docNodeAnnKW ltycl (Just GHC.AnnData) $ 
                docSeq
                [ appSep $ docLitS "data"
                , docForceSingleline lhsContextDoc
                , appSep $ docLit nameStr
                , tyVarLine
                , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                ]
              )
              consDoc
            , -- data
              --   Show a =>
              --   D
              --   = forall a
              --   . Show a =>
              --     D a
              -- This alternative is only for -XDatatypeContexts.
              -- But I think it is rather unlikely this will trigger without
              -- -XDataTypeContexts, especially with the `docNonBottomSpacing`
              -- above, so while not strictly necessary, this should not
              -- hurt.
              docAddBaseY BrIndentRegular $ docPar
              (docLitS "data")
              (docLines
                [ lhsContextDoc
                , -- TODO92 docNodeAnnKW ltycl (Just GHC.AnnData) $
                  docSeq [appSep $ docLit nameStr, tyVarLine, docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]]
                , consDoc
                ]
              )
            ]
      _ -> briDocByExactNoComment ltycl

  _ -> briDocByExactNoComment ltycl

layoutHsTyPats
  :: [LHsTypeArg GhcPs] -> [ToBriDocM BriDocNumbered]
layoutHsTyPats pats = pats <&> \case
  HsValArg tm -> layoutType tm
  HsTypeArg _l ty -> docSeq [docLit $ Text.pack "@", layoutType ty]
    -- we ignore the SourceLoc here.. this LPat not being (L _ Pat{}) change
    -- is a bit strange. Hopefully this does not ignore any important
    -- annotations.
  HsArgPar _l -> error "brittany internal error: HsArgPar{}"

createContextDoc :: HsContext GhcPs -> ToBriDocM BriDocNumbered
createContextDoc [] = docEmpty
createContextDoc [t] =
  docSeq [layoutType t, docSeparator, docLitS "=>", docSeparator]
createContextDoc (t1 : tR) = do
  t1Doc <- shareDoc $ layoutType t1
  tRDocs <- tR `forM` (shareDoc . layoutType)
  docAlt
    [ docSeq
      [ docLitS "("
      , docForceSingleline $ docSeq $ List.intersperse
        docCommaSep
        (t1Doc : tRDocs)
      , docLitS ") =>"
      , docSeparator
      ]
    , docLines $ join
      [ [docSeq [docLitS "(", docSeparator, t1Doc]]
      , tRDocs <&> \tRDoc -> docSeq [docLitS ",", docSeparator, tRDoc]
      , [docLitS ") =>", docSeparator]
      ]
    ]

createBndrDoc :: [LHsTyVarBndr flag GhcPs] -> ToBriDocM BriDocNumbered
createBndrDoc bs = do
  tyVarDocs <- bs `forM` \case
    (L _ (UserTyVar _ _ext vname)) -> return $ (lrdrNameToText vname, Nothing)
    (L _ (KindedTyVar _ _ext lrdrName kind)) -> do
      d <- shareDoc $ layoutType kind
      return $ (lrdrNameToText lrdrName, Just $ d)
  docSeq $ List.intersperse docSeparator $ tyVarDocs <&> \(vname, mKind) ->
    case mKind of
      Nothing -> docLit vname
      Just kind -> docSeq
        [ docLitS "("
        , docLit vname
        , docSeparator
        , docLitS "::"
        , docSeparator
        , kind
        , docLitS ")"
        ]

createDerivingPar
  :: HsDeriving GhcPs -> ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
createDerivingPar derivs mainDoc = do
  case derivs of
    [] -> mainDoc
    types ->
      docPar mainDoc
        $ docEnsureIndent BrIndentRegular
        $ docLines
        -- TODO92 $ docWrapNode derivs
        $ derivingClauseDoc
        <$> types

derivingClauseDoc :: LHsDerivingClause GhcPs -> ToBriDocM BriDocNumbered
derivingClauseDoc (L _ (HsDerivingClause epAnn mStrategy types)) =
  case types of
    L _ (DctSingle _ ty) ->
      let
        (lhsStrategy, rhsStrategy) =
          maybe (docEmpty, docEmpty) strategyLeftRight mStrategy
      in docSeq
        [ docDeriving
        , docHandleComms types $ lhsStrategy
        , docSeparator
        , docHandleListElemComms layoutSigType ty -- TODO92 `docHandleRemaining types` here ?
          -- \case
          --     HsIB _ t -> layoutType t
        , rhsStrategy
        ]
    (L (SrcSpanAnn _multiEpAnn _) (DctMulti NoExtField [])) -> docSeq []
    (L (SrcSpanAnn multiEpAnn _) (DctMulti NoExtField ts)) ->
      let
        tsLength = length ts
        whenMoreThan1Type val = if tsLength > 1 then docLitS val else docLitS ""
        (lhsStrategy, rhsStrategy) =
          maybe (docEmpty, docEmpty) strategyLeftRight mStrategy
        posClose = case multiEpAnn of
          EpAnn _ (AnnContext _ _ [s]) _ -> Just $ epaLocationRealSrcSpanStart s
          _ -> Nothing
      in docSeq
        [ docDeriving
        , docHandleComms types $ lhsStrategy
        , docSeparator
        , whenMoreThan1Type "("
        , docSeq -- TODO92 `docHandleRemaining types` here ?
        $ List.intersperse docCommaSep
        $ ts <&> docHandleListElemComms layoutSigType
        , docHandleComms posClose $ whenMoreThan1Type ")"
        , rhsStrategy
        ]
 where
  posDeriving = obtainAnnPos epAnn AnnDeriving
  docDeriving = docHandleComms epAnn $ docHandleComms posDeriving $ docLitS "deriving"
  strategyLeftRight = \case
    (L _ (StockStrategy _)) -> (docLitS " stock", docEmpty)
    (L _ (AnyclassStrategy _)) -> (docLitS " anyclass", docEmpty)
    (L _ (NewtypeStrategy _)) -> (docLitS " newtype", docEmpty)
    _lVia@(L _ (ViaStrategy (XViaStrategyPs viaEpAnn viaType))) ->
      ( docEmpty
      , docSeq
        [ docHandleComms viaEpAnn $ docLitS " via"
        , docSeparator
        , docHandleListElemComms layoutSigType viaType
        ]
      )


createDetailsDoc
  :: Text -> HsConDeclH98Details GhcPs -> (ToBriDocM BriDocNumbered)
createDetailsDoc consNameStr details = case details of
  PrefixCon _ args -> do
    indentPolicy <- mAsk <&> _conf_layout .> _lconfig_indentPolicy .> confUnpack
    let
      singleLine = docSeq
        [ docLit consNameStr
        , docSeparator
        , docForceSingleline
        $ docSeq
        $ List.intersperse docSeparator
        $ fmap hsScaledThing args
        <&> layoutType
        ]
      leftIndented =
        docSetParSpacing
          . docAddBaseY BrIndentRegular
          . docPar (docLit consNameStr)
          . docLines
          $ layoutType
          <$> fmap hsScaledThing args
      multiAppended = docSeq
        [ docLit consNameStr
        , docSeparator
        , docSetBaseY $ docLines $ layoutType <$> fmap hsScaledThing args
        ]
      multiIndented = docSetBaseY $ docAddBaseY BrIndentRegular $ docPar
        (docLit consNameStr)
        (docLines $ layoutType <$> fmap hsScaledThing args)
    case indentPolicy of
      IndentPolicyLeft -> docAlt [singleLine, leftIndented]
      IndentPolicyMultiple -> docAlt [singleLine, multiAppended, leftIndented]
      IndentPolicyFree ->
        docAlt [singleLine, multiAppended, multiIndented, leftIndented]
  RecCon (L _ []) ->
    docSeq [docLit consNameStr, docSeparator, docLit $ Text.pack "{}"]
  RecCon (L (SrcSpanAnn epAnn _) fields@(_ : _)) -> do
    let posOpen = obtainAnnPos epAnn AnnOpenC
    let posClose = obtainAnnPos epAnn AnnCloseC
    let ((fName1, fType1), fDocR) = case mkFieldDocs fields of
          (doc1:docR) -> (doc1, docR)
          _ -> error "cannot happen (TM)"
    -- allowSingleline <- mAsk <&> _conf_layout .> _lconfig_allowSinglelineRecord .> confUnpack
    let allowSingleline = False
    docAddBaseY BrIndentRegular $ runFilteredAlternative $ do
        -- single-line: { i :: Int, b :: Bool }
      addAlternativeCond allowSingleline $ docSeq
        [ docLit consNameStr
        , docSeparator
        , docHandleComms posOpen $ docLitS "{"
        , docSeparator
        , docForceSingleline
        $ docHandleComms epAnn
        $ docSeq
        $ join
        $ [fName1, docSeparator, docLitS "::", docSeparator, fType1]
        : [ [ docLitS ","
            , docSeparator
            , fName
            , docSeparator
            , docLitS "::"
            , docSeparator
            , fType
            ]
          | (fName, fType) <- fDocR
          ]
        , docSeparator
        , docHandleComms posClose $ docLitS "}"
        ]
      addAlternative $ docPar
        (docLit consNameStr)
        (docNonBottomSpacingS $ docLines
          [ docAlt
            [ docCols
              ColRecDecl
              [ docHandleComms posOpen $ appSep (docLitS "{")
              , docHandleComms epAnn $ appSep $ docForceSingleline fName1
              , docSeq [docLitS "::", docSeparator]
              , docForceSingleline $ fType1
              ]
            , docSeq
              [ docHandleComms posOpen $ docLitS "{"
              , docSeparator
              , docSetBaseY $ docAddBaseY BrIndentRegular $ docPar
                fName1
                (docSeq [docLitS "::", docSeparator, fType1])
              ]
            ]
          , docLines $ fDocR <&> \(fName, fType) ->
            docAlt
              [ docCols
                ColRecDecl
                [ docCommaSep
                , appSep $ docForceSingleline fName
                , docSeq [docLitS "::", docSeparator]
                , docForceSingleline fType
                ]
              , docSeq
                [ docLitS ","
                , docSeparator
                , docSetBaseY $ docAddBaseY BrIndentRegular $ docPar
                  fName
                  (docSeq [docLitS "::", docSeparator, fType])
                ]
              ]
          , docHandleComms posClose $ docLitS "}"
          ]
        )
  InfixCon arg1 arg2 -> docSeq
    [ layoutType $ hsScaledThing arg1
    , docSeparator
    , docLit consNameStr
    , docSeparator
    , layoutType $ hsScaledThing arg2
    ]
 where
  mkFieldDocs
    :: [LConDeclField GhcPs]
    -> [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
  mkFieldDocs = map createNamesAndTypeDoc

createForallDoc
  :: [LHsTyVarBndr flag GhcPs] -> Maybe (ToBriDocM BriDocNumbered)
createForallDoc [] = Nothing
createForallDoc lhsTyVarBndrs =
  Just $ docSeq [docLitS "forall ", createBndrDoc lhsTyVarBndrs]

createNamesAndTypeDoc
  :: LConDeclField GhcPs
  -> (ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)
createNamesAndTypeDoc lField@(L _ (ConDeclField epAnn names t _)) =
  ( docFlushCommsPost False posColon
    $ docHandleComms posStart
    $ docHandleComms epAnn
    $ docSeq
        [ docSeq $ List.intersperse docCommaSep $ names <&> \case
            L _ (FieldOcc _ fieldName) ->
              docLit =<< lrdrNameToTextAnn fieldName
        ]
  , docFlushCommsPost True posComma (layoutType t)
  )
 where
  (posStart, posComma) = obtainListElemStartCommaLocs lField
  posColon             = obtainAnnPos epAnn AnnDcolon