{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs #-}
module Generators where

import Test.QuickCheck
import Data.Proxy
import Data.Coerce

data Expr a where
  Const :: a -> Expr a
  Map :: Show b => Expr b -> (b -> a) -> Expr a
  Filter :: Expr a -> (a -> Bool) -> Expr a
  Join :: (Show a, Show b) => Expr a -> Expr b -> Expr (a, b)


instance Show a => Show (Expr a) where
  show (Const a) = "(Const " ++ show a ++ ")"
  show (Map e _) = "(Map f " ++ show e ++ ")"
  show (Filter e _) = "(Filter p " ++ show e ++ ")"
  show (Join l r) = "(Join " ++ show l ++ " " ++ show r ++ ")"



instance (Arbitrary a, CoArbitrary a, Show a) => Arbitrary (Expr a) where
  arbitrary = oneof (generators (Proxy @(Expr a)))

newtype TupledExpr a = Tupled {
  unTupled :: (Expr a)
}
instance (Arbitrary a, CoArbitrary a,
          Arbitrary b, CoArbitrary b,
          Show a, Show b) =>
          Arbitrary (TupledExpr (a, b)) where
  arbitrary = coerce
            $ oneof (joinGen : generators (Proxy @(TupledExpr (a,b))))


generators :: forall b a. (Coercible b (Expr a), Arbitrary b, CoArbitrary a, Arbitrary a, Show a) => Proxy b -> [Gen (Expr a)]
generators t = [
    constGen
  , filterGen t
  , mapGen t
  , mapGen (Proxy @(TupledExpr (a, Bool)))
  , mapGen (Proxy @(TupledExpr ((a, Bool), a)))
  ]

constGen :: Arbitrary a => Gen (Expr a)
constGen = Const <$> arbitrary

filterGen :: forall b a. (Coercible b (Expr a), Arbitrary b, CoArbitrary a) => Proxy b -> Gen (Expr a)
filterGen _ = do
  e <- coerce $ arbitrary @b
  p <- arbitrary @(a -> Bool)
  pure (Filter e p)

mapGen :: forall b x a. (Coercible b (Expr x), Arbitrary b, CoArbitrary x, Arbitrary a, Show x) => Proxy b -> Gen (Expr a)
mapGen _ = do
  arg <- coerce $ arbitrary @b
  f <- arbitrary @(x -> a)
  pure (Map arg f)

joinGen :: forall a b. (Arbitrary (Expr a), Arbitrary (Expr b), Show b, Show a) => Gen (Expr (a,b))
joinGen = do
  left <- arbitrary @(Expr a)
  right <- arbitrary @(Expr b)
  pure (Join left right)