{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

#if __GLASGOW_HASKELL__ >= 805
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE TypeInType #-}
#endif

import Control.Monad
import Control.Monad.ST
import Data.Primitive
import Data.Word
import Data.Proxy (Proxy(..))
import GHC.Int
import GHC.IO
import GHC.Exts
import Data.Function (on)
import Control.Applicative (Const(..))
import PrimLaws (primLaws)

import Data.Functor.Identity (Identity(..))
import qualified Data.Monoid as Monoid
import Data.Ord (Down(..))
#if MIN_VERSION_base(4,9,0)
import Data.Semigroup (stimes)
import qualified Data.Semigroup as Semigroup
#endif
#if !(MIN_VERSION_base(4,11,0))
import Data.Monoid ((<>))
#endif
#if __GLASGOW_HASKELL__ >= 805
import Foreign.Storable (Storable)
#endif
import Data.Orphans ()

import Test.Tasty (defaultMain,testGroup,TestTree)
import Test.QuickCheck (Arbitrary,Arbitrary1,Gen,CoArbitrary,Function,(===),(==>))
import qualified Test.Tasty.QuickCheck as TQC
import qualified Test.QuickCheck as QC
import qualified Test.QuickCheck.Classes.Base as QCC
import qualified Test.QuickCheck.Classes.Base.IsList as QCCL
import qualified Data.List as L

main :: IO ()
main = do
  testArray
  testByteArray
  defaultMain $ testGroup "properties"
    [ testGroup "Array"
      [ lawsToTest (QCC.eqLaws (Proxy :: Proxy (Array Int)))
      , lawsToTest (QCC.ordLaws (Proxy :: Proxy (Array Int)))
      , lawsToTest (QCC.monoidLaws (Proxy :: Proxy (Array Int)))
      , lawsToTest (QCC.showReadLaws (Proxy :: Proxy (Array Int)))
      , lawsToTest (QCC.functorLaws (Proxy1 :: Proxy1 Array))
      , lawsToTest (QCC.applicativeLaws (Proxy1 :: Proxy1 Array))
      , lawsToTest (QCC.monadLaws (Proxy1 :: Proxy1 Array))
      , lawsToTest (QCC.foldableLaws (Proxy1 :: Proxy1 Array))
      , lawsToTest (QCC.traversableLaws (Proxy1 :: Proxy1 Array))
      , lawsToTest (QCC.isListLaws (Proxy :: Proxy (Array Int)))
      , TQC.testProperty "mapArray'" (QCCL.mapProp int16 int32 mapArray')
      , TQC.testProperty "*>" $ \(xs :: Array Int) (ys :: Array Int) -> toList (xs *> ys) === (toList xs *> toList ys)
      , TQC.testProperty "<*" $ \(xs :: Array Int) (ys :: Array Int) -> toList (xs <* ys) === (toList xs <* toList ys)
      ]
    , testGroup "SmallArray"
      [ lawsToTest (QCC.eqLaws (Proxy :: Proxy (SmallArray Int)))
      , lawsToTest (QCC.ordLaws (Proxy :: Proxy (SmallArray Int)))
      , lawsToTest (QCC.monoidLaws (Proxy :: Proxy (SmallArray Int)))
      , lawsToTest (QCC.showReadLaws (Proxy :: Proxy (Array Int)))
      , lawsToTest (QCC.functorLaws (Proxy1 :: Proxy1 SmallArray))
      , lawsToTest (QCC.applicativeLaws (Proxy1 :: Proxy1 SmallArray))
      , lawsToTest (QCC.monadLaws (Proxy1 :: Proxy1 SmallArray))
      , lawsToTest (QCC.foldableLaws (Proxy1 :: Proxy1 SmallArray))
      , lawsToTest (QCC.traversableLaws (Proxy1 :: Proxy1 SmallArray))
      , lawsToTest (QCC.isListLaws (Proxy :: Proxy (SmallArray Int)))
      , TQC.testProperty "mapSmallArray'" (QCCL.mapProp int16 int32 mapSmallArray')
      , TQC.testProperty "*>" $ \(xs :: SmallArray Int) (ys :: SmallArray Int) -> toList (xs *> ys) === (toList xs *> toList ys)
      , TQC.testProperty "<*" $ \(xs :: SmallArray Int) (ys :: SmallArray Int) -> toList (xs <* ys) === (toList xs <* toList ys)
      ]
    , testGroup "ByteArray"
      [ testGroup "Ordering"
        [ TQC.testProperty "equality" byteArrayEqProp
        , TQC.testProperty "compare" byteArrayCompareProp
      , testGroup "Filling"
        [ TQC.testProperty "Int8" (setByteArrayProp (Proxy :: Proxy Int8))
        , TQC.testProperty "Int16" (setByteArrayProp (Proxy :: Proxy Int16))
        , TQC.testProperty "Int32" (setByteArrayProp (Proxy :: Proxy Int32))
        , TQC.testProperty "Int64" (setByteArrayProp (Proxy :: Proxy Int64))
        , TQC.testProperty "Int" (setByteArrayProp (Proxy :: Proxy Int))
        , TQC.testProperty "Word8" (setByteArrayProp (Proxy :: Proxy Word8))
        , TQC.testProperty "Word16" (setByteArrayProp (Proxy :: Proxy Word16))
        , TQC.testProperty "Word32" (setByteArrayProp (Proxy :: Proxy Word32))
        , TQC.testProperty "Word64" (setByteArrayProp (Proxy :: Proxy Word64))
        , TQC.testProperty "Word" (setByteArrayProp (Proxy :: Proxy Word))
        ]
      ]
      , testGroup "Resize"
        [ TQC.testProperty "shrink" byteArrayShrinkProp
        , TQC.testProperty "grow" byteArrayGrowProp
        ]
      , lawsToTest (QCC.eqLaws (Proxy :: Proxy ByteArray))
      , lawsToTest (QCC.ordLaws (Proxy :: Proxy ByteArray))
      , lawsToTest (QCC.showReadLaws (Proxy :: Proxy (Array Int)))
      , lawsToTest (QCC.isListLaws (Proxy :: Proxy ByteArray))
      , TQC.testProperty "foldrByteArray" (QCCL.foldrProp word8 foldrByteArray)
      ]
    , testGroup "PrimArray"
      [ lawsToTest (QCC.eqLaws (Proxy :: Proxy (PrimArray Word16)))
      , lawsToTest (QCC.ordLaws (Proxy :: Proxy (PrimArray Word16)))
      , lawsToTest (QCC.monoidLaws (Proxy :: Proxy (PrimArray Word16)))
      , lawsToTest (QCC.isListLaws (Proxy :: Proxy (PrimArray Word16)))
      , TQC.testProperty "foldrPrimArray" (QCCL.foldrProp int16 foldrPrimArray)
      , TQC.testProperty "foldrPrimArray'" (QCCL.foldrProp int16 foldrPrimArray')
      , TQC.testProperty "foldlPrimArray" (QCCL.foldlProp int16 foldlPrimArray)
      , TQC.testProperty "foldlPrimArray'" (QCCL.foldlProp int16 foldlPrimArray')
      , TQC.testProperty "foldlPrimArrayM'" (QCCL.foldlMProp int16 foldlPrimArrayM')
      , TQC.testProperty "mapPrimArray" (QCCL.mapProp int16 int32 mapPrimArray)
      , TQC.testProperty "traversePrimArray" (QCCL.traverseProp int16 int32 traversePrimArray)
      , TQC.testProperty "traversePrimArrayP" (QCCL.traverseProp int16 int32 traversePrimArrayP)
      , TQC.testProperty "imapPrimArray" (QCCL.imapProp int16 int32 imapPrimArray)
      , TQC.testProperty "itraversePrimArray" (QCCL.imapMProp int16 int32 itraversePrimArray)
      , TQC.testProperty "itraversePrimArrayP" (QCCL.imapMProp int16 int32 itraversePrimArrayP)
      , TQC.testProperty "generatePrimArray" (QCCL.generateProp int16 generatePrimArray)
      , TQC.testProperty "generatePrimArrayA" (QCCL.generateMProp int16 generatePrimArrayA)
      , TQC.testProperty "generatePrimArrayP" (QCCL.generateMProp int16 generatePrimArrayP)
      , TQC.testProperty "replicatePrimArray" (QCCL.replicateProp int16 replicatePrimArray)
      , TQC.testProperty "replicatePrimArrayA" (QCCL.replicateMProp int16 replicatePrimArrayA)
      , TQC.testProperty "replicatePrimArrayP" (QCCL.replicateMProp int16 replicatePrimArrayP)
      , TQC.testProperty "filterPrimArray" (QCCL.filterProp int16 filterPrimArray)
      , TQC.testProperty "filterPrimArrayA" (QCCL.filterMProp int16 filterPrimArrayA)
      , TQC.testProperty "filterPrimArrayP" (QCCL.filterMProp int16 filterPrimArrayP)
      , TQC.testProperty "mapMaybePrimArray" (QCCL.mapMaybeProp int16 int32 mapMaybePrimArray)
      , TQC.testProperty "mapMaybePrimArrayA" (QCCL.mapMaybeMProp int16 int32 mapMaybePrimArrayA)
      , TQC.testProperty "mapMaybePrimArrayP" (QCCL.mapMaybeMProp int16 int32 mapMaybePrimArrayP)
      ]
    , testGroup "DefaultSetMethod"
      [ lawsToTest (primLaws (Proxy :: Proxy DefaultSetMethod))
      ]
#if __GLASGOW_HASKELL__ >= 805
    , testGroup "PrimStorable"
      [ lawsToTest (QCC.storableLaws (Proxy :: Proxy Derived))
      ]
#endif
    , testGroup "Prim"
      [ renameLawsToTest "Word" (primLaws (Proxy :: Proxy Word))
      , renameLawsToTest "Word8" (primLaws (Proxy :: Proxy Word8))
      , renameLawsToTest "Word16" (primLaws (Proxy :: Proxy Word16))
      , renameLawsToTest "Word32" (primLaws (Proxy :: Proxy Word32))
      , renameLawsToTest "Word64" (primLaws (Proxy :: Proxy Word64))
      , renameLawsToTest "Int" (primLaws (Proxy :: Proxy Int))
      , renameLawsToTest "Int8" (primLaws (Proxy :: Proxy Int8))
      , renameLawsToTest "Int16" (primLaws (Proxy :: Proxy Int16))
      , renameLawsToTest "Int32" (primLaws (Proxy :: Proxy Int32))
      , renameLawsToTest "Int64" (primLaws (Proxy :: Proxy Int64))
      , renameLawsToTest "Const" (primLaws (Proxy :: Proxy (Const Int16 Int16)))
      , renameLawsToTest "Down" (primLaws (Proxy :: Proxy (Down Int16)))
      , renameLawsToTest "Identity" (primLaws (Proxy :: Proxy (Identity Int16)))
      , renameLawsToTest "Dual" (primLaws (Proxy :: Proxy (Monoid.Dual Int16)))
      , renameLawsToTest "Sum" (primLaws (Proxy :: Proxy (Monoid.Sum Int16)))
      , renameLawsToTest "Product" (primLaws (Proxy :: Proxy (Monoid.Product Int16)))
#if MIN_VERSION_base(4,9,0)
      , renameLawsToTest "First" (primLaws (Proxy :: Proxy (Semigroup.First Int16)))
      , renameLawsToTest "Last" (primLaws (Proxy :: Proxy (Semigroup.Last Int16)))
      , renameLawsToTest "Min" (primLaws (Proxy :: Proxy (Semigroup.Min Int16)))
      , renameLawsToTest "Max" (primLaws (Proxy :: Proxy (Semigroup.Max Int16)))
#endif
      ]
    ]

deriving instance Arbitrary a => Arbitrary (Down a)
-- Const, Dual, Sum, Product: all have Arbitrary instances defined
-- in QuickCheck itself
#if MIN_VERSION_base(4,9,0)
deriving instance Arbitrary a => Arbitrary (Semigroup.First a)
deriving instance Arbitrary a => Arbitrary (Semigroup.Last a)
deriving instance Arbitrary a => Arbitrary (Semigroup.Min a)
deriving instance Arbitrary a => Arbitrary (Semigroup.Max a)
#endif

word8 :: Proxy Word8
word8 = Proxy

int16 :: Proxy Int16
int16 = Proxy

int32 :: Proxy Int32
int32 = Proxy


setByteArrayProp :: forall a. (Prim a, Eq a, Arbitrary a, Show a) => Proxy a -> QC.Property
setByteArrayProp _ = QC.property $ \(QC.NonNegative (n :: Int)) (QC.NonNegative (off :: Int)) (QC.NonNegative (len :: Int)) (x :: a) (y :: a) ->
  (off < n && off + len <= n) ==>
  -- We use PrimArray in this test because it makes it easier to
  -- get the element-vs-byte distinction right.
  let actual = runST $ do
        m <- newPrimArray n
        forM_ (enumFromTo 0 (n - 1)) $ \ix -> writePrimArray m ix x
        setPrimArray m off len y
        unsafeFreezePrimArray m
      expected = runST $ do
        m <- newPrimArray n
        forM_ (enumFromTo 0 (n - 1)) $ \ix -> writePrimArray m ix x
        forM_ (enumFromTo off (off + len - 1)) $ \ix -> writePrimArray m ix y
        unsafeFreezePrimArray m
   in expected === actual


-- Tests that using resizeByteArray to shrink a byte array produces
-- the same results as calling Data.List.take on the list that the
-- byte array corresponds to.
byteArrayShrinkProp :: QC.Property
byteArrayShrinkProp = QC.property $ \(QC.NonNegative (n :: Int)) (QC.NonNegative (m :: Int)) ->
  let large = max n m
      small = min n m
      xs = intsLessThan large
      ys = byteArrayFromList xs
      largeBytes = large * sizeOf (undefined :: Int)
      smallBytes = small * sizeOf (undefined :: Int)
      expected = byteArrayFromList (L.take small xs)
      actual = runST $ do
        mzs0 <- newByteArray largeBytes
        copyByteArray mzs0 0 ys 0 largeBytes
        mzs1 <- resizeMutableByteArray mzs0 smallBytes
        unsafeFreezeByteArray mzs1
   in expected === actual

-- Tests that using resizeByteArray with copyByteArray (to fill in the
-- new empty space) to grow a byte array produces the same results as
-- calling Data.List.++ on the lists corresponding to the original
-- byte array and the appended byte array.
byteArrayGrowProp :: QC.Property
byteArrayGrowProp = QC.property $ \(QC.NonNegative (n :: Int)) (QC.NonNegative (m :: Int)) ->
  let large = max n m
      small = min n m
      xs1 = intsLessThan small
      xs2 = intsLessThan (large - small)
      ys1 = byteArrayFromList xs1
      ys2 = byteArrayFromList xs2
      largeBytes = large * sizeOf (undefined :: Int)
      smallBytes = small * sizeOf (undefined :: Int)
      expected = byteArrayFromList (xs1 ++ xs2)
      actual = runST $ do
        mzs0 <- newByteArray smallBytes
        copyByteArray mzs0 0 ys1 0 smallBytes
        mzs1 <- resizeMutableByteArray mzs0 largeBytes
        copyByteArray mzs1 smallBytes ys2 0 ((large - small) * sizeOf (undefined :: Int))
        unsafeFreezeByteArray mzs1
   in expected === actual

-- Tests that writing stable ptrs to a PrimArray, reading them back
-- out, and then dereferencing them gives correct results.
--stablePtrPrimProp :: QC.Property
--stablePtrPrimProp = QC.property $ \(xs :: [Integer]) -> unsafePerformIO $ do
--  ptrs <- mapM newStablePtr xs
--  let ptrs' = primArrayToList (primArrayFromList ptrs)
--  ys <- mapM deRefStablePtr ptrs'
--  mapM_ freeStablePtr ptrs'
--  return (xs === ys)

--stablePtrPrimBlockProp :: QC.Property
--stablePtrPrimBlockProp = QC.property $ \(x :: Word) (QC.NonNegative (len :: Int)) -> unsafePerformIO $ do
--  ptr <- newStablePtr x
--  let ptrs' = replicatePrimArray len ptr
--  let go ix = if ix < len
--        then do
--          n <- deRefStablePtr (indexPrimArray ptrs' ix)
--          ns <- go (ix + 1)
--          return (n : ns)
--        else return []
--  ys <- go 0
--  freeStablePtr ptr
--  return (L.replicate len x === ys)



-- Provide the non-negative integers up to the bound. For example:
--
-- >>> intsLessThan 5
-- [0,1,2,3,4]
intsLessThan :: Int -> [Int]
intsLessThan i = if i < 1
  then []
  else (i - 1) : intsLessThan (i - 1)

byteArrayCompareProp :: QC.Property
byteArrayCompareProp = QC.property $ \(xs :: [Word8]) (ys :: [Word8]) ->
  compareLengthFirst xs ys === compare (byteArrayFromList xs) (byteArrayFromList ys)

byteArrayEqProp :: QC.Property
byteArrayEqProp = QC.property $ \(xs :: [Word8]) (ys :: [Word8]) ->
  (compareLengthFirst xs ys == EQ) === (byteArrayFromList xs == byteArrayFromList ys)

compareLengthFirst :: [Word8] -> [Word8] -> Ordering
compareLengthFirst xs ys = (compare `on` length) xs ys <> compare xs ys

-- on GHC 7.4, Proxy is not polykinded, so we need this instead.
data Proxy1 (f :: * -> *) = Proxy1

lawsToTest :: QCC.Laws -> TestTree
lawsToTest (QCC.Laws name pairs) = testGroup name (map (uncurry TQC.testProperty) pairs)

renameLawsToTest :: String -> QCC.Laws -> TestTree
renameLawsToTest name (QCC.Laws _ pairs) = testGroup name (map (uncurry TQC.testProperty) pairs)

testArray :: IO ()
testArray = do
    arr <- newArray 1 'A'
    let unit =
            case writeArray arr 0 'B' of
                IO f ->
                    case f realWorld# of
                        (# _, _ #) -> ()
    c1 <- readArray arr 0
    return $! unit
    c2 <- readArray arr 0
    if c1 == 'A' && c2 == 'B'
        then return ()
        else error $ "Expected AB, got: " ++ show (c1, c2)

testByteArray :: IO ()
testByteArray = do
    let arr1 = mkByteArray ([0xde, 0xad, 0xbe, 0xef] :: [Word8])
        arr2 = mkByteArray ([0xde, 0xad, 0xbe, 0xef] :: [Word8])
        arr3 = mkByteArray ([0xde, 0xad, 0xbe, 0xee] :: [Word8])
        arr4 = mkByteArray ([0xde, 0xad, 0xbe, 0xdd] :: [Word8])
        arr5 = mkByteArray ([0xde, 0xad, 0xbe, 0xef, 0xde, 0xad, 0xbe, 0xdd] :: [Word8])
        arr6 = mkByteArray ([0xde, 0xad, 0x00, 0x01, 0xb0] :: [Word8])
    when (show arr1 /= "[0xde, 0xad, 0xbe, 0xef]") $
        fail $ "ByteArray Show incorrect: "++show arr1
    when (show arr6 /= "[0xde, 0xad, 0x00, 0x01, 0xb0]") $
        fail $ "ByteArray Show incorrect: "++ show arr6
    when (compareByteArrays arr3 1 arr4 1 3 /= GT) $
        fail $ "arr3[1,3] should be greater than arr4[1,3]"
    when (compareByteArrays arr3 0 arr4 1 3 /= GT) $
        fail $ "arr3[0,3] should be greater than arr4[1,3]"
    when (compareByteArrays arr5 1 arr2 1 3 /= EQ) $
        fail $ "arr3[1,3] should be equal to than arr4[1,3]"
    unless (arr1 > arr3) $
        fail $ "ByteArray Ord incorrect"
    unless (arr1 == arr2) $
        fail $ "ByteArray Eq incorrect"
    unless (mappend arr1 arr4 == arr5) $
        fail $ "ByteArray Monoid mappend incorrect"
    unless (mappend arr1 (mappend arr3 arr4) == mappend (mappend arr1 arr3) arr4) $
        fail $ "ByteArray Monoid mappend not associative"
    unless (mconcat [arr1,arr2,arr3,arr4,arr5] == (arr1 <> arr2 <> arr3 <> arr4 <> arr5)) $
        fail $ "ByteArray Monoid mconcat incorrect"
#if MIN_VERSION_base(4,9,0)
    unless (stimes (3 :: Int) arr4 == (arr4 <> arr4 <> arr4)) $
        fail $ "ByteArray Semigroup stimes incorrect"
#endif

mkByteArray :: Prim a => [a] -> ByteArray
mkByteArray xs = runST $ do
    marr <- newByteArray (length xs * sizeOf (head xs))
    sequence_ $ zipWith (writeByteArray marr) [0..] xs
    unsafeFreezeByteArray marr

instance Arbitrary1 Array where
  liftArbitrary elemGen = fmap fromList (QC.liftArbitrary elemGen)

instance Arbitrary a => Arbitrary (Array a) where
  arbitrary = fmap fromList QC.arbitrary

instance Arbitrary1 SmallArray where
  liftArbitrary elemGen = fmap smallArrayFromList (QC.liftArbitrary elemGen)

instance Arbitrary a => Arbitrary (SmallArray a) where
  arbitrary = fmap smallArrayFromList QC.arbitrary

instance Arbitrary ByteArray where
  arbitrary = do
    xs <- QC.arbitrary :: Gen [Word8]
    return $ runST $ do
      a <- newByteArray (L.length xs)
      iforM_ xs $ \ix x -> do
        writeByteArray a ix x
      unsafeFreezeByteArray a

instance (Arbitrary a, Prim a) => Arbitrary (PrimArray a) where
  arbitrary = do
    xs <- QC.arbitrary :: Gen [a]
    return $ runST $ do
      a <- newPrimArray (L.length xs)
      iforM_ xs $ \ix x -> do
        writePrimArray a ix x
      unsafeFreezePrimArray a



instance (Prim a, CoArbitrary a) => CoArbitrary (PrimArray a) where
  coarbitrary x = QC.coarbitrary (primArrayToList x)

instance (Prim a, Function a) => Function (PrimArray a) where
  function = QC.functionMap primArrayToList primArrayFromList

iforM_ :: Monad m => [a] -> (Int -> a -> m b) -> m ()
iforM_ xs0 f = go 0 xs0 where
  go !_ [] = return ()
  go !ix (x : xs) = f ix x >> go (ix + 1) xs

newtype DefaultSetMethod = DefaultSetMethod Int16
  deriving (Eq,Show,Arbitrary)

instance Prim DefaultSetMethod where
  sizeOf# _ = sizeOf# (undefined :: Int16)
  alignment# _ = alignment# (undefined :: Int16)
  indexByteArray# arr ix = DefaultSetMethod (indexByteArray# arr ix)
  readByteArray# arr ix s0 = case readByteArray# arr ix s0 of
    (# s1, n #) -> (# s1, DefaultSetMethod n #)
  writeByteArray# arr ix (DefaultSetMethod n) s0 = writeByteArray# arr ix n s0
  setByteArray# = defaultSetByteArray#
  indexOffAddr# addr off = DefaultSetMethod (indexOffAddr# addr off)
  readOffAddr# addr off s0 = case readOffAddr# addr off s0 of
    (# s1, n #) -> (# s1, DefaultSetMethod n #)
  writeOffAddr# addr off (DefaultSetMethod n) s0 = writeOffAddr# addr off n s0
  setOffAddr# = defaultSetOffAddr#

#if __GLASGOW_HASKELL__ >= 805
newtype Derived = Derived Int16
  deriving stock (Eq, Show)
  deriving newtype (Arbitrary, Prim)
  deriving Storable via (PrimStorable Derived)
#endif
