{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE StrictData #-}

module HaskellCodeExplorer.AST.RenamedSource
  ( NameOccurrence(..)
  , namesFromRenamedSource
  ) where

-- import BasicTypes (TupleSort(..))
import GHC.Types.Basic (TupleSort(..))
-- import BooleanFormula (BooleanFormula(..))
import GHC.Data.BooleanFormula (BooleanFormula(..))
import Data.Generics (Data, everything, extQ, mkQ)
import Data.Maybe (mapMaybe)
import qualified Data.Text as T (Text)
import GHC
  ( AmbiguousFieldOcc(..)
  , ConDecl(..)
  , ConDeclField(..)
#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
#else
  , DataFamInstDecl(..)
#endif
  , FamilyDecl(..)
  , FieldOcc(..)
  , FixitySig(..)
  , ForeignDecl(..)
  , FunDep(..)
  , GenLocated(..)
  , getLocA
  , HsBindLR(..)
  , HsExpr(..)
#if MIN_VERSION_GLASGOW_HASKELL(8,4,1,0)
  , HsPatSynDetails
#else
  , HsPatSynDetails(..)
#endif
  , HsRecField'(..)
  , HsTupleSort(..)
  , HsTyLit(..)
  , HsTyVarBndr(..)
  , HsType(..)
  , IE(..)
  , LHsBindLR
  , LHsExpr
#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
#else
  , LHsQTyVars(..)
#endif
  , LHsType
  , LPat
  , LSig
  , LTyClDecl
  , Located
  , HsBracket(..)
#if MIN_VERSION_GLASGOW_HASKELL(8,2,2,0)
  , HsMatchContext(..)
  , Match(..)
#else
  , m_fixity
  , MatchFixity(..)
#endif
  , MatchGroup(..)
  , Name
  , Pat(..)
  , PatSynBind(..)
  , reLocN
  , Sig(..)
  , TyClDecl(..)
#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
  , FamEqn(..)
  , HsDataDefn(..)
#else
  , TyFamEqn(..)
#endif
  , Type
  , RoleAnnotDecl(..)
  , InjectivityAnn (..)
  , unLoc
  )
#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
-- import HsExtension (GhcRn)
import GHC.Hs.Extension (GhcRn)
#endif
import HaskellCodeExplorer.GhcUtils (hsPatSynDetails, ieLocNames)
import Prelude hiding (span)
-- import TysWiredIn
import GHC.Builtin.Types
  ( nilDataConName
  , tupleTyConName
--  , typeNatKind
  , naturalTy
  , typeSymbolKind
  )
-- import SrcLoc
import GHC.Types.SrcLoc
  ( mkRealSrcSpan
  , mkRealSrcLoc
  , realSrcSpanEnd
  , realSrcSpanStart
  , srcLocCol
  , srcLocFile
  , srcLocLine
  , SrcSpan(..)
  )
data NameOccurrence
  = NameOccurrence { locatedName :: Located (Maybe Name)
                   , description :: T.Text
                   , isBinder :: Bool }
  | TyLitOccurrence { locatedName :: Located (Maybe Name)
                    , description :: T.Text
                    , kind :: Type }

-- | Here we are only interested in a small subset of all AST nodes, so it is
-- convenient to use generic functions
namesFromRenamedSource :: (Data a) => a -> [NameOccurrence]
namesFromRenamedSource =
  everything
    (++)
    ([]
      `mkQ` hsExprNames
     `extQ` matchGroupNames
     `extQ` bindNames
     `extQ` patNames
     `extQ` sigNames
     `extQ` hsTypeNames
     `extQ` tyClDeclNames
     `extQ` familyDeclNames
     `extQ` familyEqNames
     `extQ` dataEqNames
     `extQ` conDeclNames
     `extQ` importNames
     `extQ` hsTyVarBndrNames
     `extQ` hsPatSynDetailsNames
     `extQ` conDeclFieldNames
     `extQ` hsRecFieldExprNames
     `extQ` hsRecAmbFieldExprNames
     `extQ` hsRecFieldPatNames
     `extQ` foreignDeclNames
     `extQ` roleAnnotationNames
     `extQ` injectivityAnnotationNames
    )

fieldOccName :: Bool -> FieldOcc GhcRn -> NameOccurrence
fieldOccName isBinder (FieldOcc name located) =
  NameOccurrence
    { locatedName = L (getLocA located) (Just name)
    , description = "FieldOcc"
    , isBinder = isBinder
    }

conDeclFieldNames :: ConDeclField GhcRn -> [NameOccurrence]
conDeclFieldNames ConDeclField {..} =
  map (fieldOccName True . unLoc) cd_fld_names

hsRecFieldExprNames :: HsRecField' (FieldOcc GhcRn) (LHsExpr GhcRn) -> [NameOccurrence]
hsRecFieldExprNames HsRecField {..} = [fieldOccName False $ unLoc hsRecFieldLbl]

hsRecAmbFieldExprNames :: HsRecField' (AmbiguousFieldOcc GhcRn) (LHsExpr GhcRn) -> [NameOccurrence]
hsRecAmbFieldExprNames HsRecField {..} =
  let (L span recField) = hsRecFieldLbl
      mbName =
        case recField of
          Ambiguous _ _ -> Nothing
          Unambiguous name _ -> Just name
   in [ NameOccurrence
          { locatedName = L span mbName
          , description = "AmbiguousFieldOcc"
          , isBinder = False
          }
      ]

#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
hsRecFieldPatNames :: HsRecField' (FieldOcc GhcRn) (LPat GhcRn) -> [NameOccurrence]
#else
hsRecFieldPatNames :: HsRecField' (FieldOcc Name) (LPat Name) -> [NameOccurrence]
#endif
hsRecFieldPatNames HsRecField {..} = [fieldOccName False $ unLoc hsRecFieldLbl]

hsExprNames :: LHsExpr GhcRn -> [NameOccurrence]
hsExprNames (L _span (HsVar _ name)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "HsVar"
    , isBinder = False
    }
  ]
