content: module Dhall.LSP.Backend.Typing (annotateLet, exprAt, srcAt, typeAt) where
  
  import Dhall.Context   (Context, empty, insert)
  import Dhall.Core
      ( Binding (..)
      , Expr (..)
      , FunctionBinding (..)
      , Var (..)
      , normalize
      , shift
      , subExpressions
      , subst
      )
  import Dhall.Parser    (Src (..))
  import Dhall.TypeCheck (TypeError (..), typeWithA)
  
  import Control.Applicative ((<|>))
  import Control.Lens        (toListOf)
  import Data.Bifunctor      (first)
  import Data.Void           (Void, absurd)
  
  import Dhall.LSP.Backend.Dhall       (WellTyped, fromWellTyped)
  import Dhall.LSP.Backend.Diagnostics (Position, Range (..), rangeFromDhall)
  import Dhall.LSP.Backend.Parsing
      ( getForallIdentifier
      , getLamIdentifier
      , getLetAnnot
      , getLetIdentifier
      , getLetInner
      )
  
  -- | Find the type of the subexpression at the given position. Assumes that the
  --   input expression is well-typed. Also returns the Src descriptor containing
  --   that subexpression if possible.
  typeAt :: Position -> WellTyped -> Either String (Maybe Src, Expr Src Void)
  typeAt pos expr = do
    expr' <- case splitMultiLetSrc (fromWellTyped expr) of
               Just e -> return e
               Nothing -> Left "The impossible happened: failed to split let\
                                \ blocks when preprocessing for typeAt'."
    (mSrc, typ) <- first show $ typeAt' pos empty expr'
    case mSrc of
      Just src -> return (Just src, normalize typ)
      Nothing -> return (srcAt pos expr', normalize typ)
  
  typeAt' :: Position -> Context (Expr Src Void) -> Expr Src Void -> Either (TypeError Src Void) (Maybe Src, Expr Src Void)
  -- the user hovered over the bound name in a let expression
  typeAt' pos ctx (Note src (Let (Binding { value = a }) _)) | pos `inside` getLetIdentifier src = do
    typ <- typeWithA absurd ctx a
    return (Just $ getLetIdentifier src, typ)
  
  -- "..." in a lambda expression
  typeAt' pos _ctx (Note src (Lam _ FunctionBinding { functionBindingAnnotation = _A} _))
    | Just src' <- getLamIdentifier src
    , pos `inside` src' =
    return (Just src', _A)
  
  -- "..." in a forall expression
  typeAt' pos _ctx (Note src (Pi _ _ _A _)) | Just src' <- getForallIdentifier src
                                          , pos `inside` src' =
    return (Just src', _A)
  
  typeAt' pos ctx (Let (Binding { variable = x, value = a }) e@(Note src _)) | pos `inside` src = do
    _ <- typeWithA absurd ctx a
    let a' = shift 1 (V x 0) (normalize a)
    typeAt' pos ctx (shift (-1) (V x 0) (subst (V x 0) a' e))
  
  typeAt' pos ctx (Lam _ FunctionBinding { functionBindingVariable = x, functionBindingAnnotation = _A} b@(Note src _))
    | pos `inside` src = do
        let _A' = Dhall.Core.normalize _A
            ctx' = fmap (shift 1 (V x 0)) (insert x _A' ctx)
        typeAt' pos ctx' b
  
  typeAt' pos ctx (Pi _ x _A  _B@(Note src _)) | pos `inside` src = do
    let _A' = Dhall.Core.normalize _A
        ctx' = fmap (shift 1 (V x 0)) (insert x _A' ctx)
    typeAt' pos ctx' _B
  
  -- peel off a single Note constructor
  typeAt' pos ctx (Note _ expr) = typeAt' pos ctx expr
  
  -- catch-all
  typeAt' pos ctx expr = do
    let subExprs = toListOf subExpressions expr
    case [ (src, e) | (Note src e) <- subExprs, pos `inside` src ] of
      [] -> do typ <- typeWithA absurd ctx expr  -- return type of whole subexpression
               return (Nothing, typ)
      ((src, e):_) -> typeAt' pos ctx (Note src e)  -- continue with leaf-expression
  
  
  -- | Find the smallest Note-wrapped expression at the given position.
  exprAt :: Position -> Expr Src a -> Maybe (Expr Src a)
  exprAt pos e = do e' <- splitMultiLetSrc e
                    exprAt' pos e'
  
  exprAt' :: Position -> Expr Src a -> Maybe (Expr Src a)
  exprAt' pos e@(Note _ expr) = exprAt pos expr <|> Just e
  exprAt' pos expr =
    let subExprs = toListOf subExpressions expr
    in case [ (src, e) | (Note src e) <- subExprs, pos `inside` src ] of
      [] -> Nothing
      ((src,e) : _) -> exprAt' pos e <|> Just (Note src e)
  
  
  -- | Find the smallest Src annotation containing the given position.
  srcAt :: Position -> Expr Src a -> Maybe Src
  srcAt pos expr = do Note src _ <- exprAt pos expr
                      return src
  
  
  -- | Given a well-typed expression and a position find the let binder at that
  --   position (if there is one) and return the type annotation to be inserted
  --   (potentially replacing the existing one). If something goes wrong returns a
  --   textual error message.
  annotateLet :: Position -> WellTyped -> Either String (Src, Expr Src Void)
  annotateLet pos expr = do
    expr' <- case splitMultiLetSrc (fromWellTyped expr) of
               Just e -> return e
               Nothing -> Left "The impossible happened: failed to split let\
                                \ blocks when preprocessing for annotateLet'."
    annotateLet' pos empty expr'
  
  
  annotateLet' :: Position -> Context (Expr Src Void) -> Expr Src Void
               -> Either String (Src, Expr Src Void)
  -- the input only contains singleton lets
  annotateLet' pos ctx (Note src e@(Let (Binding { value = a }) _))
    | not $ any (pos `inside`) [ src' | Note src' _ <- toListOf subExpressions e ]
    = do _A <- first show $ typeWithA absurd ctx a
         srcAnnot <- case getLetAnnot src of
                       Just x -> return x
                       Nothing -> Left "The impossible happened: failed\
                                       \ to re-parse a Let expression."
         return (srcAnnot, normalize _A)
  
  -- binders, see typeAt'
  annotateLet' pos ctx (Let (Binding { variable = x, value = a }) e@(Note src _)) | pos `inside` src = do
    _ <- first show $ typeWithA absurd ctx a
    let a' = shift 1 (V x 0) (normalize a)
    annotateLet' pos ctx (shift (-1) (V x 0) (subst (V x 0) a' e))
  
  annotateLet' pos ctx (Lam _ FunctionBinding{ functionBindingVariable = x, functionBindingAnnotation = _A } b@(Note src _)) | pos `inside` src = do
    let _A' = Dhall.Core.normalize _A
        ctx' = fmap (shift 1 (V x 0)) (insert x _A' ctx)
    annotateLet' pos ctx' b
  
  annotateLet' pos ctx (Pi _ x _A _B@(Note src _)) | pos `inside` src = do
    let _A' = Dhall.Core.normalize _A
        ctx' = fmap (shift 1 (V x 0)) (insert x _A' ctx)
    annotateLet' pos ctx' _B
  
  -- we need to unfold Notes to make progress
  annotateLet' pos ctx (Note _ expr) =
    annotateLet' pos ctx expr
  
  -- catch-all
  annotateLet' pos ctx expr = do
    let subExprs = toListOf subExpressions expr
    case [ Note src e | (Note src e) <- subExprs, pos `inside` src ] of
      [e] -> annotateLet' pos ctx e
      _ -> Left "You weren't pointing at a let binder!"
  
  -- Make sure all lets in a multilet are annotated with their source information
  splitMultiLetSrc :: Expr Src a -> Maybe (Expr Src a)
  splitMultiLetSrc (Note src (Let b (Let b' e))) = do
    src' <- getLetInner src
    splitMultiLetSrc (Note src (Let b (Note src' (Let b' e))))
  splitMultiLetSrc expr = subExpressions splitMultiLetSrc expr
  
  -- Check if range lies completely inside a given subexpression.
  -- This version takes trailing whitespace into account
  -- (c.f. `sanitiseRange` from Backend.Diangostics).
  inside :: Position -> Src -> Bool
  inside pos src = left <= pos && pos < right
    where Range left right = rangeFromDhall src
  