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

module Hedgehog.Gen.Michelson.Untyped
  ( genInternalByteString
  , genVar
  , genTyVar
  , genExpandedOp
  , genExtInstrAbstract
  , genPrintComment
  , genStackRef
  , genTestAssert
  , genStackTypePattern
  , genInstrAbstract
  , genContract
  , genContract'
  , genEntriesOrder
  , genValue
  , genValue'
  , genElt
  , genParameterType
  , genType
  , genEpName
  , genAnnotation
  , genT
  ) where

import Prelude hiding (EQ, GT, LT)

import Data.Text qualified as T
import Hedgehog (MonadGen(GenBase), Range)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range

import Morley.Michelson.Untyped

import Hedgehog.Gen.Michelson (genErrorSrcPos, genMText)
import Hedgehog.Range.Defaults

genInternalByteString :: MonadGen m => Range Length -> m InternalByteString
genInternalByteString rangeLen = InternalByteString <$> Gen.bytes (unLength <$> rangeLen)

genVar :: MonadGen m => m Var
genVar = Var <$> genSmallText

genTyVar :: (MonadGen m, GenBase m ~ Identity) => m TyVar
genTyVar = Gen.choice [VarID <$> genVar, TyCon <$> genType]

genExpandedOp :: (MonadGen m, GenBase m ~ Identity) => m ExpandedOp
genExpandedOp = Gen.recursive Gen.choice
  -- non-recursive constructors
  [
    -- NB: When used together, genExpandedOp & genInstrAbstract are mutually recursive.
    -- So we use a generator of InstrAbstract that does not contain any ops as an escape hatch.
    PrimEx <$> Gen.choice instrAbstractWithoutOp
  ]
  -- recursive constructors
  [ PrimEx <$> genInstrAbstract genExpandedOp
  , SeqEx <$> genSmallList genExpandedOp
  , Gen.subtermM genExpandedOp $ \expandedOp -> WithSrcEx <$> genErrorSrcPos <*> pure expandedOp
  ]

genExtInstrAbstract :: (MonadGen m, GenBase m ~ Identity) => m op -> m (ExtInstrAbstract op)
genExtInstrAbstract genOp = Gen.choice
  [ STACKTYPE <$> genStackTypePattern
  , UTEST_ASSERT <$> genTestAssert genOp
  , UPRINT <$> genPrintComment def
  , UCOMMENT <$> genSmallText
  ]

genPrintComment :: MonadGen m => Range TinyLength -> m PrintComment
genPrintComment rangeLen = PrintComment <$> Gen.list (unTinyLength <$> rangeLen)
  (Gen.either genSmallText (genStackRef def))

genStackRef :: MonadGen m => Range StackRef -> m StackRef
genStackRef range = StackRef <$> Gen.integral (unStackRef <$> range)
  where unStackRef (StackRef x) = x

genTestAssert :: MonadGen m => m op -> m (TestAssert op)
genTestAssert genOp = TestAssert <$> genSmallText <*> genPrintComment def <*> genSmallList genOp

genStackTypePattern :: (MonadGen m, GenBase m ~ Identity) => m StackTypePattern
genStackTypePattern = Gen.recursive Gen.choice
  [ pure StkEmpty, pure StkRest ]
  [ Gen.subtermM genStackTypePattern $ \stp -> StkCons <$> genTyVar <*> pure stp ]

genInstrAbstract :: (MonadGen m, GenBase m ~ Identity) => m op -> m (InstrAbstract op)
genInstrAbstract genOp =
  Gen.choice $ instrAbstractWithOp genOp <> instrAbstractWithoutOp

instrAbstractWithOp :: (MonadGen m, GenBase m ~ Identity) => m op -> [m (InstrAbstract op)]
instrAbstractWithOp genOp =
  [ EXT <$> genExtInstrAbstract genOp
  , PUSH <$> genAnnotation <*> genType <*> genValue' genOp
  , IF_NONE <$> genSmallList genOp <*> genSmallList genOp
  , IF_LEFT <$> genSmallList genOp <*> genSmallList genOp
  , IF_CONS <$> genSmallList genOp <*> genSmallList genOp
  , MAP <$> genAnnotation <*> genSmallList genOp
  , ITER <$> genSmallList genOp
  , IF <$> genSmallList genOp <*> genSmallList genOp
  , LOOP <$> genSmallList genOp
  , LOOP_LEFT <$> genSmallList genOp
  , LAMBDA <$> genAnnotation <*> genType <*> genType <*> genSmallList genOp
  , DIP <$> genSmallList genOp
  , DIPN <$> Gen.word Range.linearBounded <*> genSmallList genOp
  , CREATE_CONTRACT <$> genAnnotation <*> genAnnotation <*> genContract' genOp
  ]

