{-# LANGUAGE TypeFamilies #-}

module Descript.BasicInj.Process.Inspect.SymbolAt
  ( symbolAt
  , symbolAt'
  ) where

import qualified Descript.BasicInj.Traverse.Term as T
import Descript.BasicInj.Traverse
import Descript.BasicInj.Data
import Descript.Misc
import Data.Monoid as M
import Data.Foldable
import Data.List
import qualified Data.List.NonEmpty as NonEmpty

data SymbolAt = SymbolAt Loc

instance Fold SymbolAt where
  type Res SymbolAt = First (SymbolEnv SrcAnn)
  type FAnn SymbolAt = SrcAnn

  fonTerm T.ModulePath (SymbolAt loc) (ModulePath _ xs)
    = fold $ zipWith (tryModulePathElemAt loc) (inits xs') xs'
    where xs' = NonEmpty.toList xs
  fonTerm T.RecordHead (SymbolAt loc) x = tryRecordHeadAt loc x
  fonTerm T.RecordType (SymbolAt loc) (RecordType _ head' props)
    = foldMap (tryPropertyKeyAt loc head') props
  fonTerm T.GenRecord (SymbolAt loc) (Record _ head' props)
    = foldMap (tryPropertyKeyAt loc head' . propertyKey) props
  fonTerm T.PathElem (SymbolAt loc) (PathElem _ propKey' headKey')
    = tryPropertyKeyAt loc headKey' propKey'
  fonTerm _ _ _ = First Nothing -- mempty

-- | Finds the symbol at the given location, if it exists.
symbolAt :: Loc -> Source SrcAnn -> Maybe (SymbolEnv SrcAnn)
symbolAt = symbolAt' T.Source

-- | Finds the symbol at the given location, if it exists.
symbolAt' :: TTerm t -> Loc -> t SrcAnn -> Maybe (SymbolEnv SrcAnn)
symbolAt' term loc = getFirst . foldTerm term (SymbolAt loc)

tryModulePathElemAt :: Loc
                    -> [Symbol SrcAnn]
                    -> Symbol SrcAnn
                    -> First (SymbolEnv SrcAnn)
tryModulePathElemAt loc xs = fmap (EnvModulePathElem xs) . trySymbolAt loc

tryRecordHeadAt :: Loc
                -> FSymbol SrcAnn
                -> First (SymbolEnv SrcAnn)
tryRecordHeadAt loc (FSymbol scope sym)
  = EnvRecordHead . FSymbol scope <$> trySymbolAt loc sym

tryPropertyKeyAt :: Loc
                 -> FSymbol SrcAnn
                 -> Symbol SrcAnn
                 -> First (SymbolEnv SrcAnn)
tryPropertyKeyAt loc head' = fmap (EnvPropertyKey head') . trySymbolAt loc

trySymbolAt :: Loc -> Symbol SrcAnn -> First (Symbol SrcAnn)
trySymbolAt loc sym
  | loc `isInRange` srcRange (getAnn sym) = First $ Just sym
  | otherwise = First Nothing