import Data.Maybe
import Data.Functor
import Language.Haskell.Exts

fmap2 :: (Functor f, Functor g) => (a -> b) -> f (g a) -> f (g b)
fmap2 = fmap . fmap

infixl 4 <$$>
(<$$>) :: (Functor f, Functor g) => (a -> b) -> f (g a) -> f (g b)
(<$$>) = fmap2


parseFile' :: FilePath -> IO (Module SrcSpanInfo)
parseFile' f = do
  (ParseOk ast) <- parseFile f
  return ast

removeLoc :: Functor f => f SrcSpanInfo -> f ()
removeLoc = void


mapDecl :: (Decl l -> Decl l) -> Module l -> Module l
mapDecl f (Module l h p i ds) = Module l h p i $ map f ds

renamePatBinds :: Decl l -> Decl l
renamePatBinds (PatBind l p r b) = (PatBind l p' r b)
  where
    p' = case p of
           (PVar l (Ident l' i)) -> (PVar l (Ident l' ("orig_" ++ i)))
           a -> a
renamePatBinds d = d

pullTest :: Decl l -> Maybe (Name l)
pullTest (PatBind _ (PVar _ a@(Ident _ _)) (UnGuardedRhs _ (App _ (App _ (Var _ (UnQual _ (Ident _ "checkExpect"))) _) _)) _) = Just a
pullTest (PatBind _ (PVar _ a@(Ident _ _)) (UnGuardedRhs _ (InfixApp _ (App _ (Var _ (UnQual _ (Ident _ "checkExpect"))) _) (QVarOp _ (UnQual _ (Symbol _ "$"))) _)) _) = Just a
pullTest _ = Nothing

pullTests :: [Decl l] -> [Name l]
pullTests = catMaybes . map pullTest

pullTestsFromModule :: Module l -> [Name l]
pullTestsFromModule (Module _ _ _ _ ds) = pullTests ds