{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}

module Test.Generators where

import qualified Data.Aeson.KeyMap as X
import qualified Data.List as L
import qualified Data.List.NonEmpty as NE
import qualified Data.Map as M
import qualified Data.SemVer as SemVer
import qualified Data.Text as T
import Database.Bloodhound
import Generic.Random
import Test.ApproxEq
import Test.Import

instance Arbitrary NominalDiffTime where
  arbitrary = fromInteger <$> arbitrary

#if !MIN_VERSION_QuickCheck(2,8,0)
instance (Arbitrary k, Ord k, Arbitrary v) => Arbitrary (M.Map k v) where
  arbitrary = M.fromList <$> arbitrary
#endif

instance Arbitrary Text where
  arbitrary = T.pack <$> arbitrary

instance Arbitrary UTCTime where
  arbitrary =
    UTCTime
      <$> arbitrary
      <*> (fromRational . toRational <$> choose (0 :: Double, 86400))

instance Arbitrary Day where
  arbitrary =
    ModifiedJulianDay . (2000 +) <$> arbitrary
  shrink =
    (ModifiedJulianDay <$>) . shrink . toModifiedJulianDay

#if !MIN_VERSION_QuickCheck(2,9,0)
instance Arbitrary a => Arbitrary (NonEmpty a) where
  arbitrary = liftA2 (:|) arbitrary arbitrary
#endif

arbitraryScore :: Gen Score
arbitraryScore = fmap getPositive <$> arbitrary

instance (Arbitrary a, Typeable a) => Arbitrary (Hit a) where
  arbitrary =
    Hit
      <$> arbitrary
      <*> arbitrary
      <*> arbitraryScore
      <*> arbitrary
      <*> return Nothing
      <*> arbitrary
      <*> arbitrary
      <*> return Nothing

instance Arbitrary HitFields where
  arbitrary = pure (HitFields M.empty)
  shrink = const []

instance Arbitrary HitsTotalRelation where
  arbitrary = oneof [pure HTR_EQ, pure HTR_GTE]

instance Arbitrary HitsTotal where
  arbitrary = do
    tot <- getPositive <$> arbitrary
    relation_ <- arbitrary
    return $ HitsTotal tot relation_

instance (Arbitrary a, Typeable a) => Arbitrary (SearchHits a) where
  arbitrary = reduceSize $ do
    tot <- arbitrary
    score <- arbitraryScore
    hs <- arbitrary
    return $ SearchHits tot score hs

reduceSize :: Gen a -> Gen a
reduceSize f = sized $ \n -> resize (n `div` 2) f

arbitraryAlphaNum :: Gen Char
arbitraryAlphaNum =
  oneof
    [ choose ('a', 'z'),
      choose ('A', 'Z'),
      choose ('0', '9')
    ]

instance Arbitrary RoutingValue where
  arbitrary = RoutingValue . T.pack <$> listOf1 arbitraryAlphaNum

instance Arbitrary AliasRouting where
  arbitrary =
    oneof
      [ allAlias,
        one,
        theOther,
        both'
      ]
    where
      one =
        GranularAliasRouting
          <$> (Just <$> arbitrary)
          <*> pure Nothing
      theOther =
        GranularAliasRouting Nothing
          <$> (Just <$> arbitrary)
      both' =
        GranularAliasRouting
          <$> (Just <$> arbitrary)
          <*> (Just <$> arbitrary)
      allAlias = AllAliasRouting <$> arbitrary

instance Arbitrary FieldName where
  arbitrary =
    FieldName
      . T.pack
      <$> listOf1 arbitraryAlphaNum

#if MIN_VERSION_base(4,10,0)
-- Test.QuickCheck.Modifiers

qcNonEmptyToNonEmpty :: NonEmptyList a -> NonEmpty a
qcNonEmptyToNonEmpty (NonEmpty (a : xs)) = (a :| xs)
qcNonEmptyToNonEmpty (NonEmpty []) = error "NonEmpty was empty!"

instance Arbitrary a => Arbitrary (NonEmpty a) where
  arbitrary =
        qcNonEmptyToNonEmpty
    <$> arbitrary
#endif
instance Arbitrary ScriptFields where
  arbitrary =
    pure $
      ScriptFields $
        X.fromList []

  shrink = const []

