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

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

import           Data.Generics                  ( Data
                                                , everything
                                                , extQ
                                                , mkQ
                                                )
import           Data.Maybe                     ( mapMaybe )
import qualified Data.Text                     as T
                                                ( Text )
import           GHC                            ( AmbiguousFieldOcc(..)
                                                , ConDecl(..)
                                                , ConDeclField(..)
                                                , FamEqn(..)
                                                , FamilyDecl(..)
                                                , FieldOcc(..)
                                                , FixitySig(..)
                                                , ForeignDecl(..)
                                                , FunDep(..)
                                                , GenLocated(..)
                                                , HsBindLR(..)
                                                , HsBracket(..)
                                                , HsDataDefn(..)
                                                , HsExpr(..)
                                                , HsMatchContext(..)
                                                , HsPatSynDetails
                                                , HsRecField'(..)
                                                , HsTupleSort(..)
                                                , HsTyLit(..)
                                                , HsTyVarBndr(..)
                                                , HsType(..)
                                                , IE(..)
                                                , InjectivityAnn(..)
                                                , LHsBindLR
                                                , LHsExpr
                                                , LHsType
                                                , LPat
                                                , LSig
                                                , LTyClDecl
                                                , Located
                                                , Match(..)
                                                , MatchGroup(..)
                                                , Name
                                                , Pat(..)
                                                , PatSynBind(..)
                                                , RoleAnnotDecl(..)
                                                , Sig(..)
                                                , TyClDecl(..)
                                                , Type
                                                , getLocA
                                                , reLocN
                                                , unLoc
                                                )
import           GHC.Builtin.Types              ( naturalTy
                                                , nilDataConName
                                                , tupleTyConName
                                                , typeSymbolKind
                                                )
import           GHC.Data.BooleanFormula        ( BooleanFormula(..) )
import           GHC.Hs.Extension               ( GhcRn )
import           GHC.Types.Basic                ( TupleSort(..) )
import           GHC.Types.SrcLoc               ( SrcSpan(..)
                                                , mkRealSrcLoc
                                                , mkRealSrcSpan
                                                , realSrcSpanEnd
                                                , realSrcSpanStart
                                                , srcLocCol
                                                , srcLocFile
                                                , srcLocLine
                                                )
import           HaskellCodeExplorer.GhcUtils   ( hsPatSynDetails
                                                , ieLocNames
                                                )
import           Prelude                 hiding ( span )
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
                       }
      ]

hsRecFieldPatNames
  :: HsRecField' (FieldOcc GhcRn) (LPat GhcRn) -> [NameOccurrence]
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

importNames :: IE GhcRn -> [NameOccurrence]
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
                   }
  ]


tyClDeclNames :: LTyClDecl GhcRn -> [NameOccurrence]
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
                     }
    ]

foreignDeclNames :: ForeignDecl GhcRn -> [NameOccurrence]
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
                                   }