-- | This module contains utilities for calculating positions and offsets. While
-- tokens are annotated with ranges, CST nodes are not, but they can be
-- dynamically derived with the functions in this module, which will return the
-- first and last tokens for a given node.

module Language.PureScript.CST.Positions where

import Prelude

import Data.Foldable (foldl')
import qualified Data.List.NonEmpty as NE
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Data.Void (Void)
import qualified Data.Text as Text
import Language.PureScript.CST.Types

advanceToken :: SourcePos -> Token -> SourcePos
advanceToken pos = applyDelta pos . tokenDelta

advanceLeading :: SourcePos -> [Comment LineFeed] -> SourcePos
advanceLeading = foldl' $ \a -> applyDelta a . commentDelta lineDelta

advanceTrailing :: SourcePos -> [Comment Void] -> SourcePos
advanceTrailing = foldl' $ \a -> applyDelta a . commentDelta (const (0, 0))

tokenDelta :: Token -> (Int, Int)
tokenDelta = \case
  TokLeftParen             -> (0, 1)
  TokRightParen            -> (0, 1)
  TokLeftBrace             -> (0, 1)
  TokRightBrace            -> (0, 1)
  TokLeftSquare            -> (0, 1)
  TokRightSquare           -> (0, 1)
  TokLeftArrow ASCII       -> (0, 2)
  TokLeftArrow Unicode     -> (0, 1)
  TokRightArrow ASCII      -> (0, 2)
  TokRightArrow Unicode    -> (0, 1)
  TokRightFatArrow ASCII   -> (0, 2)
  TokRightFatArrow Unicode -> (0, 1)
  TokDoubleColon ASCII     -> (0, 2)
  TokDoubleColon Unicode   -> (0, 1)
  TokForall ASCII          -> (0, 6)
  TokForall Unicode        -> (0, 1)
  TokEquals                -> (0, 1)
  TokPipe                  -> (0, 1)
  TokTick                  -> (0, 1)
  TokDot                   -> (0, 1)
  TokComma                 -> (0, 1)
  TokUnderscore            -> (0, 1)
  TokBackslash             -> (0, 1)
  TokLowerName qual name   -> (0, qualDelta qual + Text.length name)
  TokUpperName qual name   -> (0, qualDelta qual + Text.length name)
  TokOperator qual sym     -> (0, qualDelta qual + Text.length sym)
  TokSymbolName qual sym   -> (0, qualDelta qual + Text.length sym + 2)
  TokSymbolArr Unicode     -> (0, 3)
  TokSymbolArr ASCII       -> (0, 4)
  TokHole hole             -> (0, Text.length hole + 1)
  TokChar raw _            -> (0, Text.length raw + 2)
  TokInt raw _             -> (0, Text.length raw)
  TokNumber raw _          -> (0, Text.length raw)
  TokString raw _          -> multiLine 1 $ textDelta raw
  TokRawString raw         -> multiLine 3 $ textDelta raw
  TokLayoutStart           -> (0, 0)
  TokLayoutSep             -> (0, 0)
  TokLayoutEnd             -> (0, 0)
  TokEof                   -> (0, 0)

qualDelta :: [Text] -> Int
qualDelta = foldr ((+) . (+ 1) . Text.length) 0

multiLine :: Int -> (Int, Int) -> (Int, Int)
multiLine n (0, c) = (0, c + n + n)
multiLine n (l, c) = (l, c + n)

commentDelta :: (a -> (Int, Int)) -> Comment a -> (Int, Int)
commentDelta k = \case
  Comment raw -> textDelta raw
  Space n -> (0, n)
  Line a -> k a

lineDelta :: LineFeed -> (Int, Int)
lineDelta _ = (1, 1)

textDelta :: Text -> (Int, Int)
textDelta = Text.foldl' go (0, 0)
  where
  go (!l, !c) = \case
    '\n' -> (l + 1, 1)
    _    -> (l, c + 1)

applyDelta :: SourcePos -> (Int, Int) -> SourcePos
applyDelta (SourcePos l c) = \case
  (0, n) -> SourcePos l (c + n)
  (k, d) -> SourcePos (l + k) d

sepLast :: Separated a -> a
sepLast (Separated hd []) = hd
sepLast (Separated _ tl) = snd $ last tl

type TokenRange = (SourceToken, SourceToken)

toSourceRange :: TokenRange -> SourceRange
toSourceRange (a, b) = widen (srcRange a) (srcRange b)

widen :: SourceRange -> SourceRange -> SourceRange
widen (SourceRange s1 _) (SourceRange _ e2) = SourceRange s1 e2

srcRange :: SourceToken -> SourceRange
srcRange = tokRange . tokAnn

nameRange :: Name a -> TokenRange
nameRange a = (nameTok a, nameTok a)

qualRange :: QualifiedName a -> TokenRange
qualRange a = (qualTok a, qualTok a)

wrappedRange :: Wrapped a -> TokenRange
wrappedRange Wrapped { wrpOpen, wrpClose } = (wrpOpen, wrpClose)

moduleRange :: Module a -> TokenRange
moduleRange Module { modKeyword, modWhere, modImports, modDecls } =
  case (modImports, modDecls) of
    ([], []) -> (modKeyword, modWhere)
    (is, []) -> (modKeyword, snd . importDeclRange $ last is)
    (_,  ds) -> (modKeyword, snd . declRange $ last ds)

exportRange :: Export a -> TokenRange
exportRange = \case
  ExportValue _ a -> nameRange a
  ExportOp _ a -> nameRange a
  ExportType _ a b
    | Just b' <- b -> (nameTok a, snd $ dataMembersRange b')
    | otherwise -> nameRange a
  ExportTypeOp _ a b -> (a, nameTok b)
  ExportClass _ a b -> (a, nameTok b)
  ExportModule _ a b -> (a, nameTok b)