instance Arbitrary ScriptParams where
  arbitrary =
    pure $
      ScriptParams $
        X.fromList
          [ ("a", Number 42),
            ("b", String "forty two")
          ]

  shrink = const []

instance Arbitrary RegexpFlags where
  arbitrary =
    oneof
      [ pure AllRegexpFlags,
        pure NoRegexpFlags,
        SomeRegexpFlags <$> genUniqueFlags
      ]
    where
      genUniqueFlags =
        NE.fromList . L.nub
          <$> listOf1 arbitrary

instance Arbitrary IndexAliasCreate where
  arbitrary =
    IndexAliasCreate
      <$> arbitrary
      <*> reduceSize arbitrary

instance Arbitrary Query where
  arbitrary =
    reduceSize $
      oneof
        [ TermQuery <$> arbitrary <*> arbitrary,
          TermsQuery <$> arbitrary <*> arbitrary,
          QueryMatchQuery <$> arbitrary,
          QueryMultiMatchQuery <$> arbitrary,
          QueryBoolQuery <$> arbitrary,
          QueryBoostingQuery <$> arbitrary,
          QueryCommonTermsQuery <$> arbitrary,
          ConstantScoreQuery <$> arbitrary <*> arbitrary,
          QueryDisMaxQuery <$> arbitrary,
          QueryFuzzyLikeThisQuery <$> arbitrary,
          QueryFuzzyLikeFieldQuery <$> arbitrary,
          QueryFuzzyQuery <$> arbitrary,
          QueryHasChildQuery <$> arbitrary,
          QueryHasParentQuery <$> arbitrary,
          IdsQuery <$> arbitrary,
          QueryIndicesQuery <$> arbitrary,
          MatchAllQuery <$> arbitrary,
          QueryMoreLikeThisQuery <$> arbitrary,
          QueryMoreLikeThisFieldQuery <$> arbitrary,
          QueryNestedQuery <$> arbitrary,
          QueryPrefixQuery <$> arbitrary,
          QueryQueryStringQuery <$> arbitrary,
          QuerySimpleQueryStringQuery <$> arbitrary,
          QueryRangeQuery <$> arbitrary,
          QueryRegexpQuery <$> arbitrary
        ]

-- TODO: Implement shrink
-- shrink = genericShrink

instance Arbitrary Filter where
  arbitrary =
    Filter <$> arbitrary
  shrink (Filter q) =
    Filter <$> shrink q

instance Arbitrary ReplicaBounds where
  arbitrary =
    oneof
      [ replicasBounded,
        replicasLowerBounded,
        pure ReplicasUnbounded
      ]
    where
      replicasBounded = do
        Positive a <- arbitrary
        Positive b <- arbitrary
        return (ReplicasBounded a b)
      replicasLowerBounded = do
        Positive a <- arbitrary
        return (ReplicasLowerBounded a)

instance Arbitrary NodeAttrName where
  arbitrary =
    NodeAttrName
      . T.pack
      <$> listOf1 arbitraryAlphaNum

instance Arbitrary NodeAttrFilter where
  arbitrary = do
    n <- arbitrary
    xs <- listOf1 (listOf1 arbitraryAlphaNum)
    let (s, ss) = unpackConsPartial xs
        ts = T.pack <$> s :| ss
    return (NodeAttrFilter n ts)
    where
      -- listOf1 means this shouldn't blow up.
      unpackConsPartial (x : xs) = (x, xs)
      unpackConsPartial _ = error "unpackConsPartial failed but shouldn't have"

instance Arbitrary VersionNumber where
  arbitrary = do
    major <- posInt
    minor <- posInt
    patch <- posInt
    return $ VersionNumber $ SemVer.version major minor patch [] []
    where
      posInt = getPositive <$> arbitrary

instance Arbitrary TemplateQueryKeyValuePairs where
  arbitrary = TemplateQueryKeyValuePairs . X.fromList <$> arbitrary
  shrink (TemplateQueryKeyValuePairs x) = map (TemplateQueryKeyValuePairs . X.fromList) . shrink $ X.toList x

instance Arbitrary IndexName where arbitrary = genericArbitraryU

instance Arbitrary DocId where arbitrary = genericArbitraryU

instance Arbitrary Version where arbitrary = genericArbitraryU

instance Arbitrary BuildHash where arbitrary = genericArbitraryU

