-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Hedgehog.Gen.Morley.Micheline
  ( -- * 'Expression' generators
    genExpression
  , genExpressionInt
  , genExpressionString
  , genExpressionBytes
  , genExpressionSeq
  , genExpressionPrim

    -- * Generic 'Exp' generators
  , genExp
  , genExpInt
  , genExpString
  , genExpBytes
  , genExpSeq
  , genExpPrim
  , genMichelinePrimAp
  , genExprAnnotation
  ) where

import Hedgehog (MonadGen(GenBase))
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range (Range)

import Hedgehog.Gen.Michelson.Untyped (genAnnotation)
import Hedgehog.Range.Defaults
import Morley.Micheline.Expression

----------------------------------------------------------------------------
-- Expression generators
----------------------------------------------------------------------------

genExpression :: forall m. (MonadGen m, GenBase m ~ Identity) => m Expression
genExpression = genExp Nothing Nothing (mkUniformExpExtras Gen.enumBounded)

genExpressionInt :: MonadGen m => Range ExpressionInt -> m Expression
genExpressionInt = genExpInt (pure ())

genExpressionString :: MonadGen m => Range SmallLength -> m Expression
genExpressionString = genExpString (pure ())

genExpressionBytes :: MonadGen m => Range Length -> m Expression
genExpressionBytes = genExpBytes (pure ())

genExpressionSeq :: (MonadGen m, GenBase m ~ Identity) => Range SmallLength -> m Expression
genExpressionSeq = genExpSeq genExpression (pure ())

genExpressionPrim :: (MonadGen m, GenBase m ~ Identity) => m Expression
genExpressionPrim = genExpPrim genExpression (pure ())

----------------------------------------------------------------------------
-- Generic Exp generators
----------------------------------------------------------------------------

-- | Generate extended expression given the generators for all the extension
-- points.
--
-- In case your expression has no extra constructors, avoid supplying
-- @Just 'Gen.discard'@ as that would cause the generator to give up
-- periodically (supply 'Nothing' instead).
genExp
  :: forall x m. (MonadGen m, GenBase m ~ Identity)
  => Maybe (m (XExp x))
     -- ^ Non-recursive extra constructors
  -> Maybe (m (XExp x))
     -- ^ Recursive extra constructors (that can generate @Exp@ inside)
  -> ExpExtras m x
  -> m (Exp x)
genExp mGenXCon mGenXConRec xGens = Gen.recursive Gen.choice
  ( maybe id (:) (ExpX <<$>> mGenXCon)
    [ genExpInt (eeInt xGens) def
    , genExpString (eeString xGens) def
    , genExpBytes (eeBytes xGens) def
    ]
  )
  ( maybe id (:) (ExpX <<$>> mGenXConRec)
    [genExpSeq runRec (eeSeq xGens) def, genExpPrim runRec (eePrim xGens)]
  )
  where
    runRec = Gen.subterm (genExp mGenXCon mGenXConRec xGens) id

genExpInt :: MonadGen m => m (XExpInt x) -> Range ExpressionInt -> m (Exp x)
genExpInt genX range =
  ExpInt <$> genX <*> Gen.integral (unExpressionInt <$> range)

genExpString :: MonadGen m => m (XExpString x) -> Range SmallLength -> m (Exp x)
genExpString genX rangeLen =
  ExpString <$> genX <*> (Gen.text (unSmallLength <$> rangeLen) Gen.unicodeAll)

genExpBytes :: MonadGen m => m (XExpBytes x) -> Range Length -> m (Exp x)
genExpBytes genX rangeLen =
  ExpBytes <$> genX <*> Gen.bytes (unLength <$> rangeLen)

genExpSeq
  :: MonadGen m
  => m (Exp x) -> m (XExpSeq x) -> Range SmallLength -> m (Exp x)
genExpSeq doGenExp genX rangeLen =
  ExpSeq <$> genX <*> Gen.list (unSmallLength <$> rangeLen) doGenExp

genExpPrim
  :: (MonadGen m, GenBase m ~ Identity)
  => m (Exp x) -> m (XExpPrim x) -> m (Exp x)
genExpPrim doGenExp genX = ExpPrim <$> genX <*> genMichelinePrimAp doGenExp

genMichelinePrimAp :: forall x m. (MonadGen m, GenBase m ~ Identity) => m (Exp x) -> m (MichelinePrimAp x)
genMichelinePrimAp doGenExp =
  MichelinePrimAp
    <$> genMichelinePrimitive
    <*> Gen.list (unSmallLength <$> def) doGenExp
    <*> genAnnots
  where
    genMichelinePrimitive = MichelinePrimitive <$> (Gen.element $ toList michelsonPrimitive)
    genAnnots = Gen.list (unSmallLength <$> def) genExprAnnotation

genExprAnnotation :: (MonadGen m, GenBase m ~ Identity) => m Annotation
genExprAnnotation =  Gen.choice
  [genAnnotationType, genAnnotationVariable, genAnnotationField]
  where
    genAnnotationType = AnnotationType <$> genAnnotation
    genAnnotationVariable = AnnotationVariable <$> genAnnotation
    genAnnotationField = AnnotationField <$> genAnnotation