importDeclRange :: ImportDecl a -> TokenRange
importDeclRange ImportDecl { impKeyword, impModule, impNames, impQual }
  | Just (_, modName) <- impQual = (impKeyword, nameTok modName)
  | Just (_, imports) <- impNames = (impKeyword, wrpClose imports)
  | otherwise = (impKeyword, nameTok impModule)

importRange :: Import a -> TokenRange
importRange = \case
  ImportValue _ a -> nameRange a
  ImportOp _ a -> nameRange a
  ImportType _ a b
    | Just b' <- b -> (nameTok a, snd $ dataMembersRange b')
    | otherwise -> nameRange a
  ImportTypeOp _ a b -> (a, nameTok b)
  ImportClass _ a b -> (a, nameTok b)

dataMembersRange :: DataMembers a -> TokenRange
dataMembersRange = \case
  DataAll _ a -> (a, a)
  DataEnumerated _ (Wrapped a _ b) -> (a, b)

declRange :: Declaration a -> TokenRange
declRange = \case
  DeclData _ hd ctors
    | Just (_, cs) <- ctors -> (fst start, snd . dataCtorRange $ sepLast cs)
    | otherwise -> start
    where start = dataHeadRange hd
  DeclType _ a _ b -> (fst $ dataHeadRange a,  snd $ typeRange b)
  DeclNewtype _ a _ _ b -> (fst $ dataHeadRange a, snd $ typeRange b)
  DeclClass _ hd body
    | Just (_, ts) <- body -> (fst start, snd . typeRange . lblValue $ NE.last ts)
    | otherwise -> start
    where start = classHeadRange hd
  DeclInstanceChain _ a -> (fst . instanceRange $ sepHead a, snd . instanceRange $ sepLast a)
  DeclDerive _ a _ b -> (a, snd $ instanceHeadRange b)
  DeclKindSignature _ a (Labeled _ _ b) -> (a, snd $ typeRange b)
  DeclSignature _ (Labeled a _ b) -> (nameTok a, snd $ typeRange b)
  DeclValue _ a -> valueBindingFieldsRange a
  DeclFixity _ (FixityFields a _ (FixityValue _ _ b)) -> (fst a, nameTok b)
  DeclFixity _ (FixityFields a _ (FixityType _ _ _ b)) -> (fst a, nameTok b)
  DeclForeign _ a _ b -> (a, snd $ foreignRange b)
  DeclRole _ a _ _ b -> (a, roleTok $ NE.last b)