instrAbstractWithoutOp :: (MonadGen m, GenBase m ~ Identity) => [m (InstrAbstract op)]
instrAbstractWithoutOp =
  [ DROPN <$> Gen.word Range.linearBounded
  , pure DROP
  , DUP <$> genAnnotation
  , DUPN <$> genAnnotation <*> Gen.word Range.linearBounded
  , pure SWAP
  , DIG <$> Gen.word Range.linearBounded
  , DUG <$> Gen.word Range.linearBounded
  , SOME <$> genAnnotation <*> genAnnotation
  , NONE <$> genAnnotation <*> genAnnotation <*> genType
  , UNIT <$> genAnnotation <*> genAnnotation
  , PAIR <$> genAnnotation <*> genAnnotation <*> genAnnotation <*> genAnnotation
  , PAIRN <$> genAnnotation <*> Gen.word Range.linearBounded
  , UNPAIRN <$> Gen.word Range.linearBounded
  , CAR <$> genAnnotation <*> genAnnotation
  , CDR <$> genAnnotation <*> genAnnotation
  , LEFT <$> genAnnotation <*> genAnnotation <*> genAnnotation <*> genAnnotation <*> genType
  , RIGHT <$> genAnnotation <*> genAnnotation <*> genAnnotation <*> genAnnotation <*> genType
  , NIL <$> genAnnotation <*> genAnnotation <*> genType
  , CONS <$> genAnnotation
  , SIZE <$> genAnnotation
  , EMPTY_SET <$> genAnnotation <*> genAnnotation <*> genType
  , EMPTY_MAP <$> genAnnotation <*> genAnnotation <*> genType <*> genType
  , EMPTY_BIG_MAP <$> genAnnotation <*> genAnnotation <*> genType <*> genType
  , MEM <$> genAnnotation
  , GET <$> genAnnotation
  , GETN <$> genAnnotation <*> Gen.word Range.linearBounded
  , UPDATE <$> genAnnotation
  , EXEC <$> genAnnotation
  , APPLY <$> genAnnotation
  , pure FAILWITH
  , CAST <$> genAnnotation <*> genType
  , RENAME <$> genAnnotation
  , PACK <$> genAnnotation
  , UNPACK <$> genAnnotation <*> genAnnotation <*> genType
  , CONCAT <$> genAnnotation
  , SLICE <$> genAnnotation
  , ISNAT <$> genAnnotation
  , ADD <$> genAnnotation
  , SUB <$> genAnnotation
  , SUB_MUTEZ <$> genAnnotation
  , MUL <$> genAnnotation
  , EDIV <$> genAnnotation
  , ABS <$> genAnnotation
  , NEG <$> genAnnotation
  , LSL <$> genAnnotation
  , LSR <$> genAnnotation
  , OR <$> genAnnotation
  , AND <$> genAnnotation
  , XOR <$> genAnnotation
  , NOT <$> genAnnotation
  , COMPARE <$> genAnnotation
  , EQ <$> genAnnotation
  , NEQ <$> genAnnotation
  , LT <$> genAnnotation
  , GT <$> genAnnotation
  , LE <$> genAnnotation
  , GE <$> genAnnotation
  , INT <$> genAnnotation
  , SELF <$> genAnnotation <*> genAnnotation
  , CONTRACT <$> genAnnotation <*> genAnnotation <*> genType
  , TRANSFER_TOKENS <$> genAnnotation
  , SET_DELEGATE <$> genAnnotation
  , IMPLICIT_ACCOUNT <$> genAnnotation
  , NOW <$> genAnnotation
  , AMOUNT <$> genAnnotation
  , BALANCE <$> genAnnotation
  , CHECK_SIGNATURE <$> genAnnotation
  , SHA256 <$> genAnnotation
  , SHA512 <$> genAnnotation
  , BLAKE2B <$> genAnnotation
  , SHA3 <$> genAnnotation
  , KECCAK <$> genAnnotation
  , HASH_KEY <$> genAnnotation
  , SOURCE <$> genAnnotation
  , SENDER <$> genAnnotation
  , ADDRESS <$> genAnnotation
  , CHAIN_ID <$> genAnnotation
  , LEVEL <$> genAnnotation
  , SELF_ADDRESS <$> genAnnotation
  ]

genContract :: (MonadGen m, GenBase m ~ Identity) => m Contract
genContract = genContract' genExpandedOp

genContract' :: (MonadGen m, GenBase m ~ Identity) => m op -> m (Contract' op)
genContract' genOp = Contract
  <$> genParameterType
  <*> genType
  <*> genSmallList genOp
  <*> genEntriesOrder
  <*> Gen.list (Range.exponential 0 2) (genSomeView genOp)

genEntriesOrder :: (MonadGen m) => m EntriesOrder
genEntriesOrder = Gen.enumBounded

genViewName :: (MonadGen m, GenBase m ~ Identity) => m ViewName
genViewName = unsafe . mkViewName <$> Gen.text
  (Range.linear 0 viewNameMaxLength)
  (Gen.filter isValidViewNameChar Gen.ascii)

