{-# LANGUAGE LambdaCase
           , ViewPatterns
           #-}
-- | Functions that convert the value and function definitions of the GHC AST to corresponding elements in the Haskell-tools AST representation
module Language.Haskell.Tools.AST.FromGHC.Binds where

import Control.Monad.Reader

import SrcLoc as GHC
import RdrName as GHC
import HsBinds as GHC
import HsExpr as GHC
import BasicTypes as GHC
import ApiAnnotation as GHC
import Bag as GHC
import Outputable as GHC
import HsPat as GHC
import HsDecls as GHC
import HsTypes as GHC
import OccName as GHC
import Name as GHC
import BooleanFormula as GHC

import Data.List

import Language.Haskell.Tools.AST.FromGHC.Base
import Language.Haskell.Tools.AST.FromGHC.Exprs
import Language.Haskell.Tools.AST.FromGHC.Patterns
import Language.Haskell.Tools.AST.FromGHC.Types
import Language.Haskell.Tools.AST.FromGHC.Kinds
import Language.Haskell.Tools.AST.FromGHC.Monad
import Language.Haskell.Tools.AST.FromGHC.Utils
import Language.Haskell.Tools.AST.FromGHC.GHCUtils

import Language.Haskell.Tools.AST (Ann(..), AnnMaybe(..), AnnList(..), Dom, RangeStage)
import qualified Language.Haskell.Tools.AST as AST

trfBind :: TransformName n r => Located (HsBind n) -> Trf (Ann AST.ValueBind (Dom r) RangeStage)
trfBind = trfLocNoSema trfBind'
  
trfBind' :: TransformName n r => HsBind n -> Trf (AST.ValueBind (Dom r) RangeStage)
-- a value binding (not a function)
trfBind' (FunBind { fun_id = id, fun_matches = MG { mg_alts = unLoc -> [L matchLoc (Match { m_pats = [], m_grhss = GRHSs [L rhsLoc (GRHS [] expr)] (unLoc -> locals) })]} }) 
  = AST.SimpleBind <$> copyAnnot AST.VarPat (define $ trfName id)
                   <*> addToScope locals (annLocNoSema (combineSrcSpans (getLoc expr) <$> tokenLoc AnnEqual) (AST.UnguardedRhs <$> trfExpr expr))
                   <*> addToScope locals (trfWhereLocalBinds locals)
trfBind' (FunBind id (MG (unLoc -> matches) _ _ _) _ _ _) = AST.FunBind <$> makeNonemptyIndentedList (mapM (trfMatch (unLoc id)) matches)
trfBind' (PatBind pat (GRHSs rhs (unLoc -> locals)) _ _ _) = AST.SimpleBind <$> trfPattern pat <*> trfRhss rhs <*> trfWhereLocalBinds locals
trfBind' (AbsBinds _ _ _ _ _) = error "AbsBinds are not allowed as an input to the conversion (they are generated by the type checker)"
trfBind' (PatSynBind _) = error "Pattern synonym bindings should be recognized on the declaration level"

trfMatch :: TransformName n r => n -> Located (Match n (LHsExpr n)) -> Trf (Ann AST.Match (Dom r) RangeStage)
trfMatch id = trfLocNoSema (trfMatch' id)

trfMatch' :: TransformName n r => n -> Match n (LHsExpr n) -> Trf (AST.Match (Dom r) RangeStage)
trfMatch' name (Match funid pats typ (GRHSs rhss (unLoc -> locBinds)))
  -- TODO: add the optional typ to pats
  = AST.Match <$> trfMatchLhs name funid pats
              <*> addToScope pats (trfRhss rhss)
              <*> addToScope pats (trfWhereLocalBinds locBinds)

trfMatchLhs :: TransformName n r => n -> MatchFixity n -> [LPat n] -> Trf (Ann AST.MatchLhs (Dom r) RangeStage)
trfMatchLhs name fb pats 
  = do implicitIdLoc <- mkSrcSpan <$> atTheStart <*> atTheStart
       closeLoc <- srcSpanStart <$> (combineSrcSpans <$> tokenLoc AnnEqual <*> tokenLoc AnnVbar)
       let (n, isInfix) = case fb of NonFunBindMatch -> (L implicitIdLoc name, False)
                                     FunBindMatch n inf -> (n, inf)
       args <- mapM trfPattern pats
       annLocNoSema (mkSrcSpan <$> atTheStart <*> (pure closeLoc)) $
         case (args, isInfix) of 
            (left:right:rest, True) -> AST.InfixLhs left <$> define (trfOperator n) <*> pure right <*> makeList " " (pure closeLoc) (pure rest)
            _                       -> AST.NormalLhs <$> define (trfName n) <*> makeList " " (pure closeLoc) (pure args)

trfRhss :: TransformName n r => [Located (GRHS n (LHsExpr n))] -> Trf (Ann AST.Rhs (Dom r) RangeStage)
-- the original location on the GRHS misleadingly contains the local bindings
trfRhss [unLoc -> GRHS [] body] = annLocNoSema (combineSrcSpans (getLoc body) <$> tokenBefore (srcSpanStart $ getLoc body) AnnEqual) 
                                         (AST.UnguardedRhs <$> trfExpr body)
trfRhss rhss = annLocNoSema (pure $ collectLocs rhss) 
                      (AST.GuardedRhss . nonemptyAnnList <$> mapM trfGuardedRhs rhss)
                      
trfGuardedRhs :: TransformName n r => Located (GRHS n (LHsExpr n)) -> Trf (Ann AST.GuardedRhs (Dom r) RangeStage)
trfGuardedRhs = trfLocNoSema $ \(GRHS guards body) 
  -> AST.GuardedRhs . nonemptyAnnList <$> trfScopedSequence trfRhsGuard guards <*> addToScope guards (trfExpr body)
  
trfRhsGuard :: TransformName n r => Located (Stmt n (LHsExpr n)) -> Trf (Ann AST.RhsGuard (Dom r) RangeStage)
trfRhsGuard = trfLocNoSema trfRhsGuard'
  
trfRhsGuard' :: TransformName n r => Stmt n (LHsExpr n) -> Trf (AST.RhsGuard (Dom r) RangeStage)
trfRhsGuard' (BindStmt pat body _ _ _) = AST.GuardBind <$> trfPattern pat <*> trfExpr body
trfRhsGuard' (BodyStmt body _ _ _) = AST.GuardCheck <$> trfExpr body
trfRhsGuard' (LetStmt (unLoc -> binds)) = AST.GuardLet <$> trfLocalBinds binds
  
trfWhereLocalBinds :: TransformName n r => HsLocalBinds n -> Trf (AnnMaybe AST.LocalBinds (Dom r) RangeStage)
trfWhereLocalBinds EmptyLocalBinds = nothing "" "" atTheEnd
trfWhereLocalBinds binds
  = makeJust <$> annLocNoSema (combineSrcSpans (getBindLocs binds) <$> tokenLoc AnnWhere) (AST.LocalBinds <$> addToScope binds (trfLocalBinds binds))

getBindLocs :: HsLocalBinds n -> SrcSpan
getBindLocs (HsValBinds (ValBindsIn binds sigs)) = foldLocs $ map getLoc (bagToList binds) ++ map getLoc sigs
getBindLocs (HsValBinds (ValBindsOut binds sigs)) = foldLocs $ map getLoc (concatMap (bagToList . snd) binds) ++ map getLoc sigs
getBindLocs (HsIPBinds (IPBinds binds _)) = foldLocs $ map getLoc binds
  
trfLocalBinds :: TransformName n r => HsLocalBinds n -> Trf (AnnList AST.LocalBind (Dom r) RangeStage)
trfLocalBinds (HsValBinds (ValBindsIn binds sigs)) 
  = makeIndentedList (after AnnWhere)
      (orderDefs <$> ((++) <$> mapM (copyAnnot AST.LocalValBind . trfBind) (bagToList binds) 
                           <*> mapM trfLocalSig sigs))
trfLocalBinds (HsValBinds (ValBindsOut binds sigs)) 
  = makeIndentedList (after AnnWhere)
      (orderDefs <$> ((++) <$> (concat <$> mapM (mapM (copyAnnot AST.LocalValBind . trfBind) . bagToList . snd) binds)
                           <*> mapM trfLocalSig sigs))
trfLocalBinds (HsIPBinds (IPBinds binds _))
  = makeIndentedList (after AnnWhere) (mapM trfIpBind binds)

trfIpBind :: TransformName n r => Located (IPBind n) -> Trf (Ann AST.LocalBind (Dom r) RangeStage)
trfIpBind = trfLocNoSema $ \case
  IPBind (Left (L l ipname)) expr -> AST.LocalValBind <$> (annContNoSema $ AST.SimpleBind <$> focusOn l (annContNoSema (AST.VarPat <$> define (trfImplicitName ipname)))
                                                                                          <*> annFromNoSema AnnEqual (AST.UnguardedRhs <$> trfExpr expr)
                                                                                          <*> nothing " " "" atTheEnd)
             
trfLocalSig :: TransformName n r => Located (Sig n) -> Trf (Ann AST.LocalBind (Dom r) RangeStage)
trfLocalSig = trfLocNoSema $ \case
  ts@(TypeSig {}) -> AST.LocalSignature <$> annContNoSema (trfTypeSig' ts)
  (FixSig fs) -> AST.LocalFixity <$> annContNoSema (trfFixitySig fs)
  
trfTypeSig :: TransformName n r => Located (Sig n) -> Trf (Ann AST.TypeSignature (Dom r) RangeStage)
trfTypeSig = trfLocNoSema trfTypeSig'

trfTypeSig' :: TransformName n r => Sig n -> Trf (AST.TypeSignature (Dom r) RangeStage)
trfTypeSig' (TypeSig names typ) 
  = defineTypeVars $ AST.TypeSignature <$> makeNonemptyList ", " (mapM trfName names) <*> trfType (hswc_body $ hsib_body typ)
  
trfFixitySig :: TransformName n r => FixitySig n -> Trf (AST.FixitySignature (Dom r) RangeStage)
trfFixitySig (FixitySig names (Fixity _ prec dir)) 
  = AST.FixitySignature <$> transformDir dir
                        <*> annLocNoSema (tokenLoc AnnVal) (pure $ AST.Precedence prec) 
                        <*> (nonemptyAnnList . nub <$> mapM trfOperator names)
  where transformDir InfixL = directionChar (pure AST.AssocLeft)
        transformDir InfixR = directionChar (pure AST.AssocRight)
        transformDir InfixN = annLocNoSema (srcLocSpan . srcSpanEnd <$> tokenLoc AnnInfix) (pure AST.AssocNone)
        
        directionChar = annLocNoSema ((\l -> mkSrcSpan (updateCol (subtract 1) l) l) . srcSpanEnd <$> tokenLoc AnnInfix)
   
trfRewriteRule :: TransformName n r => Located (RuleDecl n) -> Trf (Ann AST.Rule (Dom r) RangeStage)
trfRewriteRule = trfLocNoSema $ \(HsRule (L nameLoc (_, ruleName)) act bndrs left _ right _) ->
  AST.Rule <$> trfFastString (L nameLoc ruleName) 
           <*> trfPhase (before AnnForall) act
           <*> makeNonemptyList " " (mapM trfRuleBndr bndrs)
           <*> trfExpr left
           <*> trfExpr right

trfRuleBndr :: TransformName n r =>  Located (RuleBndr n) -> Trf (Ann AST.TyVar (Dom r) RangeStage)
trfRuleBndr = trfLocNoSema $ \case (RuleBndr n) -> AST.TyVarDecl <$> trfName n <*> nothing " " "" atTheEnd
                                   (RuleBndrSig n k) -> AST.TyVarDecl <$> trfName n <*> (makeJust <$> (trfKindSig' (hswc_body $ hsib_body k)))

trfMinimalFormula :: TransformName n r => Located (BooleanFormula (Located n)) -> Trf (Ann AST.MinimalFormula (Dom r) RangeStage)
trfMinimalFormula = trfLocNoSema trfMinimalFormula'

trfMinimalFormula' :: TransformName n r => BooleanFormula (Located n) -> Trf (AST.MinimalFormula (Dom r) RangeStage)
trfMinimalFormula' (Var name) = AST.MinimalName <$> trfName name
trfMinimalFormula' (And formulas) = AST.MinimalAnd <$> trfAnnList " & " trfMinimalFormula' formulas
trfMinimalFormula' (Or formulas) = AST.MinimalOr <$> trfAnnList " | " trfMinimalFormula' formulas
trfMinimalFormula' (Parens formula) = AST.MinimalParen <$> trfMinimalFormula formula