hsExprNames lhe@(L _ (ExplicitList _ exprs))
  | null exprs =
    [ NameOccurrence
      { locatedName = L (getLocA lhe) $ Just nilDataConName
      , description = "ExplicitList"
      , isBinder = False
      }
    ]
  | otherwise = []
hsExprNames (L _span (RecordCon _ name _)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "RecordCon"
    , isBinder = False
    }
  ]
hsExprNames (L _span (HsRecFld _ (Unambiguous name located))) =
  [ NameOccurrence
    { locatedName = L (getLocA located) (Just name)
    , description = "HsRecFld"
    , isBinder = False
    }
  ]
hsExprNames (L _span (HsRecFld _ (Ambiguous _name located))) =
  [ NameOccurrence
    { locatedName = L (getLocA located) Nothing
    , description = "HsRecFld"
    , isBinder = False
    }
  ]
hsExprNames lhr@(L _ (HsRnBracketOut _ (VarBr _ quote name) _)) =
  case getLocA lhr of
    RealSrcSpan realSpan _ ->
      let start = realSrcSpanStart realSpan
          end = realSrcSpanEnd realSpan
          offset =
            if quote
              then 1 -- 'x
              else 2 -- ''T
          start' =
            mkRealSrcLoc
              (srcLocFile start)
              (srcLocLine start)
              (srcLocCol start + offset)
          span' = RealSrcSpan (mkRealSrcSpan start' end) Nothing
       in [ NameOccurrence
              { locatedName = L span' (Just $ unLoc name)
              , description = "VarBr"
              , isBinder = False
              }
          ]
    _ -> []
hsExprNames _ = []

matchGroupNames :: MatchGroup GhcRn (LHsExpr GhcRn) -> [NameOccurrence]
matchGroupNames =
  mapMaybe (fmap toNameOcc . matchContextName . m_ctxt . unLoc) .
  unLoc . mg_alts
  where
    --matchContextName :: HsMatchContext Name -> Maybe (Located Name)
    matchContextName (FunRhs name _ _bool) = Just name
    matchContextName _ = Nothing
    --toNameOcc :: LIdP GhcRn -> NameOccurrence
    toNameOcc n =
      NameOccurrence
        {locatedName = Just <$> reLocN n, description = "Match", isBinder = True}

bindNames :: LHsBindLR GhcRn GhcRn -> [NameOccurrence]
bindNames (L _span (PatSynBind _ PSB {..})) =
  [ NameOccurrence
      { locatedName = Just <$> reLocN psb_id
      , description = "PatSynBind"
      , isBinder = True
      }
  ]
bindNames _ = []

hsPatSynDetailsNames :: HsPatSynDetails GhcRn -> [NameOccurrence]
hsPatSynDetailsNames =
  map
    (\name ->
       NameOccurrence
         { locatedName = Just <$> name
         , description = "HsPatSynDetails"
         , isBinder = True
         }) .
  hsPatSynDetails