instance Arbitrary IndexAliasRouting where arbitrary = genericArbitraryU

instance Arbitrary ShardCount where arbitrary = genericArbitraryU

instance Arbitrary ReplicaCount where arbitrary = genericArbitraryU

instance Arbitrary TemplateName where arbitrary = genericArbitraryU

instance Arbitrary IndexPattern where arbitrary = genericArbitraryU

instance Arbitrary QueryString where arbitrary = genericArbitraryU

instance Arbitrary CacheName where arbitrary = genericArbitraryU

instance Arbitrary CacheKey where arbitrary = genericArbitraryU

instance Arbitrary Existence where arbitrary = genericArbitraryU

instance Arbitrary CutoffFrequency where arbitrary = genericArbitraryU

instance Arbitrary Analyzer where arbitrary = genericArbitraryU

instance Arbitrary MaxExpansions where arbitrary = genericArbitraryU

instance Arbitrary Lenient where arbitrary = genericArbitraryU

instance Arbitrary Tiebreaker where arbitrary = genericArbitraryU

instance Arbitrary Boost where arbitrary = genericArbitraryU

instance Arbitrary BoostTerms where arbitrary = genericArbitraryU

instance Arbitrary MinimumMatch where arbitrary = genericArbitraryU

instance Arbitrary DisableCoord where arbitrary = genericArbitraryU

instance Arbitrary IgnoreTermFrequency where arbitrary = genericArbitraryU

instance Arbitrary MinimumTermFrequency where arbitrary = genericArbitraryU

instance Arbitrary MaxQueryTerms where arbitrary = genericArbitraryU

instance Arbitrary Fuzziness where arbitrary = genericArbitraryU

instance Arbitrary PrefixLength where arbitrary = genericArbitraryU

instance Arbitrary RelationName where arbitrary = genericArbitraryU

instance Arbitrary PercentMatch where arbitrary = genericArbitraryU

instance Arbitrary StopWord where arbitrary = genericArbitraryU

instance Arbitrary QueryPath where arbitrary = genericArbitraryU

instance Arbitrary AllowLeadingWildcard where arbitrary = genericArbitraryU

instance Arbitrary LowercaseExpanded where arbitrary = genericArbitraryU

instance Arbitrary EnablePositionIncrements where arbitrary = genericArbitraryU

instance Arbitrary AnalyzeWildcard where arbitrary = genericArbitraryU

instance Arbitrary GeneratePhraseQueries where arbitrary = genericArbitraryU

instance Arbitrary Locale where arbitrary = genericArbitraryU

instance Arbitrary MaxWordLength where arbitrary = genericArbitraryU

instance Arbitrary MinWordLength where arbitrary = genericArbitraryU

instance Arbitrary PhraseSlop where arbitrary = genericArbitraryU

instance Arbitrary MinDocFrequency where arbitrary = genericArbitraryU

instance Arbitrary MaxDocFrequency where arbitrary = genericArbitraryU

instance Arbitrary Regexp where arbitrary = genericArbitraryU

instance Arbitrary SimpleQueryStringQuery where arbitrary = genericArbitraryU

instance Arbitrary FieldOrFields where arbitrary = genericArbitraryU

instance Arbitrary SimpleQueryFlag where arbitrary = genericArbitraryU

instance Arbitrary RegexpQuery where arbitrary = genericArbitraryU

instance Arbitrary QueryStringQuery where arbitrary = genericArbitraryU

instance Arbitrary RangeQuery where arbitrary = genericArbitraryU

instance Arbitrary RangeValue where arbitrary = genericArbitraryU

instance Arbitrary PrefixQuery where arbitrary = genericArbitraryU

instance Arbitrary NestedQuery where arbitrary = genericArbitraryU

instance Arbitrary MoreLikeThisFieldQuery where arbitrary = genericArbitraryU

instance Arbitrary MoreLikeThisQuery where arbitrary = genericArbitraryU

instance Arbitrary IndicesQuery where arbitrary = genericArbitraryU

instance Arbitrary IgnoreUnmapped where arbitrary = genericArbitraryU

instance Arbitrary MinChildren where arbitrary = genericArbitraryU

instance Arbitrary MaxChildren where arbitrary = genericArbitraryU

instance Arbitrary AggregateParentScore where arbitrary = genericArbitraryU