dataHeadRange :: DataHead a -> TokenRange
dataHeadRange (DataHead kw name vars)
  | [] <- vars = (kw, nameTok name)
  | otherwise = (kw, snd . typeVarBindingRange $ last vars)

dataCtorRange :: DataCtor a -> TokenRange
dataCtorRange (DataCtor _ name fields)
  | [] <- fields = nameRange name
  | otherwise = (nameTok name, snd . typeRange $ last fields)

classHeadRange :: ClassHead a -> TokenRange
classHeadRange (ClassHead kw _ name vars fdeps)
  | Just (_, fs) <- fdeps = (kw, snd .classFundepRange $ sepLast fs)
  | [] <- vars = (kw, snd $ nameRange name)
  | otherwise = (kw, snd . typeVarBindingRange $ last vars)

classFundepRange :: ClassFundep -> TokenRange
classFundepRange = \case
  FundepDetermined arr bs -> (arr, nameTok $ NE.last bs)
  FundepDetermines as _ bs -> (nameTok $ NE.head as, nameTok $ NE.last bs)

instanceRange :: Instance a -> TokenRange
instanceRange (Instance hd bd)
  | Just (_, ts) <- bd = (fst start, snd . instanceBindingRange $ NE.last ts)
  | otherwise = start
  where start = instanceHeadRange hd

instanceHeadRange :: InstanceHead a -> TokenRange
instanceHeadRange (InstanceHead kw _ _ cls types)
  | [] <- types = (kw, qualTok cls)
  | otherwise = (kw, snd . typeRange $ last types)

instanceBindingRange :: InstanceBinding a -> TokenRange
instanceBindingRange = \case
  InstanceBindingSignature _ (Labeled a _ b) -> (nameTok a, snd $ typeRange b)
  InstanceBindingName _ a -> valueBindingFieldsRange a

foreignRange :: Foreign a -> TokenRange
foreignRange = \case
  ForeignValue (Labeled a _ b) -> (nameTok a, snd $ typeRange b)
  ForeignData a (Labeled _ _ b) -> (a, snd $ typeRange b)
  ForeignKind a b -> (a, nameTok b)

valueBindingFieldsRange :: ValueBindingFields a -> TokenRange
valueBindingFieldsRange (ValueBindingFields a _ b) = (nameTok a, snd $ guardedRange b)

guardedRange :: Guarded a -> TokenRange
guardedRange = \case
  Unconditional a b -> (a, snd $ whereRange b)
  Guarded as -> (fst . guardedExprRange $ NE.head as, snd . guardedExprRange $ NE.last as)

guardedExprRange :: GuardedExpr a -> TokenRange
guardedExprRange (GuardedExpr a _ _ b) = (a, snd $ whereRange b)

whereRange :: Where a -> TokenRange
whereRange (Where a bs)
  | Just (_, ls) <- bs = (fst $ exprRange a, snd . letBindingRange $ NE.last ls)
  | otherwise = exprRange a

typeRange :: Type a -> TokenRange
typeRange = \case
  TypeVar _ a -> nameRange a
  TypeConstructor _ a -> qualRange a
  TypeWildcard _ a -> (a, a)
  TypeHole _ a -> nameRange a
  TypeString _ a _ -> (a, a)
  TypeInt _ a b _ -> (fromMaybe b a, b)
  TypeRow _ a -> wrappedRange a
  TypeRecord _ a -> wrappedRange a
  TypeForall _ a _ _ b -> (a, snd $ typeRange b)
  TypeKinded _ a _ b -> (fst $ typeRange a, snd $ typeRange b)
  TypeApp _ a b -> (fst $ typeRange a, snd $ typeRange b)
  TypeOp _ a _ b -> (fst $ typeRange a, snd $ typeRange b)
  TypeOpName _ a -> qualRange a
  TypeArr _ a _ b -> (fst $ typeRange a, snd $ typeRange b)
  TypeArrName _ a -> (a, a)
  TypeConstrained _ a _ b -> (fst $ constraintRange a, snd $ typeRange b)
  TypeParens _ a -> wrappedRange a
  TypeUnaryRow _ a b -> (a, snd $ typeRange b)

