{-# LANGUAGE TemplateHaskell #-}

module Test.OITestGenerator (
    AxiomResult, (===>), (=!=),
    Axiom, axiom,
    Arg, Op, op, arg, but, only, withConstraint,
    withGens,
    generate_oi_tests, generate_axiom's_tests, generate_single_test,
    generate_basic_tests,
    show_all_tests,
) where

import Control.Applicative
import Control.Monad
import Data.List
import Language.Haskell.TH
import Prelude hiding (exp)
import Test.OITestGenerator.Axiom
import Test.OITestGenerator.GenHelper
import Test.OITestGenerator.Op
import Test.OITestGenerator.OpArg
import Test.QuickCheck

{- Functions to export -}

-- | Derives one test from each 'Axiom' by translating it to a QuickCheck
-- 'Property'. This is done basically by replacing '=!=' with '==' and '===>'
-- with 'QuickCheck.==>'.
--
-- It returns an expression of type @[Property]@, i.e.
--
-- > $(generate_basic_tests axs) :: [Property]
generate_basic_tests :: [Axiom] -> ExpQ
generate_basic_tests = liftM ListE . mapM gen_basic_test

-- | Generate all operation invariance tests that typecheck. For example, given
-- an operation
--
-- > f :: A -> A -> B -> C
--
-- and an axiom
--
-- > ax :: D -> AxiomResult A
-- > ax x = lhs x =!= rhs x
--
-- the following tests are generated:
--
-- > property $ \y2 y3 -> f (lhs x) y2 y3 == f (rhs x) y2 y3
-- > property $ \y1 y3 -> f y1 (lhs x) y3 == f y1 (rhs x) y3
--
-- but not
--
-- > property $ \y1 y2 -> f y1 y2 (lhs x) == f y1 y2 (rhs x)
--
-- because of the types.
--
-- All combinations of operations, their respective arguments and axioms are tried.
--
-- It returns an expression of type @[Property]@, i.e.
--
-- > $(generate_oi_tests axs ops) :: [Property]
generate_oi_tests :: [Axiom] -> [Op] -> ExpQ
generate_oi_tests = gen_tests

-- | A convenience function for calling 'generate_oi_tests' with only one axiom.
--
-- It returns an expression of type @[Property]@, i.e.
--
-- > $(generate_axiom's_tests ax ops) :: [Property]
generate_axiom's_tests :: Axiom -> [Op] -> ExpQ
generate_axiom's_tests ax = gen_tests [ax]

-- | Calls 'generate_oi_tests' with only the given 'Axiom' and 'Op'. It checks
-- if there is exactly one test that can be generated and throws an error
-- otherwise. If multiple applications of the given operation to the given
-- axioms are possible (i.e. it has multiple arguments with a matching type),
-- the allowed arguments have to be restricted using 'but' or 'only'.
--
-- It returns an expression of type @Property@, i.e.
--
-- > $(generate_single_test ax op) :: Property
generate_single_test :: Axiom -> Op -> ExpQ
generate_single_test ax opn = op2opArgs opn
    >>= filter_Ops ax
    >>= \args ->
        case args of
            [op_arg] -> generate_test ax op_arg
            _        -> do let axnm = axiom_name ax
                           let opnm = op_name opn
                           axtp <- return_type_of_axiom axnm
                           opargs <- op_args opn
                           oparg_types <- liftM (zip opargs) $ mapM (flip type_of_argi opnm) opargs
                           fail $ "Invalid number of feasible Op arguments. "
                            ++ "Expected 1, but got " ++ show (length args)
                            ++ ( if (length args) == 0
                                 then ""
                                 else ", namely " ++ show (map opArg_argi args)
                               )
                            ++ "."
                            ++ "\n"
                            ++ "Type of axiom " ++ show axnm ++ "'s result:\n"
                            ++ "         " ++ pprint axtp ++ "\n"
                            ++ "Type of operation arguments " ++ show opnm ++ ":\n"
                            ++ (concatMap (\(i,tp) -> "  Arg#" ++ show i ++ ": " ++ pprint tp ++ "\n") oparg_types) ++ "\n" 

-- | @show_all_tests Nothing axs ops@ returns a 'String' typed expression. Print
-- the resulting string to get a piece of code, containing one property for each
-- possible operation invariance test, defined via 'generate_single_test'.
--
-- The first argument can replace the default function generating a name for
-- each test, which is
--
-- > \opn argi ax -> opn ++ show argi ++ "_" ++ ax
--
-- . An example output looks like this:
--
-- > enqueue1_q3 :: Property
-- > enqueue1_q3 = $(generate_single_test (axiom 'q3) (op 'enqueue `only` 1))
show_all_tests :: Maybe (String -> Int -> String -> String) -> [Name] -> [Name] -> ExpQ
show_all_tests Nothing axs ops = show_all_tests (Just def_test_name) axs ops
show_all_tests (Just fun) axs ops =
    cross_ops_axs ops axs
    >>= return . concatMap (\(opn, argi, ax) ->
        let name = fun (nameBase opn) argi (nameBase ax)
        in  name ++ " :: Property\n" ++
            name ++ " = " ++ "$(" ++ intercalate " "
            ["generate_single_test",
            "(axiom '" ++ nameBase ax ++ ")",
            "(op '" ++ nameBase opn ++ " `only` " ++ show argi ++ ")"]
            ++ ")\n\n"
    )
    >>= stringE

{- Private functions -}

simple_dec :: String -> Exp -> Dec
simple_dec nm exp = ValD (VarP (mkName nm)) (NormalB exp) []

def_test_name :: String -> Int -> String -> String
def_test_name opn argi ax = opn ++ show argi ++ "_" ++ ax

gen_label :: Int -> Name -> Name -> String
gen_label argi opn ax = show opn ++ "@" ++ show argi ++ "/" ++ show ax

concatMapM :: (Monad m) => (a -> m [b]) -> [a] -> m [b]
concatMapM f = liftM concat . mapM f

(.>) :: (a -> b) -> (b -> c) -> a -> c
(.>) = flip (.)

gen_tests_per_ax :: Axiom -> [Op] -> Q [Exp]
gen_tests_per_ax ax ops =
        cross_Ops_args ops
        >>= filter_Ops ax
        >>= mapM (generate_test ax)

gen_tests :: [Axiom] -> [Op] -> ExpQ
gen_tests axs ops = concatMapM (flip gen_tests_per_ax ops) axs
                    >>= return . ListE

cross_ops_axs :: [Name] -> [Name] -> Q [(Name, Int, Name)]
cross_ops_axs ops axs = concatMapM (\ax
    -> cross_ops_args ops
    >>= filter_ops ax
    >>= mapM (\(argi, opn) -> return (opn, argi, ax))
    ) axs

cross_ops_args :: [Name] -> Q [(Int, Name)]
cross_ops_args = concatMapM cross_op_args

cross_op_args :: Name -> Q [(Int, Name)]
cross_op_args opn = typeof opn >>= \tp -> return $ zip (enumFromTo 1 $ num_args tp) (repeat opn)

cross_Ops_args :: [Op] -> Q [OpArg]
cross_Ops_args = concatMapM cross_Op_args

cross_Op_args :: Op -> Q [OpArg]
cross_Op_args = op2opArgs

delete_i :: Int -> [a] -> [a]
delete_i i xs = take i xs ++ drop (i+1) xs

apply_gens_to' :: ExpQ -> [Name] -> Int -> ExpQ
apply_gens_to' prop [] _ = prop
apply_gens_to' prop (x:xs) n = [| $app (forAll $(varE x)) $(apply_gens_to' prop xs (n+1)) |]
    where app = liftAn n