genSomeView :: (MonadGen m, GenBase m ~ Identity) => m op -> m (View' op)
genSomeView genOp = View
  <$> genViewName
  <*> genType
  <*> genType
  <*> genSmallList genOp

genValue :: (MonadGen m, GenBase m ~ Identity) => m Value
genValue = genValue' genExpandedOp

genValueInt :: MonadGen m => Range ValueInt -> m (Value' op)
genValueInt range = ValueInt <$> Gen.integral (unValueInt <$> range)

genValueString :: MonadGen m => Range Length -> m (Value' op)
genValueString = fmap ValueString . genMText

genValueBytes :: MonadGen m => Range Length -> m (Value' op)
genValueBytes = fmap ValueBytes . genInternalByteString

genValue' :: MonadGen m => m op -> m (Value' op)
genValue' genOp = Gen.recursive Gen.choice
  -- non-recursive constructors
  [ genValueInt def
  , genValueString def
  , genValueBytes def
  , pure ValueUnit
  , pure ValueTrue
  , pure ValueFalse
  , pure ValueNone
  , pure ValueNil
  ]
  -- recursive constructors
  [ Gen.subterm2 (genValue' genOp) (genValue' genOp) ValuePair
  , Gen.subterm (genValue' genOp) ValueLeft
  , Gen.subterm (genValue' genOp) ValueRight
  , Gen.subterm (genValue' genOp) ValueSome
  , ValueSeq <$> genSmallNonEmpty (genValue' genOp)
  , ValueMap <$> genSmallNonEmpty (genElt genOp)
  , ValueLambda <$> genSmallNonEmpty genOp
  ]

genElt :: MonadGen m => m op -> m (Elt op)
genElt genOp = Elt <$> genValue' genOp <*> genValue' genOp

genParameterType :: (MonadGen m, GenBase m ~ Identity) => m ParameterType
genParameterType = ParameterType <$> genType <*> genAnnotation

-- | Generate a 'Ty'.
genType :: (MonadGen m, GenBase m ~ Identity) => m Ty
genType = Ty <$> genT <*> genAnnotation

genComparableType :: (MonadGen m, GenBase m ~ Identity) => m Ty
genComparableType = Ty <$> genComparableT <*> genAnnotation

genEpName :: (MonadGen m, GenBase m ~ Identity) => m EpName
genEpName = Gen.mapMaybe (rightToMaybe . epNameFromRefAnn) genAnnotation

genAnnotation :: forall m a. (MonadGen m, GenBase m ~ Identity) => m (Annotation a)
genAnnotation = Gen.mapMaybe (rightToMaybe . mkAnnotation) genAnnotationText
  where
    genAnnotationText :: m Text
    genAnnotationText = Gen.frequency
      [ (1, pure "")
      , (17, T.cons <$> annStart <*> Gen.text (Range.linear 0 100) annBodyChar)
      , (1, pure specialFieldAnn)
      , (1, Gen.element specialVarAnns)
      ]

    annStart :: m Char
    annStart = Gen.choice [ pure '_', Gen.alphaNum ]

    annBodyChar :: m Char
    annBodyChar = Gen.choice [ annStart, Gen.element (".%@" :: String) ]

-- | Generate a 'T'.
genT :: (MonadGen m, GenBase m ~ Identity) => m T
genT =
  Gen.recursive Gen.choice
    -- non-recursive constructors
    [ pure TKey
    , pure TUnit
    , pure TSignature
    , pure TChainId
    , pure TOperation
    , pure TInt
    , pure TNat
    , pure TString
    , pure TBytes
    , pure TMutez
    , pure TBool
    , pure TKeyHash
    , pure TTimestamp
    , pure TAddress
    , pure TChest
    , pure TChestKey
    ]
    -- recursive constructors
    [ TOption <$> genType
    , TList <$> genType
    , TSet <$> genComparableType
    , TContract <$> genType
    , TPair <$> genAnnotation <*> genAnnotation <*> genAnnotation
            <*> genAnnotation <*> genType <*> genType
    , TOr <$> genAnnotation <*> genAnnotation <*> genType <*> genType
    , TLambda <$> genType <*> genType
    , TMap <$> genComparableType <*> genType
    , TBigMap <$> genComparableType <*> genType
    ]

-- We don't need to pass @mustBeValid@ here because comparable types are always valid.
genComparableT :: (MonadGen m, GenBase m ~ Identity) => m T
genComparableT =
  Gen.recursive Gen.choice
    [ pure TInt
    , pure TNat
    , pure TString
    , pure TBytes
    , pure TMutez
    , pure TBool
    , pure TKeyHash
    , pure TTimestamp
    , pure TAddress
    ]
    [ TPair <$> genAnnotation <*> genAnnotation <*> genAnnotation
            <*> genAnnotation <*> genComparableType <*> genComparableType
    ]

smallCollectionRange :: Range Int
smallCollectionRange = Range.linear 0 3

genSmallList :: MonadGen m => m a -> m [a]
genSmallList = Gen.list smallCollectionRange

genSmallNonEmpty :: MonadGen m => m a -> m (NonEmpty a)
genSmallNonEmpty = Gen.nonEmpty smallCollectionRange

genSmallText :: MonadGen m => m Text
genSmallText = Gen.text (Range.linear 0 10) Gen.unicodeAll