#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
importNames :: IE GhcRn -> [NameOccurrence]
#else
importNames :: IE Name -> [NameOccurrence]
#endif
importNames =
  map
    (\name ->
        NameOccurrence
        { locatedName = Just <$> name
        , description = "IE"
        , isBinder = False
        }) .
  ieLocNames


patNames :: LPat GhcRn -> [NameOccurrence]
patNames (L _span (VarPat _ name)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "VarPat"
    , isBinder = True
    }
  ]
patNames (L _span (ConPat _ name _)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "ConPatIn"
    , isBinder = False
    }
  ]
patNames (L _span (AsPat _ name _)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "AsPat"
    , isBinder = True
    }
  ]
patNames (L _span (NPlusKPat _ name _ _ _ _)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "NPlusKPat"
    , isBinder = True
    }
  ]
patNames _ = []


sigNames :: LSig GhcRn -> [NameOccurrence]
sigNames (L _span (TypeSig _ names _)) =
  map
    (\n ->
        NameOccurrence
        { locatedName = Just <$> reLocN n
        , description = "TypeSig"
        , isBinder = False
        })
    names

sigNames (L _span (PatSynSig _ names _)) =
  map (\name -> NameOccurrence (Just <$> reLocN name) "PatSynSig" False) names
sigNames (L _span (ClassOpSig _ _ names _)) =
  map
    (\n ->
        NameOccurrence
        { locatedName = Just <$> reLocN n
        , description = "ClassOpSig"
        , isBinder = True
        })
    names
sigNames (L _span (FixSig _ (FixitySig _ names _))) =
  map
    (\n ->
        NameOccurrence
        { locatedName = Just <$> reLocN n
        , description = "FixitySig"
        , isBinder = False
        })
    names
sigNames (L _span (InlineSig _ name _)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "InlineSig"
    , isBinder = False
    }
  ]
sigNames (L _span (SpecSig _ name _ _)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "SpecSig"
    , isBinder = False
    }
  ]
sigNames (L _span (MinimalSig _ _ (L _ boolFormula))) =
  map
    (\n ->
        NameOccurrence
        { locatedName = Just <$> reLocN n
        , description = "MinimalSig"
        , isBinder = False
        }) .
  boolFormulaNames $
  boolFormula
  where
    boolFormulaNames :: BooleanFormula name -> [name]
    boolFormulaNames (Var a) = [a]
    boolFormulaNames (And fs) = concatMap (boolFormulaNames . unLoc) fs
    boolFormulaNames (Or fs) = concatMap (boolFormulaNames . unLoc) fs
    boolFormulaNames (Parens (L _ f)) = boolFormulaNames f
sigNames (L _ _) = []

hsTypeNames :: LHsType GhcRn -> [NameOccurrence]
hsTypeNames (L _span (HsTyVar _ _promoted name)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "HsTyVar"
    , isBinder = False
    }
  ]
hsTypeNames lht@(L _ (HsTyLit _ lit)) =
  let kind =
        case lit of
          HsNumTy _ _ -> naturalTy
          HsStrTy _ _ -> typeSymbolKind
          HsCharTy _ _ -> typeSymbolKind
  in [ TyLitOccurrence
       { locatedName = L (getLocA lht) Nothing
       , description = "HsTyLit"
       , kind = kind
       }
     ]
hsTypeNames (L _span (HsOpTy _ _ name _)) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN name
    , description = "HsOpTy"
    , isBinder = False
    }
  ]
hsTypeNames lht@(L _ (HsTupleTy _ tupleSort types))
  | null types =
    let sort =
          case tupleSort of
            HsUnboxedTuple -> UnboxedTuple
            HsBoxedOrConstraintTuple -> BoxedTuple
    in [ NameOccurrence
         { locatedName = L (getLocA lht) (Just $ tupleTyConName sort 0)
         , description = "HsTupleTy"
         , isBinder = False
         }
       ]
  | otherwise = []
--https://ghc.haskell.org/trac/ghc/ticket/13737
--hsTypeNames (L span (HsExplicitListTy _kind types)) = ...
--hsTypeNames (L span (HsExplicitTupleTy _kind types)) = ...
hsTypeNames _ = []