apply_gens_to :: ExpQ -> [Name] -> ExpQ
apply_gens_to prop xs = apply_gens_to' prop xs 0

gen_basic_test :: Axiom -> ExpQ
gen_basic_test ax =
    let axnm        = axiom_name ax
        test_name   = show axnm
        axE         = varE axnm
        app         = num_args_name axnm >>= liftAn
        app2        = num_args_name axnm >>= liftA2n
    in [|
            let  lhs         = $app ar_lhs   $axE
                 rhs         = $app ar_rhs   $axE
                 axiom_cond  = $app ar_cond  $axE
            in  label test_name $ property $
                $app2 (==>) axiom_cond $
                $app2 (==) lhs rhs
        |]

generate_test :: Axiom -> OpArg -> ExpQ
generate_test ax opn =
    do
        let axnm = axiom_name ax
        let opnm = opArg_name opn
        let argi = opArg_argi opn
        argn_ax <- num_args_name axnm
        argn_op <- num_args_name opnm
        let axE = varE axnm
        let opE = varE opnm
        ax_custom_gens <- axiom_gens ax
        let op_custom_gens = opArg_gens opn
        let ax_gens = if length ax_custom_gens > 0
                      then ax_custom_gens
                      else replicate argn_ax 'arbitrary
        let opn_gens = if length op_custom_gens > 0
                       then delete_i (argi - 1) op_custom_gens
                       else replicate (argn_op - 1) 'arbitrary
        let has_constr = opArg_has_constraint opn
        let op_constr = if has_constr
                        then varE $ opArg_constraint opn
                        else [| $(num_args_name opnm >>= constN) True |]
        let all_gens = ax_gens ++ opn_gens
        let fargi = flip_n_to_1 argi
        let app = liftAn argn_ax
        let app2 = liftA2n (argn_ax + argn_op - 1)
        let eat_spare_cond_args = eat_args (argn_ax+1) (argn_ax + argn_op - 1)
        let test_name = gen_label argi opnm axnm
        [|
            let lhs = $app ar_lhs $axE
                rhs = $app ar_rhs $axE
                app_op_constraint = $app ($fargi $op_constr)
                lhs_feasible = app_op_constraint lhs
                rhs_feasible = app_op_constraint rhs
                op_app = $app ($fargi $opE)
                op_cond = $app2 both_or_none lhs_feasible rhs_feasible
                axiom_cond = $eat_spare_cond_args ($app ar_cond $axE)
            in label test_name $ property $
                $(apply_gens_to
                    [| $app2 (==>) axiom_cond $
                       $app2 (==>) op_cond $
                       $app2 (==) (op_app lhs) (op_app rhs)
                    |]
                    all_gens)
         |]