constraintRange :: Constraint a -> TokenRange
constraintRange = \case
  Constraint _ name args
    | [] <- args -> qualRange name
    | otherwise -> (qualTok name, snd . typeRange $ last args)
  ConstraintParens _ wrp -> wrappedRange wrp

typeVarBindingRange :: TypeVarBinding a -> TokenRange
typeVarBindingRange = \case
  TypeVarKinded a -> wrappedRange a
  TypeVarName a -> nameRange a

exprRange :: Expr a -> TokenRange
exprRange = \case
  ExprHole _ a -> nameRange a
  ExprSection _ a -> (a, a)
  ExprIdent _ a -> qualRange a
  ExprConstructor _ a -> qualRange a
  ExprBoolean _ a _ -> (a, a)
  ExprChar _ a _ -> (a, a)
  ExprString _ a _ -> (a, a)
  ExprNumber _ a _ -> (a, a)
  ExprArray _ a -> wrappedRange a
  ExprRecord _ a -> wrappedRange a
  ExprParens _ a -> wrappedRange a
  ExprTyped _ a _ b -> (fst $ exprRange a, snd $ typeRange b)
  ExprInfix _ a _ b -> (fst $ exprRange a, snd $ exprRange b)
  ExprOp _ a _ b -> (fst $ exprRange a, snd $ exprRange b)
  ExprOpName _ a -> qualRange a
  ExprNegate _ a b -> (a, snd $ exprRange b)
  ExprRecordAccessor _ (RecordAccessor a _ b) -> (fst $ exprRange a, lblTok $ sepLast b)
  ExprRecordUpdate _ a b -> (fst $ exprRange a, snd $ wrappedRange b)
  ExprApp _ a b -> (fst $ exprRange a, snd $ exprRange b)
  ExprLambda _ (Lambda a _ _ b) -> (a, snd $ exprRange b)
  ExprIf _ (IfThenElse a _ _ _ _ b) -> (a, snd $ exprRange b)
  ExprCase _ (CaseOf a _ _ c) -> (a, snd . guardedRange . snd $ NE.last c)
  ExprLet _ (LetIn a _ _ b) -> (a, snd $ exprRange b)
  ExprDo _ (DoBlock a b) -> (a,  snd . doStatementRange $ NE.last b)
  ExprAdo _ (AdoBlock a _ _ b) -> (a, snd $ exprRange b)

letBindingRange :: LetBinding a -> TokenRange
letBindingRange = \case
  LetBindingSignature _ (Labeled a _ b) -> (nameTok a, snd $ typeRange b)
  LetBindingName _ a -> valueBindingFieldsRange a
  LetBindingPattern _ a _ b -> (fst $ binderRange a, snd $ whereRange b)

doStatementRange :: DoStatement a -> TokenRange
doStatementRange = \case
  DoLet a bs -> (a, snd . letBindingRange $ NE.last bs)
  DoDiscard a -> exprRange a
  DoBind a _ b -> (fst $ binderRange a, snd $ exprRange b)

binderRange :: Binder a -> TokenRange
binderRange = \case
  BinderWildcard _ a -> (a, a)
  BinderVar _ a -> nameRange a
  BinderNamed _ a _ b -> (nameTok a, snd $ binderRange b)
  BinderConstructor _ a bs
    | [] <- bs -> qualRange a
    | otherwise -> (qualTok a, snd . binderRange $ last bs)
  BinderBoolean _ a _ -> (a, a)
  BinderChar _ a _ -> (a, a)
  BinderString _ a _ -> (a, a)
  BinderNumber _ a b _
    | Just a' <- a -> (a', b)
    | otherwise -> (b, b)
  BinderArray _ a -> wrappedRange a
  BinderRecord _ a -> wrappedRange a
  BinderParens _ a -> wrappedRange a
  BinderTyped _ a _ b -> (fst $ binderRange a, snd $ typeRange b)
  BinderOp _ a _ b -> (fst $ binderRange a, snd $ binderRange b)

recordUpdateRange :: RecordUpdate a -> TokenRange
recordUpdateRange = \case
  RecordUpdateLeaf a _ b -> (lblTok a, snd $ exprRange b)
  RecordUpdateBranch a (Wrapped _ _ b) -> (lblTok a, b)