instance Arbitrary HasParentQuery where arbitrary = genericArbitraryU

instance Arbitrary HasChildQuery where arbitrary = genericArbitraryU

instance Arbitrary FuzzyQuery where arbitrary = genericArbitraryU

instance Arbitrary FuzzyLikeFieldQuery where arbitrary = genericArbitraryU

instance Arbitrary FuzzyLikeThisQuery where arbitrary = genericArbitraryU

instance Arbitrary DisMaxQuery where arbitrary = genericArbitraryU

instance Arbitrary CommonTermsQuery where arbitrary = genericArbitraryU

instance Arbitrary DistanceRange where arbitrary = genericArbitraryU

instance Arbitrary MultiMatchQuery where arbitrary = genericArbitraryU

instance Arbitrary LessThanD where arbitrary = genericArbitraryU

instance Arbitrary LessThanEqD where arbitrary = genericArbitraryU

instance Arbitrary GreaterThanD where arbitrary = genericArbitraryU

instance Arbitrary GreaterThanEqD where arbitrary = genericArbitraryU

instance Arbitrary LessThan where arbitrary = genericArbitraryU

instance Arbitrary LessThanEq where arbitrary = genericArbitraryU

instance Arbitrary GreaterThan where arbitrary = genericArbitraryU

instance Arbitrary GreaterThanEq where arbitrary = genericArbitraryU

instance Arbitrary GeoPoint where arbitrary = genericArbitraryU

instance Arbitrary NullValue where arbitrary = genericArbitraryU

instance Arbitrary MinimumMatchHighLow where arbitrary = genericArbitraryU

instance Arbitrary CommonMinimumMatch where arbitrary = genericArbitraryU

instance Arbitrary BoostingQuery where arbitrary = genericArbitraryU

instance Arbitrary BoolQuery where arbitrary = genericArbitraryU

instance Arbitrary MatchQuery where arbitrary = genericArbitraryU

instance Arbitrary MultiMatchQueryType where arbitrary = genericArbitraryU

instance Arbitrary BooleanOperator where arbitrary = genericArbitraryU

instance Arbitrary ZeroTermsQuery where arbitrary = genericArbitraryU

instance Arbitrary MatchQueryType where arbitrary = genericArbitraryU

instance Arbitrary SearchAliasRouting where arbitrary = genericArbitraryU

instance Arbitrary ScoreType where arbitrary = genericArbitraryU

instance Arbitrary Distance where arbitrary = genericArbitraryU

instance Arbitrary DistanceUnit where arbitrary = genericArbitraryU

instance Arbitrary DistanceType where arbitrary = genericArbitraryU

instance Arbitrary OptimizeBbox where arbitrary = genericArbitraryU

instance Arbitrary GeoBoundingBoxConstraint where arbitrary = genericArbitraryU

instance Arbitrary GeoFilterType where arbitrary = genericArbitraryU

instance Arbitrary GeoBoundingBox where arbitrary = genericArbitraryU

instance Arbitrary LatLon where arbitrary = genericArbitraryU

instance Arbitrary RangeExecution where arbitrary = genericArbitraryU

instance Arbitrary RegexpFlag where arbitrary = genericArbitraryU

instance Arbitrary BoolMatch where arbitrary = genericArbitraryU

instance Arbitrary Term where arbitrary = genericArbitraryU

instance Arbitrary IndexMappingsLimits where arbitrary = genericArbitraryU

instance Arbitrary IndexSettings where arbitrary = genericArbitraryU

instance Arbitrary TokenChar where arbitrary = genericArbitraryU

instance Arbitrary Ngram where arbitrary = genericArbitraryU

instance Arbitrary TokenizerDefinition where arbitrary = genericArbitraryU

instance Arbitrary TokenFilter where arbitrary = genericArbitraryU

instance Arbitrary NgramFilter where arbitrary = genericArbitraryU

instance Arbitrary EdgeNgramFilterSide where arbitrary = genericArbitraryU

instance Arbitrary TokenFilterDefinition where arbitrary = genericArbitraryU

instance Arbitrary Language where arbitrary = genericArbitraryU

instance Arbitrary Shingle where arbitrary = genericArbitraryU

instance Arbitrary CharFilter where arbitrary = genericArbitraryU

instance Arbitrary AnalyzerDefinition where arbitrary = genericArbitraryU

