{-# LANGUAGE CPP              #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}
{-# LANGUAGE LambdaCase       #-}
{-# LANGUAGE PatternSynonyms  #-}
{-# LANGUAGE RankNTypes       #-}
{-# LANGUAGE RecordWildCards  #-}
{-# LANGUAGE TupleSections    #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-missing-signatures #-}
{-# OPTIONS_GHC -Wno-overlapping-patterns #-}
module Ide.Plugin.GHC where

import           Data.Functor                            ((<&>))
import           Data.List.Extra                         (stripInfix)
import qualified Data.List.NonEmpty                      as NE
import qualified Data.Text                               as T
import           Development.IDE
import           Development.IDE.GHC.Compat
import           Development.IDE.GHC.Compat.ExactPrint
import           Ide.PluginUtils                         (subRange)
import           Language.Haskell.GHC.ExactPrint.Parsers (parseDecl)

import           GHC.Parser.Annotation                   (AddEpAnn (..),
                                                          Anchor (Anchor),
                                                          AnchorOperation (MovedAnchor),
                                                          DeltaPos (..),
                                                          EpAnn (..),
                                                          EpAnnComments (EpaComments),
                                                          EpaLocation (EpaDelta),
                                                          SrcSpanAnn' (SrcSpanAnn),
                                                          spanAsAnchor)
#if MIN_VERSION_ghc(9,5,0)
import           GHC.Parser.Annotation                   (TokenLocation (..))
#endif
import           Language.Haskell.GHC.ExactPrint         (showAst)

type GP = GhcPass Parsed

-- | Check if a given range is in the range of located item
inRange :: HasSrcSpan a => Range -> a -> Bool
inRange range s = maybe False (subRange range) (srcSpanToRange (getLoc s))

-- | Get data decl and its location
getDataDecl :: LHsDecl GP -> Maybe (LTyClDecl GP)
getDataDecl (L l (TyClD _ d@DataDecl{})) = Just (L l d)
getDataDecl _                            = Nothing

isConDeclH98 :: ConDecl GP -> Bool
isConDeclH98 ConDeclH98{} = True
isConDeclH98 _            = False

isH98DataDecl :: LTyClDecl GP -> Bool
isH98DataDecl (L _ decl@DataDecl{}) =
    any (isConDeclH98 . (\(L _ r) -> r)) (dd_cons $ tcdDataDefn decl)
isH98DataDecl _ = False

-- | Convert H98 data type definition to GADT's
h98ToGADTDecl :: TyClDecl GP -> TyClDecl GP
h98ToGADTDecl = \case
    DataDecl{..} -> DataDecl
        { tcdDataDefn = updateDefn tcdLName tcdTyVars tcdDataDefn
        , ..
        }
    x -> x
    where
        updateDefn dataName tyVars = \case
            HsDataDefn{..} -> HsDataDefn
                { dd_cons =
                    mapX (h98ToGADTConDecl dataName tyVars (wrapCtxt dd_ctxt)) <$> dd_cons
                , dd_ctxt = emptyCtxt -- Context can't appear at the data name in GADT
                , ..
                }
            x -> x

-- | Convert H98 data constructor to GADT data constructor
h98ToGADTConDecl ::
    LIdP GP -- ^Type constructor name,
            -- used for constructing final result type in GADT
    -> LHsQTyVars GP
            -- ^Type variable names
            -- used for constructing final result type in GADT
    -> Maybe (LHsContext GP)
            -- ^Data type context
    -> ConDecl GP
    -> ConDecl GP
h98ToGADTConDecl dataName tyVars ctxt = \case
    ConDeclH98{..} ->
        ConDeclGADT
            con_ext
#if MIN_VERSION_ghc(9,5,0)
            (NE.singleton con_name)
#else
            [con_name]
#endif

#if MIN_VERSION_ghc(9,5,0)
            (L NoTokenLoc HsNormalTok)
#endif
            -- Ignore all existential type variable since GADT not needed
            implicitTyVars
            (mergeContext ctxt con_mb_cxt)
            (renderDetails con_args)
            renderResultTy
            con_doc
    x -> x
    where
        -- Parameters in the data constructor
        renderDetails :: HsConDeclH98Details GP -> HsConDeclGADTDetails GP
        renderDetails (PrefixCon _ args)   = PrefixConGADT args
        renderDetails (InfixCon arg1 arg2) = PrefixConGADT [arg1, arg2]
#if MIN_VERSION_ghc(9,3,0)
        renderDetails (RecCon recs)        = RecConGADT recs noHsUniTok
#else
        renderDetails (RecCon recs)        = RecConGADT recs
#endif


        -- | Construct GADT result type
        renderResultTy :: LHsType GP
        renderResultTy = case tyVars of
            -- Without type variable
            HsQTvs _ []   -> wrappedDataName
            -- With type variable
            HsQTvs _ vars -> foldl go wrappedDataName vars
            _             -> wrappedDataName
            where
                wrappedDataName = wrap (HsTyVar noUsed NotPromoted dataName)

                -- Bundle data name with type vars by `HsAppTy`
                go acc (L _(UserTyVar' var)) =
                    wrap
                        (HsAppTy noExtField acc
                            (wrap (HsTyVar noUsed NotPromoted var)))
                go acc _ = acc

        -- Merge data type context and constructor type context
        mergeContext :: Maybe (LHsContext GP) -> Maybe (LHsContext GP) -> Maybe (LHsContext GP)
        mergeContext ctxt1 ctxt2 =
            (wrap . map wrap) . map unParTy
                <$> (getContextType ctxt1 <> getContextType ctxt2)
            where
                getContextType :: Maybe (LHsContext GP) -> Maybe [HsType GP]
                getContextType ctxt = map unWrap . unWrap <$> ctxt

                -- Unparen the outmost, it only occurs at the outmost
                -- for a valid type.
                --
                -- Note for context paren rule:
                --
                -- If only one element, it __can__ have a paren type.
                -- If not, there can't have a parent type.
                unParTy :: HsType GP -> HsType GP
                unParTy (HsParTy _ ty) = unWrap ty
                unParTy x              = x
{- |
We use `printOutputable` to print H98 data decl as GADT syntax,
this print is not perfect, it will:

1. Make data name and the `where` key word in different lines.
2. Make the whole data decl prints in one line if there is only one data constructor.
3. The ident size of every data constructor depends on its origin
   format, and may have different ident size between constructors.

Hence, we first use `printOutputable` to get an initial GADT syntax,
then use `ghc-exactprint` to parse the initial result, and finally
adjust the details that mentioned above.

The adjustment includes:

1. Make the `where` key word at the same line of data name.
2. Remove the extra blank line caused by adjustment of `where`.
3. Make every data constructor start with a new line and 2 spaces
-}
prettyGADTDecl :: DynFlags -> TyClDecl GP -> Either String String
prettyGADTDecl df decl =
    let old = printOutputable decl
        hsDecl = parseDecl df "unused" (T.unpack old)
        tycld = adjustTyClD hsDecl
    in removeExtraEmptyLine . exactPrint <$> tycld
    where
        adjustTyClD = \case
                Right (L _ (TyClD _ tycld)) -> Right $ adjustDataDecl tycld
                Right x -> Left $ "Expect TyClD but got " <> showAst x
#if MIN_VERSION_ghc(9,3,0)
                Left err -> Left $ printWithoutUniques err
#else
                Left err -> Left $ show err
#endif

        adjustDataDecl DataDecl{..} = DataDecl
            { tcdDExt = adjustWhere tcdDExt
            , tcdDataDefn = tcdDataDefn
                { dd_cons =
                      fmap adjustCon (dd_cons tcdDataDefn)
                }
            , ..
            }
        adjustDataDecl x = x

        -- Make every data constructor start with a new line and 2 spaces
        adjustCon :: LConDecl GP -> LConDecl GP
        adjustCon (L (SrcSpanAnn _ loc) r) =
            L (SrcSpanAnn (EpAnn (go (spanAsAnchor loc)) (AnnListItem []) (EpaComments [])) loc) r
            where
                go (Anchor a _) = Anchor a (MovedAnchor (DifferentLine 1 2))

        -- Adjust where annotation to the same line of the type constructor
        adjustWhere tcdDExt = tcdDExt <&> map
            (\(AddEpAnn ann l) ->
            if ann == AnnWhere
                then AddEpAnn AnnWhere (EpaDelta (SameLine 1) [])
                else AddEpAnn ann l
            )

        -- Remove the first extra line if exist
        removeExtraEmptyLine s = case stripInfix "\n\n" s of
            Just (x, xs) -> x <> "\n" <> xs
            Nothing      -> s

wrap :: forall a. WrapXRec GP a => a -> XRec GP a
wrap = wrapXRec @GP
wrapCtxt = id
emptyCtxt = Nothing
unWrap = unXRec @GP
mapX = mapXRec @GP
noUsed = EpAnnNotUsed

pattern UserTyVar' :: LIdP pass -> HsTyVarBndr flag pass
pattern UserTyVar' s <- UserTyVar _ _ s

implicitTyVars = (wrapXRec @GP mkHsOuterImplicit)