content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveLift #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE QuasiQuotes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# OPTIONS_GHC -Wno-orphans #-}
  
  module TestUtils.Arbitrary (
    ArbitraryObject (..),
    forAllArbitraryObjects,
  ) where
  
  import Control.Monad (forM)
  import Data.Aeson (ToJSON (..), Value (..), encode)
  import qualified Data.Aeson as Aeson
  import Data.List (nub)
  import Data.Proxy (Proxy (..))
  import Data.Text (Text)
  import qualified Data.Text as Text
  import Data.Typeable (Typeable)
  import GHC.Exts (fromList)
  import GHC.TypeLits (KnownSymbol)
  import Language.Haskell.TH (ExpQ, listE, runIO)
  import Language.Haskell.TH.Instances ()
  import Language.Haskell.TH.Quote (QuasiQuoter (quoteType))
  import Language.Haskell.TH.Syntax (Lift)
  import Test.QuickCheck
  
  import Data.Aeson.Schema (IsSchema, Object, schema)
  import Data.Aeson.Schema.Key (
    IsSchemaKey (..),
    SchemaKey,
    SchemaKey' (..),
    SchemaKeyV,
    toContext,
   )
  import Data.Aeson.Schema.Type (
    Schema' (..),
    SchemaObjectMapV,
    SchemaType,
    SchemaType' (..),
    SchemaTypeV,
    SchemaV,
    showSchemaV,
    toSchemaObjectV,
   )
  import Data.Aeson.Schema.Utils.All (All (..))
  import qualified Data.Aeson.Schema.Utils.Compat as Compat
  import Data.Aeson.Schema.Utils.NameLike (NameLike (..), fromName)
  
  data ArbitraryObject where
    ArbitraryObject ::
      (IsSchema schema) =>
      Proxy (Object schema)
      -> Value
      -> SchemaV
      -> ArbitraryObject
  
  -- Show the value and schema as something that could be copied/pasted into GHCi.
  instance Show ArbitraryObject where
    show (ArbitraryObject _ v schemaV) =
      unlines
        [ "ArbitraryObject:"
        , "  " ++ show (encode v)
        , "  [schema| " ++ showSchemaV schemaV ++ " |]"
        ]
  
  -- | A Template Haskell function to generate a splice for QuickCheck tests to generate arbitrary
  --  objects with arbitrary schemas.
  --
  --  Note that for repeated runs of the test suite, the schemas will be the same, with the actual
  --  JSON values generated randomly. You need to recompile in order to generate different schemas.
  arbitraryObject :: ExpQ
  arbitraryObject = do
    arbitrarySchemas <- runIO $ genSchemaTypes 20
  
    [|oneof $(listE $ map mkSchemaGen arbitrarySchemas)|]
    where
      mkSchemaGen schemaV =
        let schemaType = quoteType schema $ showSchemaV schemaV
         in [|genSchema' (Proxy :: Proxy (Object $schemaType)) schemaV|]
  
  -- |
  -- Splices to a 'forAll' with 'arbitraryObject', outputting information about the object
  -- generated, to ensure we get good generation.
  --
  -- >>> $(forAllArbitraryObjects) :: Testable prop => ArbitraryObject -> prop
  forAllArbitraryObjects :: ExpQ
  forAllArbitraryObjects = [|forAllArbitraryObjects' $arbitraryObject|]
  
  forAllArbitraryObjects' :: Gen ArbitraryObject -> (ArbitraryObject -> Property) -> Property
  forAllArbitraryObjects' genArbitraryObject runTest =
    forAll @_ @Property genArbitraryObject $ \o@(ArbitraryObject _ _ schemaType) ->
      tabulate "Key types" (map getKeyType $ getKeys schemaType) $
        tabulate "Schema types" (getSchemaTypes schemaType) $
          tabulate "Object sizes" (map show $ getObjectSizes schemaType) $
            tabulate "Object depth" [show $ getObjectDepth schemaType] $
              runTest o
  
  {- Run time helpers -}
  
  deriving instance Lift NameLike
  deriving instance Lift SchemaV
  deriving instance Lift SchemaTypeV
  
  genSchema' ::
    forall schema.
    ( ArbitrarySchema ('SchemaObject schema)
    , IsSchema ('Schema schema)
    ) =>
    Proxy (Object ('Schema schema))
    -> SchemaV
    -> Gen ArbitraryObject
  genSchema' proxy schemaV = do
    v <- genSchema @('SchemaObject schema)
    return $ ArbitraryObject proxy v schemaV
  
  getKeyType :: SchemaKeyV -> String
  getKeyType = \case
    NormalKey _ -> "Normal"
    PhantomKey _ -> "Phantom"
  
  getKeys :: SchemaV -> [SchemaKeyV]
  getKeys = getKeys' . toSchemaObjectV
    where
      getKeys' = \case
        SchemaMaybe inner -> getKeys' inner
        SchemaTry inner -> getKeys' inner
        SchemaList inner -> getKeys' inner
        SchemaUnion schemas -> concatMap getKeys' schemas
        SchemaObject pairs -> concatMap (\(key, inner) -> key : getKeys' inner) pairs
        _ -> []
  
  getSchemaTypes :: SchemaV -> [String]
  getSchemaTypes = getSchemaTypes' . toSchemaObjectV
    where
      getSchemaTypes' = \case
        SchemaScalar name -> [fromName name]
        SchemaMaybe inner -> "SchemaMaybe" : getSchemaTypes' inner
        SchemaTry inner -> "SchemaTry" : getSchemaTypes' inner
        SchemaList inner -> "SchemaList" : getSchemaTypes' inner
        SchemaUnion schemas -> "SchemaUnion" : concatMap getSchemaTypes' schemas
        SchemaObject pairs -> "SchemaObject" : concatMap (getSchemaTypes' . snd) pairs
        SchemaInclude _ -> error "ArbitraryObject unexpectedly generated a schema that includes another schema"
  
  getObjectSizes :: SchemaV -> [Int]
  getObjectSizes = getObjectSizes' . toSchemaObjectV
    where
      getObjectSizes' = \case
        SchemaScalar _ -> []
        SchemaMaybe inner -> getObjectSizes' inner
        SchemaTry inner -> getObjectSizes' inner
        SchemaList inner -> getObjectSizes' inner
        SchemaUnion schemas -> concatMap getObjectSizes' schemas
        SchemaObject pairs -> length pairs : concatMap (getObjectSizes' . snd) pairs
        SchemaInclude _ -> error "ArbitraryObject unexpectedly generated a schema that includes another schema"
  
  getObjectDepth :: SchemaV -> Int
  getObjectDepth = getObjectDepth' . toSchemaObjectV
    where
      getObjectDepth' = \case
        SchemaScalar _ -> 0
        SchemaMaybe inner -> getObjectDepth' inner
        SchemaTry inner -> getObjectDepth' inner
        SchemaList inner -> getObjectDepth' inner
        SchemaUnion schemas -> maximum $ map getObjectDepth' schemas
        SchemaObject pairs -> 1 + maximum (map (getObjectDepth' . snd) pairs)
        SchemaInclude _ -> error "ArbitraryObject unexpectedly generated a schema that includes another schema"
  
  {- Generating schemas -}
  
  class ArbitrarySchema (schema :: SchemaType) where
    genSchema :: Gen Value
  
  instance {-# OVERLAPS #-} ArbitrarySchema ('SchemaScalar Text) where
    genSchema = toJSON <$> arbitrary @String
  
  instance (Arbitrary inner, ToJSON inner, Typeable inner) => ArbitrarySchema ('SchemaScalar inner) where
    genSchema = toJSON <$> arbitrary @inner
  
  instance (ArbitrarySchema inner) => ArbitrarySchema ('SchemaMaybe inner) where
    genSchema =
      frequency
        [ (3, genSchema @inner)
        , (1, pure Null)
        ]
  
  instance (ArbitrarySchema inner) => ArbitrarySchema ('SchemaTry inner) where
    genSchema =
      frequency
        [ (3, genSchema @inner)
        , (1, genValue)
        ]
      where
        genValue =
          oneof
            [ pure Null
            , Number . realToFrac <$> arbitrary @Double
            , Bool <$> arbitrary
            , String . Text.pack <$> arbitrary
            ]
  
  instance (ArbitrarySchema inner) => ArbitrarySchema ('SchemaList inner) where
    genSchema = Array . fromList <$> listOf (genSchema @inner)
  
  instance (All ArbitrarySchema schemas) => ArbitrarySchema ('SchemaUnion schemas) where
    genSchema = oneof $ mapAll @ArbitrarySchema @schemas genSchemaElem
      where
        genSchemaElem :: forall schema. (ArbitrarySchema schema) => Proxy schema -> Gen Value
        genSchemaElem _ = genSchema @schema
  
  instance (All ArbitraryObjectPair pairs) => ArbitrarySchema ('SchemaObject (pairs :: [(SchemaKey, SchemaType)])) where
    genSchema = Object . Compat.unions <$> genSchemaPairs
      where
        genSchemaPairs :: Gen [Aeson.Object]
        genSchemaPairs = sequence $ mapAll @ArbitraryObjectPair @pairs genSchemaPair
  
  class (IsSchemaKey (Fst pair)) => ArbitraryObjectPair (pair :: (SchemaKey, SchemaType)) where
    genSchemaPair :: Proxy pair -> Gen Aeson.Object
    genSchemaPair _ = toContext schemaKey <$> genInnerSchema @pair
      where
        schemaKey = toSchemaKeyV $ Proxy @(Fst pair)
  
    genInnerSchema :: Gen Value
  
  instance (IsSchemaKey key, ArbitrarySchema schema) => ArbitraryObjectPair '(key, schema) where
    genInnerSchema = genSchema @schema
  
  -- For phantom keys, Maybe is only valid for Objects. Since phantom keys parse the schema with
  -- the current object as the context, we should guarantee that this only generates objects, and
  -- not Null.
  instance
    {-# OVERLAPS #-}
    (KnownSymbol key, inner ~ 'SchemaObject a, ArbitrarySchema inner) =>
    ArbitraryObjectPair '( 'PhantomKey key, 'SchemaMaybe inner)
    where
    genInnerSchema = genSchema @inner
  
  -- For phantom keys, Try can be used on any schema, but for all non-object schemas, need to ensure
  -- we generate 'Null', because Try on a non-object schema will always be an invalid parse.
  instance
    {-# OVERLAPS #-}
    (KnownSymbol key, ArbitrarySchema ('SchemaTry inner)) =>
    ArbitraryObjectPair '( 'PhantomKey key, 'SchemaTry inner)
    where
    genInnerSchema = castNull <$> genSchema @('SchemaTry inner)
      where
        castNull inner =
          case inner of
            Object _ -> inner
            _ -> Null
  
  -- For phantom keys, Union can be used on any schemas, as long as at least one is an object schema.
  instance
    {-# OVERLAPS #-}
    (KnownSymbol key, FilterObjectSchemas schemas ~ objectSchemas, ArbitrarySchema ('SchemaUnion objectSchemas)) =>
    ArbitraryObjectPair '( 'PhantomKey key, 'SchemaUnion schemas)
    where
    genInnerSchema = genSchema @('SchemaUnion objectSchemas)
  
  {- Generating schema definitions -}
  
  genSchemaTypes :: Int -> IO [SchemaV]
  genSchemaTypes numSchemasToGenerate =
    generate $
      sequence $
        take
          numSchemasToGenerate
          [resize n arbitrary | n <- [0, 2 ..]]
  
  instance Arbitrary SchemaV where
    arbitrary = Schema <$> sized genSchemaObject
  
  -- | Generate an arbitrary schema.
  --
  --  SchemaType is a recursive definition, so we want to make sure that generating a schema will
  --  terminate, and also not take too long. The ways we account for that are:
  --   * Providing an upper bound on the depth of any object schemas in the current object (n / 2)
  --   * Providing an upper bound on the number of keys in the current object (n / 3)
  --   * Providing an upper bound on the number of schemas in a union (n / 5)
  genSchemaObject :: Int -> Gen SchemaObjectMapV
  genSchemaObject n = do
    keys <- genUniqList1 (n `div` 3) genKey
    forM keys $ \key ->
      frequency
        [ (10, genSchemaObjectPairNormal key)
        , (1, genSchemaObjectPairPhantom key)
        ]
    where
      genSchemaObject' = do
        n' <- choose (0, n `div` 2)
        SchemaObject <$> genSchemaObject n'
  
      genSchemaObjectPairNormal key = do
        schemaType <-
          frequency $
            if n == 0
              then scalarSchemaTypes
              else allSchemaTypes
        return (NormalKey key, schemaType)
  
      genSchemaObjectPairPhantom key = do
        schemaType <-
          frequency
            [ (2, SchemaMaybe <$> genSchemaObject')
            , (2, SchemaTry <$> frequency nonNullableSchemaTypes)
            , (4, genSchemaObject')
            , (1, genSchemaUnion genSchemaObject')
            ]
        return (PhantomKey key, schemaType)
  
      scalarSchemaTypes =
        [ (4, pure $ SchemaScalar $ NameRef "Bool")
        , (4, pure $ SchemaScalar $ NameRef "Int")
        , (4, pure $ SchemaScalar $ NameRef "Double")
        , (4, pure $ SchemaScalar $ NameRef "Text")
        ]
  
      nonNullableSchemaTypes =
        scalarSchemaTypes
          ++ [ (2, SchemaList <$> frequency allSchemaTypes)
             , (1, genSchemaUnion $ frequency allSchemaTypes)
             , (2, genSchemaObject')
             ]
  
      allSchemaTypes =
        nonNullableSchemaTypes
          ++ [ (2, SchemaMaybe <$> frequency nonNullableSchemaTypes)
             , (2, SchemaTry <$> frequency nonNullableSchemaTypes)
             ]
  
      -- avoid generating big unions by scaling list length
      genSchemaUnion gen = SchemaUnion <$> genUniqList1 (n `div` 5) gen
  
  -- | Generate a valid JSON key
  --  See Data.Aeson.Schema.TH.Parse.jsonKey'
  genKey :: Gen String
  genKey = listOf1 $ arbitraryPrintableChar `suchThat` (`notElem` " \"\\!?[](),.@:{}#")
  
  -- | Generate a non-empty and unique list of the given generator.
  --
  --  Takes in the max size of the list.
  genUniqList1 :: (Eq a) => Int -> Gen a -> Gen [a]
  genUniqList1 n gen = do
    k <- choose (1, max 1 n)
    take k . nub <$> infiniteListOf gen
  
  {- Helper type families -}
  
  type family Fst x where
    Fst '(a, _) = a
  
  type family FilterObjectSchemas schemas where
    FilterObjectSchemas '[] = '[]
    FilterObjectSchemas ('SchemaObject inner ': xs) = 'SchemaObject inner : FilterObjectSchemas xs
    FilterObjectSchemas (_ ': xs) = FilterObjectSchemas xs
  