both_or_none :: Bool -> Bool -> Bool
both_or_none a b |      a && b  = True
both_or_none a b | not (a || b) = False
both_or_none _ _ | otherwise    = error "Applicability constraint true for only one side of the axiom!"

eat_args :: Int -> Int -> Q Exp
eat_args from to = [| $(liftAn (from-1)) $(eat_n_args (to-from+1)) |]

eat_n_args :: Int -> ExpQ
eat_n_args = constN

constN :: Int -> ExpQ
constN n | n < 0 = fail "argument too small"
constN 0 = [| id |]
constN n = [| const . $(constN (n-1)) |]

liftAn :: Int -> ExpQ
liftAn n | n < 0 = fail "argument too small"
liftAn 0 = [| id |]
liftAn n = [| (.) . $(liftAn (n-1)) |]

liftA2n :: Int -> ExpQ
liftA2n n | n < 0 = fail "argument too small"
liftA2n 0 = [| id |]
liftA2n n = [| liftA2 . $(liftA2n (n-1)) |]

flip_n_to_1 :: Int -> ExpQ
flip_n_to_1 n | n < 1 = fail "argument too small"
flip_n_to_1 1 = [| id |]
flip_n_to_1 n = [| flip . (.) $(flip_n_to_1 (n-1)) |]

{- Allows all (i, opn) for which the i'th argument of opn is exactly the type of ax.
 - It should allow matching types.
 -}
filter_ops :: Name -> [(Int, Name)] -> Q [(Int, Name)]
filter_ops ax = filterM (is_applicable ax)

is_applicable :: Name -> (Int, Name) -> Q Bool
is_applicable ax (argi, opn) = do xt <- return_type_of_axiom ax
                                  ot <- type_of_argi argi opn
                                  return (xt == ot)

filter_Ops :: Axiom -> [OpArg] -> Q [OpArg]
filter_Ops ax = filterM (is_applicable_Op $ axiom_name ax)

is_applicable_Op :: Name -> OpArg -> Q Bool
is_applicable_Op ax opn = do xt <- return_type_of_axiom ax
                             ot <- type_of_argi (opArg_argi opn) (opArg_name opn)
                             mapM_ warn_if_type_is_polymorphic [xt, ot]
                             return (xt == ot)

warn_if_type_is_polymorphic :: Type -> Q ()
warn_if_type_is_polymorphic tp = if type_is_polymorphic tp
                                 then reportError $ "Polymorphic types are not supported!\n"
                                                 ++ "Offending type: " ++ pprint tp ++ "\n"
                                 else return ()

type_is_polymorphic :: Type -> Bool
type_is_polymorphic tp = case tp of
    ForallT vars _ tp' -> length vars > 0 || type_is_polymorphic tp'
    AppT tp' tp'' -> type_is_polymorphic tp' || type_is_polymorphic tp''
    SigT tp' _ -> type_is_polymorphic tp'
    VarT _ -> True
    ConT _ -> False
    PromotedT _ -> False
    TupleT _ -> False
    UnboxedTupleT _ -> False
    ArrowT -> False
    ListT -> False
    PromotedTupleT _ -> False
    PromotedNilT -> False
    PromotedConsT -> False
    StarT -> False
    ConstraintT -> False
    LitT _ -> False

{- Axioms return a pair of type (a, a), and we want the type of a.
 -}
return_type_of_axiom :: Name -> TypeQ
return_type_of_axiom name = return_type_of name >>= axiom_result_type

axiom_result_type :: Type -> TypeQ
axiom_result_type (AppT (ConT nm) tp) | nm == ''AxiomResult = return tp
axiom_result_type tp = fail $ "Expected ``AxiomResult a'' for some a, but got ``" ++ pprint tp ++ "''"

-- For a type (a, a) returns the type of a.
pair_type :: Type -> TypeQ
pair_type (AppT (AppT (TupleT 2) a) b) | a == b = return a
pair_type tp = fail $ "Expected a pair ``(a, a)'' but got ``" ++ pprint tp ++ "''"

return_type_of :: Name -> TypeQ
return_type_of = liftM return_type . typeof

type_of_argi :: Int -> Name -> TypeQ
type_of_argi i = liftM (arg_i_type i) . typeof

return_type :: Type -> Type
return_type (ForallT _ _ tp) = return_type tp
return_type (AppT (AppT ArrowT _) tp) = return_type tp
return_type tp = tp

arg_i_type :: Int -> Type -> Type
arg_i_type i (ForallT _ _ tp) = arg_i_type i tp
arg_i_type 1 (AppT (AppT ArrowT tp) _) = tp
arg_i_type i (AppT (AppT ArrowT _) tp) = arg_i_type (i-1) tp
arg_i_type i t = error $ "Failed pattern match in arg_i_type.\n"
    ++ "Please file a bug report and attach the following output.\n"
    ++ "i=" ++ show i ++ "\n"
    ++ show t