-- TODO: This should have a proper generator that doesn't
-- create garbage that has to be filtered out.
instance Arbitrary CharFilterDefinition where
  arbitrary =
    oneof
      [ CharFilterDefinitionMapping
          . chomp
          <$> arbitrary,
        CharFilterDefinitionPatternReplace
          <$> arbitrary
          <*> arbitrary
          <*> arbitrary
      ]
    where
      chomp =
        M.map T.strip
          . M.mapKeys (T.replace "=>" "" . T.strip)

instance Arbitrary Analysis where arbitrary = genericArbitraryU

instance Arbitrary Tokenizer where arbitrary = genericArbitraryU

instance Arbitrary Compression where arbitrary = genericArbitraryU

instance Arbitrary Bytes where arbitrary = genericArbitraryU

instance Arbitrary AllocationPolicy where arbitrary = genericArbitraryU

instance Arbitrary InitialShardCount where arbitrary = genericArbitraryU

instance Arbitrary FSType where arbitrary = genericArbitraryU

instance Arbitrary CompoundFormat where arbitrary = genericArbitraryU

instance Arbitrary FsSnapshotRepo where arbitrary = genericArbitraryU

instance Arbitrary SnapshotRepoName where arbitrary = genericArbitraryU

instance Arbitrary DirectGeneratorSuggestModeTypes where arbitrary = genericArbitraryU

instance Arbitrary DirectGenerators where arbitrary = genericArbitraryU

instance Arbitrary PhraseSuggesterCollate where arbitrary = genericArbitraryU

instance Arbitrary PhraseSuggesterHighlighter where arbitrary = genericArbitraryU

instance Arbitrary Size where arbitrary = genericArbitraryU

instance Arbitrary PhraseSuggester where arbitrary = genericArbitraryU

instance Arbitrary SuggestType where arbitrary = genericArbitraryU

instance Arbitrary Suggest where arbitrary = genericArbitraryU

instance Arbitrary FunctionScoreQuery where arbitrary = genericArbitraryU

instance Arbitrary FunctionScoreFunction where arbitrary = genericArbitraryU

instance Arbitrary FunctionScoreFunctions where arbitrary = genericArbitraryU

instance Arbitrary ComponentFunctionScoreFunction where arbitrary = genericArbitraryU

instance Arbitrary Script where arbitrary = genericArbitraryU

instance Arbitrary ScriptLanguage where arbitrary = genericArbitraryU

instance Arbitrary ScriptSource where arbitrary = genericArbitraryU

instance Arbitrary ScoreMode where arbitrary = genericArbitraryU

instance Arbitrary BoostMode where arbitrary = genericArbitraryU

instance Arbitrary Seed where arbitrary = genericArbitraryU

instance Arbitrary FieldValueFactor where arbitrary = genericArbitraryU

instance Arbitrary Weight where arbitrary = genericArbitraryU

instance Arbitrary Factor where arbitrary = genericArbitraryU

instance Arbitrary FactorMissingFieldValue where arbitrary = genericArbitraryU

instance Arbitrary FactorModifier where arbitrary = genericArbitraryU

instance Arbitrary UpdatableIndexSetting where
  arbitrary = resize 10 genericArbitraryU

newtype UpdatableIndexSetting'
  = UpdatableIndexSetting' UpdatableIndexSetting
  deriving (Show, Eq, ToJSON, FromJSON, ApproxEq, Typeable)

instance Arbitrary UpdatableIndexSetting' where
  arbitrary = do
    settings <- arbitrary
    return $
      UpdatableIndexSetting' $ case settings of
        RoutingAllocationInclude xs ->
          RoutingAllocationInclude (dropDuplicateAttrNames xs)
        RoutingAllocationExclude xs ->
          RoutingAllocationExclude (dropDuplicateAttrNames xs)
        RoutingAllocationRequire xs ->
          RoutingAllocationRequire (dropDuplicateAttrNames xs)
        x -> x
    where
      dropDuplicateAttrNames =
        NE.fromList . L.nubBy sameAttrName . NE.toList
      sameAttrName a b =
        nodeAttrFilterName a == nodeAttrFilterName b
  shrink (UpdatableIndexSetting' x) = map UpdatableIndexSetting' (shrink x)

instance Arbitrary InnerHits where arbitrary = genericArbitraryU