-- the flag in HsTyVarBndr flag GhcRn can be either () (for visible
-- cases) or Specificity (for invisible cases).  I'm guessing we only
-- need to handle the visible cases, but it's trivial to replicate
-- this function for invisible cases
hsTyVarBndrNames :: HsTyVarBndr () GhcRn -> [NameOccurrence]
hsTyVarBndrNames (UserTyVar _ _ n) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN n
    , description = "UserTyVar"
    , isBinder = True
    }
  ]
hsTyVarBndrNames (KindedTyVar _ _ n _) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN n
    , description = "KindedTyVar"
    , isBinder = True
    }
  ]


#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
tyClDeclNames :: LTyClDecl GhcRn -> [NameOccurrence]
#else
tyClDeclNames :: LTyClDecl Name -> [NameOccurrence]
#endif
tyClDeclNames (L _span DataDecl {..}) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN tcdLName
    , description = "DataDecl"
    , isBinder = True
    }
  ]
tyClDeclNames (L _span SynDecl {..}) =
  [ NameOccurrence
    { locatedName = Just <$> reLocN tcdLName
    , description = "SynDecl"
    , isBinder = True
    }
  ]
tyClDeclNames (L _span ClassDecl {..}) =
  NameOccurrence
  { locatedName = Just <$> reLocN tcdLName
  , description = "ClassDecl"
  , isBinder = True
  } :
  concatMap (go . unLoc) tcdFDs
  where
    -- toNameOcc :: Located Name -> NameOccurrence
    toNameOcc n =
      NameOccurrence
      { locatedName = Just <$> reLocN n
      , description = "FunDep"
      , isBinder = False
      }
    go (FunDep _ names1 names2) = map toNameOcc names1 ++ map toNameOcc names2
    go _ = []
tyClDeclNames _ = []

familyDeclNames :: FamilyDecl GhcRn -> [NameOccurrence]
familyDeclNames FamilyDecl {..} =
  [ NameOccurrence
    { locatedName = Just <$> reLocN fdLName
    , description = "FamilyDecl"
    , isBinder = True
    }
  ]

familyEqNames :: FamEqn GhcRn (LHsType GhcRn) -> [NameOccurrence]
familyEqNames FamEqn {feqn_tycon = tyCon} =
  [ NameOccurrence
    { locatedName = Just <$> reLocN tyCon
    , description = "FamEqn"
    , isBinder = False
    }
  ]
dataEqNames :: FamEqn GhcRn (HsDataDefn GhcRn) -> [NameOccurrence]
dataEqNames FamEqn {feqn_tycon = tyCon} =
  [ NameOccurrence
    { locatedName = Just <$> reLocN tyCon
    , description = "FamEqn"
    , isBinder = False
    }
  ]

conDeclNames :: ConDecl GhcRn -> [NameOccurrence]
conDeclNames con =
  case con of
    ConDeclGADT {con_names = names} ->
      map
        (\n ->
            NameOccurrence
            { locatedName = Just <$> reLocN n
            , description = "ConDeclGADT"
            , isBinder = True
            })
        names
    ConDeclH98 {con_name = name} ->
      [ NameOccurrence
        { locatedName = Just <$> reLocN name
        , description = "ConDeclH98"
        , isBinder = True
        }
      ]

#if MIN_VERSION_GLASGOW_HASKELL(8,4,3,0)
foreignDeclNames :: ForeignDecl GhcRn -> [NameOccurrence]
#else
foreignDeclNames :: ForeignDecl Name -> [NameOccurrence]
#endif
foreignDeclNames decl =
  [ NameOccurrence
    { locatedName = Just <$> reLocN (fd_name decl)
    , description = "ForeignDecl"
    , isBinder = True
    }
  ]

roleAnnotationNames :: RoleAnnotDecl GhcRn -> [NameOccurrence]
roleAnnotationNames (RoleAnnotDecl _ n _) =
  [ NameOccurrence
      { locatedName = Just <$> reLocN n
      , description = "RoleAnnotDecl"
      , isBinder = False
      }
  ]

injectivityAnnotationNames :: InjectivityAnn GhcRn -> [NameOccurrence]
injectivityAnnotationNames (InjectivityAnn _ lhsName rhsNames) =
  injAnnNameOcc lhsName : map injAnnNameOcc rhsNames
  where
    -- injAnnNameOcc :: GenLocated SrcSpan Name -> NameOccurrence
    injAnnNameOcc n =
      NameOccurrence
        { locatedName = Just <$> reLocN n
        , description = "InjectivityAnn"
        , isBinder = False
        }