{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}

module Main where

import Data.Maybe (isNothing)
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word (Word8, Word16, Word32, Word64)
import GHC.TypeNats
import GHC.Generics
import Data.Kind (Type)
import Data.Finite (Finite)
import Data.Proxy (Proxy(..))
import Control.Monad.Loops (andM)
import Data.Typeable (Typeable, typeRep)
import Data.String (fromString)
import Hedgehog

import qualified Hedgehog.Gen as G
import qualified Hedgehog.Range as R
import qualified Data.Vector.Sized as VS
import qualified Data.Vector.Unboxed.Sized as VUS
import qualified Data.Vector.Storable.Sized as VSS
import qualified Data.Bit as B
import qualified Data.Bit.ThreadSafe as BTS

import Data.Finitary

-- A representation of types
data SomeFinitaryRep where
  SomeFinitaryRep :: forall (a :: Type) . (Finitary a, Show a, Eq a, Typeable a, 1 <= (Cardinality a)) => Proxy a -> SomeFinitaryRep

-- A representation of test functions
data SomeTestFunction where
  SomeTestFunction :: (forall (a :: Type) . (Finitary a, Show a, Eq a, 1 <= (Cardinality a)) => Proxy a -> Property) -> SomeTestFunction

-- Some weird generic
data Foo = Bar | Baz (VS.Vector 4 Bool) | Quux Word8
  deriving (Eq, Show, Typeable, Generic, Finitary)

-- Generators
choose :: forall (a :: Type) m . (MonadGen m, Finitary a) => m a
choose = fromFinite <$> chooseFinite

chooseFinite :: forall (n :: Nat) m . (KnownNat n, MonadGen m) => m (Finite n)
chooseFinite = fromIntegral <$> G.integral (R.linear 0 limit)
  where limit = subtract @Integer 1 . fromIntegral . natVal @n $ Proxy

-- Data
allTheTypes :: [SomeFinitaryRep]
allTheTypes = [
  SomeFinitaryRep @() Proxy,
  SomeFinitaryRep @(Proxy Int) Proxy,
  SomeFinitaryRep @Bool Proxy,
  SomeFinitaryRep @B.Bit Proxy,
  SomeFinitaryRep @BTS.Bit Proxy,
  SomeFinitaryRep @Ordering Proxy,
  SomeFinitaryRep @Char Proxy,
  SomeFinitaryRep @Word8 Proxy,
  SomeFinitaryRep @Word16 Proxy,
  SomeFinitaryRep @Word32 Proxy,
  SomeFinitaryRep @Word64 Proxy,
  SomeFinitaryRep @Int16 Proxy,
  SomeFinitaryRep @Int8 Proxy,
  SomeFinitaryRep @Int32 Proxy,
  SomeFinitaryRep @Int64 Proxy,
  SomeFinitaryRep @Int Proxy,
  SomeFinitaryRep @Word Proxy,
  SomeFinitaryRep @(Maybe Word8) Proxy,
  SomeFinitaryRep @(Either Word8 Int8) Proxy,
  SomeFinitaryRep @(Word8, Int8) Proxy,
  SomeFinitaryRep @(VS.Vector 4 Bool) Proxy,
  SomeFinitaryRep @(VUS.Vector 4 Bool) Proxy,
  SomeFinitaryRep @(VSS.Vector 4 Bool) Proxy,
  SomeFinitaryRep @Foo Proxy
  ]

constructTest :: SomeTestFunction -> GroupName -> IO Bool
constructTest (SomeTestFunction testFunc) name = checkParallel . Group name $ fmap go allTheTypes
  where go (SomeFinitaryRep p) = (fromString . show . typeRep $ p, testFunc p)

-- Properties
isBijection :: forall (a :: Type) . (Finitary a, Show a, Eq a) => Proxy a -> Property
isBijection _ = withTests testLimit (property $ do x <- forAll $ choose @a
                                                   x === (fromFinite . toFinite $ x))
  where testLimit = fromIntegral . (* 2) . min 32767 . natVal @(Cardinality a) $ Proxy

startIsCorrect :: forall (a :: Type) . (Finitary a, Show a, Eq a, 1 <= (Cardinality a)) => Proxy a -> Property
startIsCorrect _ = property $ start @a === fromFinite minBound

previousStartNothing :: forall (a :: Type) . (Finitary a, 1 <= (Cardinality a)) => Proxy a -> Property
previousStartNothing _ = property $ assert (isNothing . previous $ start @a)

endNextNothing :: forall (a :: Type) . (Finitary a, 1 <= (Cardinality a)) => Proxy a -> Property
endNextNothing _ = property $ assert (isNothing . next $ end @a)

endIsCorrect :: forall (a :: Type) . (Finitary a, Show a, Eq a, 1 <= (Cardinality a)) => Proxy a -> Property
endIsCorrect _ = property $ end @a === fromFinite maxBound

allTests :: [(SomeTestFunction, GroupName)]
allTests = [
  (SomeTestFunction isBijection, "bijectivity"),
  (SomeTestFunction startIsCorrect, "start"),
  (SomeTestFunction endIsCorrect, "end"),
  (SomeTestFunction previousStartNothing, "previous + start"),
  (SomeTestFunction endNextNothing, "next + end")
  ]

main :: IO Bool
main = andM . fmap (uncurry constructTest) $ allTests