{-# 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.ToBriDocTools
import Language.Haskell.Brittany.Internal.Prelude
import Language.Haskell.Brittany.Internal.Types
import Language.Haskell.Brittany.Internal.Components.BriDoc



layoutDataDecl
  :: Maybe (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
            tyVars <- mapM 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 $ docAlt
              [ -- newtype Tagged s b = Tagged { unTagged :: b }
                docSeq
                [ appSep $ docLitS "newtype"
                , appSep $ docLit nameStr
                , appSep (docSeqSep tyVars)
                , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                , docSeparator
                , docLitS "="
                , docSeparator
                , docForceParSpacing $ docHandleComms epAnn $ rhsDoc
                ]
              , -- newtype Tagged s b
                --   = Tagged { unTagged :: b }
                -- newtype Tagged s
                --                b
                --   = Tagged { unTagged :: b }
                docAddBaseY BrIndentRegular $ docPar
                  ( docSeq
                    [ appSep $ docLitS "newtype"
                    , appSep $ docLit nameStr
                    , docAlt
                      [ docForceSingleline $ docSeq
                        [ appSep (docSeqSep tyVars)
                        , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                        ]
                      , docSetBaseY $ docLines
                        $ map docForceSingleline $ tyVars ++ patDocs
                      ]
                    ]
                  )
                  ( docSeq
                    [ docLitS "="
                    , docSeparator
                    , docHandleComms epAnn $ rhsDoc
                    ]
                  )
              , -- newtype Tagged
                --   s
                --   b
                --   = Tagged { unTagged :: b }
                docAddBaseY BrIndentRegular $ docPar
                  ( docSeq
                    [ appSep $ docLitS "newtype"
                    , appSep $ docLit nameStr
                    -- , appSep tyVarLine
                    -- , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                    ]
                  )
                  ( docLines
                    $  map (docEnsureIndent BrIndentRegular) tyVars
                    ++ map (docEnsureIndent BrIndentRegular) patDocs
                    ++ [ docSeq
                         [ docLitS "="
                         , docSeparator
                         , docHandleComms epAnn $ rhsDoc
                         ]
                       ]
                  )
              ]
          else maybe
            (error
            $  "Unsupported form of DataFamInstDecl:"
            ++ " ConDeclH98 with context"
            )
            briDocByExactNoComment
            ltycl
      _ -> maybe
            (error
            $  "Unsupported form of DataFamInstDecl:"
            ++ " ConDeclH98 with forall"
            )
            briDocByExactNoComment
            ltycl
  HsDataDefn NoExtField NewType _ _ Just{} _ _ -> maybe
    (error $ "Unsupported form of DataFamInstDecl: NewType _ _ Just _ _")
    briDocByExactNoComment
    ltycl
  HsDataDefn NoExtField NewType _ _ Nothing _ _ -> maybe
    (error $ "Unsupported form of DataFamInstDecl: NewType _ _ Nothing _ _")
    briDocByExactNoComment
    ltycl

  -- data MyData = MyData ..
  -- data MyData = MyData { .. }
  HsDataDefn NoExtField DataType ctxMay _ctype Nothing conss mDerivs -> do
    lhsContextDoc <- case ctxMay of
      Just (L _ lhsContext) -> shareDoc $ createContextDoc lhsContext
      Nothing -> pure docEmpty
    nameStr <- lrdrNameToTextAnn name
    tyVarLine <- shareDoc $ docSeqSep $ createBndrDoc bndrs
    patDocs <- mapM shareDoc $ layoutHsTyPats pats
    lhsDoc <- shareDoc $ docSeq
      [ -- TOOD92 docNodeAnnKW ltycl (Just GHC.AnnData) $ 
        appSep $ docLitS "data"
      , docForceSingleline $ lhsContextDoc
      , appSep $ docLit nameStr
      , appSep tyVarLine
      , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
      ]
    let prefixes = "=" : repeat "|"
    layoutConssResult <- mapM layoutConDecl (zip prefixes conss)
    case sequence layoutConssResult of
      Left err -> maybe (error err) briDocByExactNoComment ltycl
      Right [] -> do
        docHandleComms declEpAnn $ createDerivingPar mDerivs $ docSeq
          [ appSep $ docLitS "data"
          , lhsContextDoc
          , appSep $ docLit nameStr
          , appSep tyVarLine
          , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
          ]
      Right [(consDocSl, consDocMl)] -> do
        docHandleComms declEpAnn $ createDerivingPar mDerivs $ docAlt
          [ -- data D = forall a . Show a => D a
            docSeq [lhsDoc, consDocSl]
          , -- data D
            --   = forall a . Show a => D a
            -- data D
            --   = forall a
            --   . Show a =>
            --     D a
            docAddBaseY BrIndentRegular
            $ docPar lhsDoc
              (docNonBottomSpacing $ docAlt [consDocSl, consDocMl])
          , -- data
            --   Show a =>
            --   D
            --   = rhsDoc
            -- 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
            (-- TODO92 docNodeAnnKW ltycl (Just GHC.AnnData) $
             docLitS "data")
            (docLines
              [ lhsContextDoc
              , docSeq
                [ appSep $ docLit nameStr
                , tyVarLine
                , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                ]
              , consDocMl
              ]
            )
          ]
      Right consDocTuples -> do
        docHandleComms declEpAnn
          $ createDerivingPar mDerivs
          $ docAddBaseY BrIndentRegular
          $ docPar
            (docAlt
              [ -- data Show a => D a
                lhsDoc
              , -- data
                --   Show a =>
                --   D
                -- 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
                (-- TODO92 docNodeAnnKW ltycl (Just GHC.AnnData) $
                 docLitS "data")
                (docLines
                  [ lhsContextDoc
                  , docSeq
                    [ appSep $ docLit nameStr
                    , tyVarLine
                    , docSeq $ [x | p <- patDocs, x <- [p, docSeparator]]
                    ]
                  ]
                )
              ]
            )
            (docLines $ [docAlt [sl, ml] | (sl, ml) <- consDocTuples])
  
  HsDataDefn NoExtField DataType _ _ Just{} _ _ -> maybe
    (error $ "Unsupported form of DataFamInstDecl: DataType _ _ Just _ _")
    briDocByExactNoComment
    ltycl

layoutConDecl
  :: (String, LConDecl GhcPs)
  -> ToBriDocM
       (Either String (ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered))
layoutConDecl (prefix, L _ con) = case con of
  ConDeclH98 epAnn consName _hasExt qvars mRhsContext details _conDoc -> do
    consNameStr  <- lrdrNameToTextAnn consName
    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
    detailsDoc <- shareDoc $ createDetailsDoc consNameStr details
    let posEqual = obtainAnnPos epAnn AnnEqual
    pure $ Right
      ( docSeq
          [ docHandleComms epAnn $ docHandleComms posEqual $ docLitS prefix
          , docSeparator
          , docSetIndentLevel $ docSeq
            [ case forallDocMay of
              Nothing        -> docEmpty
              Just forallDoc -> docSeq
                [ docForceSingleline forallDoc
                , docSeparator
                , docLitS "."
                , docSeparator
                ]
            , maybe docEmpty docForceSingleline rhsContextDocMay
            , detailsDoc
            ]
          ]
      , docHandleComms epAnn
        $ docHandleComms posEqual
        $ docNonBottomSpacing
        $ case (forallDocMay, rhsContextDocMay) of
            (Just forallDoc, Just rhsContextDoc) -> docLines
              [ docSeq
                [docLitS prefix, docSeparator, docForceSingleline forallDoc]
              , docSeq
                [ docLitS "."
                , docSeparator
                , docSetBaseY $ docLines [rhsContextDoc, docSetBaseY detailsDoc]
                ]
              ]
            (Just forallDoc, Nothing) -> docLines
              [ docSeq
                [docLitS prefix, docSeparator, docForceSingleline forallDoc]
              , docSeq [docLitS ".", docSeparator, detailsDoc]
              ]
            (Nothing, Just rhsContextDoc) -> docSeq
              [ docLitS prefix
              , docSeparator
              , docSetBaseY $ docLines [rhsContextDoc, docSetBaseY detailsDoc]
              ]
            (Nothing, Nothing) -> docSeq
              [docLitS prefix, docSeparator, detailsDoc]
      )
  ConDeclGADT{} -> pure
    $ Left "Unsupported: ConDeclGADT inside DataFamInstDecl"


layoutHsTyPats
  :: [LHsTypeArg GhcPs] -> [ToBriDocM BriDocNumbered]
layoutHsTyPats pats = pats <&> \case
  HsValArg tm -> callLayouter2 layout_type False tm
  HsTypeArg _l ty ->
    docSeq [docLit $ Text.pack "@", callLayouter2 layout_type False 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 [callLayouter2 layout_type False t, docSeparator, docLitS "=>", docSeparator]
createContextDoc (t1 : tR) = do
  t1Doc <- shareDoc $ callLayouter2 layout_type False t1
  tRDocs <- tR `forM` (shareDoc . callLayouter2 layout_type False)
  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 = map $ \x -> do
  (vname, mKind) <- case x of
    (L _ (UserTyVar _ _ext vname)) -> return $ (lrdrNameToText vname, Nothing)
    (L _ (KindedTyVar _ _ext lrdrName kind)) -> do
      d <- shareDoc $ callLayouter2 layout_type False kind
      return $ (lrdrNameToText lrdrName, Just $ d)
  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 (callLayouter layout_sigType) 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 (callLayouter layout_sigType)
        , 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 (callLayouter layout_sigType) 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
        <&> callLayouter2 layout_type False
        ]
      leftIndented =
        docSetParSpacing
          . docAddBaseY BrIndentRegular
          . docPar (docLit consNameStr)
          . docLines
          $ callLayouter2 layout_type False
          <$> fmap hsScaledThing args
      multiAppended = docSeq
        [ docLit consNameStr
        , docSeparator
        , docSetBaseY
          $ docLines
          $ callLayouter2 layout_type False <$> fmap hsScaledThing args
        ]
      multiIndented = docSetBaseY $ docAddBaseY BrIndentRegular $ docPar
        (docLit consNameStr)
        (docLines $ callLayouter2 layout_type False <$> 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
    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 $ docSetParSpacing $ 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 "{"
              , docHandleComms epAnn 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
    [ callLayouter2 layout_type False $ hsScaledThing arg1
    , docSeparator
    , docLit consNameStr
    , docSeparator
    , callLayouter2 layout_type False $ 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 ", docSeqSep $ 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
      (callLayouter2 layout_type (hasAnyCommentsBelow epAnn) t)
  )
 where
  (posStart, posComma) = obtainListElemStartCommaLocs lField
  posColon             = obtainAnnPos epAnn AnnDcolon