{-# LANGUAGE TemplateHaskell, DeriveLift, StandaloneDeriving #-}

module Mooc.Th (testing, testing', timeLimit,
                isDefined, withDefined, hasType, hasType', importsOnly, show',
                reifyType, DataType(..), FieldType(..), Constructor(..), withConstructor,
                withInstance, withInstanceType, withInstanceSilent, withInstance1, withInstances1,
                classContains, defineInstance)
where

import Data.Char
import Data.List
import Data.Maybe
import Language.Haskell.TH hiding (reifyType)
import Language.Haskell.TH.Syntax hiding (reifyType)
import Language.Haskell.TH.Datatype (tvName)

-- testing presence of definitions

isDefined :: String -> Q Exp
isDefined s = do
  mn <- lookupValueName s
  let defined = isJust mn
  [|counterexample ("You haven't defined '"++s++"' yet!") defined|]

varOrCon :: Name -> Q Exp
varOrCon n@(Name (OccName s) _)
  | isUpper (head s) = return $ ConE n
  | otherwise = return $ VarE n

withDefined :: String -> Q Exp
withDefined s = do
  mn <- lookupValueName s
  case mn of
    Nothing -> [|(\_ -> counterexample ("You haven't defined '"++s++"' yet!") False)|]
    Just n -> [|(\f -> f $(varOrCon n))|]

getType :: String -> Q (Maybe Type)
getType s = do
  mn <- lookupValueName s
  case mn of Nothing -> return Nothing
             Just n -> do i <- reify n
                          case i of VarI _ typ _ -> return (Just typ)
                                    DataConI _ typ _ -> return (Just typ)
                                    _ -> return Nothing

showType :: Type -> String
showType (ConT name) = nameBase name
showType (AppT (AppT ArrowT arg) typ) = showType arg ++ " -> " ++ showType typ
showType (AppT ListT typ) = "[" ++ showType typ ++ "]"
showType (AppT t1 t2) = showType t1 ++ " " ++ showType t2  -- TODO doesn't handle parens
showType (ForallT _ _ t) = showType t -- TODO hide foralls
showType (VarT name) = nameBase name
showType t = show t
-- TODO: more cases as needed

hasType :: String -> Q Type -> Q Exp
hasType s qtyp = do
  mt <- getType s
  typ <- qtyp
  let correct = mt == Just typ
  let expected = showType typ
  let actual = case mt of Nothing -> "nothing"
                          Just t -> showType t
  if correct
    then [|\k -> counterexample s (k $(varOrCon $ mkName s))|]
    else [|\k -> counterexample ("The type of '"++s++"'\n  Expected: "++expected++"\n  Was: "++actual) False|]

hasType' :: String -> String -> Q Exp
hasType' s expected = do
  mt <- getType s
  let actual = case mt of Nothing -> "nothing"
                          Just t -> showType t
  if (actual==expected)
    then [|\k -> counterexample s (k $(varOrCon $ mkName s))|]
    else [|\k -> counterexample ("The type of '"++s++"'\n  Expected: "++expected++"\n  Was: "++actual) False|]

-- testing types

data FieldType = SimpleType String | WeirdType
  deriving (Eq,Show,Lift)

interpretFieldType (ConT n) = SimpleType $ nameBase n
interpretFieldType (VarT n) = SimpleType $ nameBase n
interpretFieldType (AppT ListT t) =
  case interpretFieldType t of
    SimpleType s -> SimpleType $ "[" ++ s ++ "]"
    _ -> WeirdType
interpretFieldType (AppT t1 t2) =
  case (interpretFieldType t1, interpretFieldType t2) of
    (SimpleType s1, SimpleType s2) -> SimpleType $ s1 ++ " " ++ s2
    _ -> WeirdType
interpretFieldType _ = WeirdType

data Constructor = Constructor String [FieldType] | Weird String
  deriving (Eq,Show,Lift)

interpretConstructor (NormalC n bts) = Constructor (nameBase n) args
  where (_,typs) = unzip bts
        args = map interpretFieldType typs
interpretConstructor c = Weird (constructorName c)

data DataType = DataType [String] [Constructor]
  deriving (Eq,Show,Lift)

lookupType' s =
    do n <- lookupTypeName s
       case n of Nothing -> return Nothing
                 Just n -> do info <- reify n
                              return $ case info of
                                         TyConI (DataD _ _ vs _ cs _) -> let cons = map interpretConstructor cs
                                                                             vars = map bndrName vs
                                                                         in Just (DataType vars cons)
                                         _ -> Nothing

reifyType :: String -> Q Exp
reifyType s = do
  info <- lookupType' s
  case info of
    Nothing -> [|\k -> counterexample ("Type "++s++" not defined!") False|]
    Just d -> [|\k -> counterexample ("The type "++s) $ k d|]

bndrName = nameBase . tvName

constructorName = nameBase . constructorName'

constructorName' :: Con -> Name
constructorName' c = case c of
  NormalC n _ -> n
  RecC n _ -> n
  InfixC _ n _ -> n
  ForallC _ _ c -> constructorName' c
  GadtC (n:_) _ _ -> n
  RecGadtC (n:_) _ _ -> n

withConstructor :: String -> String -> [String] -> Q Exp
withConstructor typName consName argTypes = do
  d <- lookupType' typName
  case d of
    Nothing -> [|\k -> counterexample ("Type "++typName++" not defined!") False|]
    Just (DataType _ cs) ->
        case filter ok cs of
          [] -> [|\k -> counterexample ("Type "++typName++" should have a constructor named "++consName) False|]
          [Weird _] -> [|\k -> counterexample ("Constructor "++consName++" of type "++typName++" is weird. Make it normal.") False|]
          [Constructor _ ts]
              | ts == map SimpleType argTypes -> [|\k -> counterexample consName (k $(varOrCon $ mkName consName))|]
              | otherwise -> [|\k -> counterexample ("Constructor "++consName++" of type "++typName++" should have fields of types "++intercalate ", " argTypes) False |]
    where ok (Constructor n _) = n == consName
          ok (Weird n) = n == consName

-- testing classes

data Instance = Found | NotFound String String | NoClass String | NoType String
  deriving (Show, Eq)

lookupInstanceResolved :: Name -> Type -> Q Instance
lookupInstanceResolved cl t = do
  b <- isInstance cl [t]
  if b
    then return Found
    else return $ NotFound (nameBase cl) (showType t)

lookupInstance :: String -> String -> Q Instance
lookupInstance cln typn = do
  cl <- lookupTypeName cln
  typ <- lookupTypeName typn
  case (cl,typ) of
    (Nothing,_) -> return $ NoClass cln
    (_,Nothing) -> return $ NoType typn
    (Just c, Just t) -> lookupInstanceResolved c (ConT t)

withInstance :: String -> String -> Q Exp -> Q Exp
withInstance cln typn val = do
  ins <- lookupInstance cln typn
  case ins of
    NoClass cln -> [|\k -> counterexample ("Class "++cln++" not found") (property False)|]
    NoType typn -> [|\k -> counterexample ("Type "++typn++" not found") (property False)|]
    NotFound cln typn -> [|\k -> counterexample ("Type "++typn++" is not an instance of class "++cln) (property False)|]
    Found -> [|\k -> k $val|]

withInstanceType :: String -> Q Type -> Q Exp -> Q Exp
withInstanceType cln qt val = do
  cl <- lookupTypeName cln
  t <- qt
  case cl of
    Nothing -> [|\k -> counterexample ("Class "++cln++" not found") (property False)|]
    Just cl -> do ins <- lookupInstanceResolved cl t
                  case ins of
                    NotFound cln typn -> [|\k -> counterexample ("Type "++typn++" is not an instance of class "++cln) (property False)|]
                    Found -> [|\k -> k $val|]

withInstanceSilent :: String -> String -> Q Exp -> String -> Q Exp
withInstanceSilent cln typn val err = do
  ins <- lookupInstance cln typn
  case ins of
    NoClass cln -> [|\_ -> counterexample err (property False)|]
    NoType typn -> [|\_ -> counterexample err (property False)|]
    NotFound cln typn -> [|\_ -> counterexample err (property False)|]
    Found -> [|\k -> k $val|]

classContains :: String -> String -> Q Exp
classContains cln varn = do
  var <- lookupValueName varn
  case var of
    Nothing -> [|counterexample ("Function "++varn++" not found") (property False)|]
    Just cl -> do
      info <- reify cl
      case info of
        (ClassOpI _ _ parent) -> let pn = nameBase parent in [|counterexample ("Function "++varn++" is in the wrong class!") (pn ?== cln)|]
        _ -> [|counterexample ("Function "++varn++" is not a method of class "++cln) (property False)|]

lookupConstructor :: String -> Q (Maybe Type)
lookupConstructor "[]" = return (Just ListT)
lookupConstructor x = (fmap.fmap) ConT $ lookupTypeName x

lookupInstance1 :: String -> String -> Q Instance
lookupInstance1 cln typn = do
  cl <- lookupTypeName cln
  cons <- lookupConstructor typn
  case (cl,cons) of
    (Nothing,_) -> return $ NoClass cln
    (_,Nothing) -> return $ NoType typn
    (Just c, Just t) -> lookupInstanceResolved c (AppT t (VarT (mkName "a")))

withInstance1 :: String -> String -> Q Exp -> Q Exp
withInstance1 cln typn val = do
  ins <- lookupInstance1 cln typn
  case ins of
    NoClass cln -> [|\k -> counterexample ("Class "++cln++" not found") (property False)|]
    NoType typn -> [|\k -> counterexample ("Type "++typn++" not found") (property False)|]
    NotFound cln typn -> [|\k -> counterexample ("Type "++typn++" is not an instance of class "++cln) (property False)|]
    Found -> [|\k -> k $val|]

withInstances1 :: String -> [String] -> Q Exp -> Q Exp
withInstances1 cln typns val = do
  instances <- mapM (lookupInstance1 cln) typns
  case filter (/=Found) instances of
    (NoClass cln:_) -> [|\k -> counterexample ("Class "++cln++" not found") (property False)|]
    (NoType typn:_) -> [|\k -> counterexample ("Type "++typn++" not found") (property False)|]
    (NotFound cln typn:_) -> [|\k -> counterexample ("Type "++typn++" is not an instance of class "++cln) (property False)|]
    _ -> [|\k -> k $val|]

defineInstance :: String -> Name -> String -> Q Exp -> Q [Dec]
defineInstance cln typn methodn body = do
  cl <- lookupTypeName cln
  --let method = Name (OccName methodn) NameS
  case cl of
    Nothing -> return []
    (Just c) -> sequence [instanceD (cxt []) (appT (conT c) (conT typn)) [funD (mkName methodn) [clause [] (normalB body) []]]]

-- verbose testing

data Call = Prefix Name [Exp] | Operator Name Exp Exp

parseCall :: Exp -> Call
parseCall (VarE name) = Prefix name []
parseCall (UnboundVarE name) = Prefix name []
parseCall (AppE f arg) = Prefix name (args ++ [arg])
  where Prefix name args = parseCall f
parseCall (InfixE (Just l) (VarE op) (Just r)) = Operator op l r
parseCall x = error $ "Unsupported: " ++ show x

show' v = showsPrec 11 v ""

showArgs :: Call -> Q Exp
showArgs (Prefix name args) = [|fstring ++ " " ++ intercalate " " $arglist|]
  where fstring = nameBase name
        arglist = foldr (\arg exp -> [| show' $arg : $exp |]) [| [] |] (map return args)
showArgs (Operator name left right) = [| show' $(return left) ++ " " ++ opstring ++ " " ++ show' $(return right) |]
  where opstring = nameBase name

timeLimit :: Int
timeLimit = 5 * 1000 * 1000 -- 5 seconds in microseconds, less than the toplevel timeout of 10 seconds

testing :: Q Exp -> Q Exp
testing call = do
  parsed <- fmap parseCall call
  [| \k -> counterexample $(showArgs parsed) (within timeLimit (k $call)) |]

-- TH pprint prints all names as qualified, let's convert the names to unqualified locals
unqualifyName (Name n _) = Name n NameS

unqualifyType :: Type -> Type
unqualifyType (AppT f x) = AppT (unqualifyType f) (unqualifyType x)
unqualifyType (ConT n) = ConT (unqualifyName n)
unqualifyType x = error $ "Unsupported: " ++ show x

unqualify :: Exp -> Exp
unqualify (VarE n) = VarE (unqualifyName n)
unqualify (UnboundVarE n) = UnboundVarE (unqualifyName n)
unqualify (ConE n) = ConE (unqualifyName n)
unqualify (AppE f x) = AppE (unqualify f) (unqualify x)
unqualify (InfixE mleft op mright) = InfixE (fmap unqualify mleft) (unqualify op) (fmap unqualify mright)
unqualify (UInfixE left op right) = UInfixE (unqualify left) (unqualify op) (unqualify right)
unqualify (LitE l) = LitE l
unqualify (TupE exps) = TupE (map (fmap unqualify) exps)
unqualify (ListE exps) = ListE (map unqualify exps)
unqualify (ArithSeqE (FromR x)) = ArithSeqE (FromR (unqualify x))
unqualify (SigE e t) = SigE (unqualify e) (unqualifyType t)
unqualify x = error $ "Unsupported: " ++ show x

testing' :: Q Exp -> Q Exp
testing' call = do
  str <- fmap (pprint.unqualify) call
  --str <- fmap show call
  [| \k -> counterexample str (within timeLimit (k $call)) |]

-- checking imports

importsOnly :: String -> [String] -> Q Exp
importsOnly name allowed = do
  (ModuleInfo pkgImports) <- reifyModule (Module (PkgName "main") (ModName name))
  let imports = [nam | (Module _ (ModName nam)) <- pkgImports]
  let forbidden = imports \\ allowed
  case forbidden of
    [] -> [|property True|]
    xs -> [|counterexample ("forbidden imports: " ++ show xs) False|]