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

{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Hedgehog.Gen.Michelson.Typed
  ( genBigMap
  , genEpAddress
  , genValueKeyHash
  , genValueMutez
  , genValueInt
  , genValueList
  , genValueUnit
  , genValuePair
  , genValueTimestamp
  , genValueChestAndKey

  , genValue
  , genValue'

  , genSimpleInstr
  ) where

import Control.Exception qualified as Ex
import Data.Singletons (Sing)
import Hedgehog (GenBase, MonadGen)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range

import Hedgehog.Gen.Michelson.Untyped (genEpName)
import Hedgehog.Gen.Tezos.Address (genAddress)
import Hedgehog.Gen.Tezos.Core (genChainId, genMutez, genTimestamp)
import Hedgehog.Gen.Tezos.Crypto (genKeyHash, genPublicKey, genSignature)
import Hedgehog.Gen.Tezos.Crypto.BLS12381 (genBls12381Fr, genBls12381G1, genBls12381G2)
import Morley.Michelson.Text (mkMText)
import Morley.Michelson.Typed
  (Instr(DROP, FAILWITH, PUSH, SWAP, Seq, UNIT), SingT(..), T(..), Value'(..), toVal)
import Morley.Michelson.Typed.Entrypoints (EpAddress(..), unsafeSepcCallRoot)
import Morley.Michelson.Typed.Haskell.Value (BigMap(..), BigMapId(..), ToT, WellTypedToT)
import Morley.Michelson.Typed.Scope
import Morley.Michelson.Typed.Value (RemFail(..))
import Morley.Tezos.Crypto.BLS12381 qualified as BLS

import Hedgehog.Gen.Tezos.Crypto.Timelock (genChestAndKey)

genBigMap
  :: forall k v m.
     (MonadGen m, Ord k, WellTypedToT k, WellTypedToT v, Comparable (ToT k))
  => m k -> m v -> m (BigMap k v)
genBigMap genK genV =
  BigMap
    <$> Gen.maybe (fmap BigMapId genBigMapId)
    <*> Gen.map (Range.linear 0 100) (liftA2 (,) genK genV)

genEpAddress :: (MonadGen m, GenBase m ~ Identity) => m EpAddress
genEpAddress = EpAddress <$> genAddress <*> genEpName

genValueKeyHash :: MonadGen m => m (Value' instr 'TKeyHash)
genValueKeyHash = VKeyHash <$> genKeyHash

genValueMutez :: MonadGen m => m (Value' instr 'TMutez)
genValueMutez = VMutez <$> genMutez

genValueInt :: MonadGen m => m (Value' instr 'TInt)
genValueInt = VInt <$> Gen.integral (Range.linearFrom 0 -1000 1000)

genValueList :: (MonadGen m, SingI a) => m (Value' instr a) -> m (Value' instr ('TList a))
genValueList genA = VList <$> Gen.list (Range.linear 0 100) genA

genValueUnit :: Applicative m => m (Value' instr 'TUnit)
genValueUnit = pure VUnit

genValuePair :: MonadGen m => m (Value' instr a) -> m (Value' instr b) -> m (Value' instr ('TPair a b))
genValuePair genA genB = VPair ... (,) <$> genA <*> genB

genValueTimestamp :: MonadGen m => m (Value' instr 'TTimestamp)
genValueTimestamp = VTimestamp <$> genTimestamp

genValueChestAndKey :: MonadGen m => m (Value' instr 'TChest, Value' instr 'TChestKey)
genValueChestAndKey = bimap VChest VChestKey <$> genChestAndKey

genValue
  :: forall t m.
      (MonadGen m, GenBase m ~ Identity, HasNoOp t, WellTyped t)
  => m (Value' Instr t)
genValue = genValue' (sing @t)

-- | Generate a simple instruction.
-- Ideally instruction generator should produce instructions containing
-- all possible primitive instructions.
-- In our case we consider only a few primitive instructions and
-- pick one from a hardcoded list. Hence we call it "simple".
-- Another limitation is that input stack and output stack types must be
-- identical and non-empty.
genSimpleInstr :: (MonadGen m, inp ~ (x ': xs), SingI x) => m (Instr inp inp)
genSimpleInstr = Gen.element
  [ UNIT `Seq` FAILWITH
  , PUSH (VInt 5) `Seq` DROP
  , UNIT `Seq` SWAP `Seq` SWAP `Seq` DROP
  ]

genBigMapId :: MonadGen m => m Natural
genBigMapId = Gen.integral (Range.constant 0 100000000)

genValue'
  :: (MonadGen m, GenBase m ~ Identity, HasNoOp t, WellTyped t)
  => Sing t -> m (Value' Instr t)
genValue' = \case
  STKey -> VKey <$> genPublicKey
  STUnit -> genValueUnit
  STSignature -> VSignature <$> genSignature
  STChainId -> VChainId <$> genChainId
  STOption st -> Gen.choice
    [ pure $ VOption Nothing
    , VOption . Just <$> genValue' st
    ]
  STList (STPair STBls12381G1 STBls12381G2) -> genBls12Pairing
  STList st -> VList <$> Gen.list (Range.linear 0 100) (genValue' st)
  STSet st -> VSet <$> Gen.set (Range.linear 0 100) (genValue' st)
  STContract (s :: SingT p) -> case (checkOpPresence s, checkNestedBigMapsPresence s) of
    (OpAbsent, NestedBigMapsAbsent) -> VContract <$> genAddress <*> pure unsafeSepcCallRoot
    _ -> Gen.discard
  STTicket (s :: SingT p) ->
    VTicket <$> genAddress <*> genValue' s
            <*> Gen.integral (Range.linearFrom 0 0 1000)
  STPair STChestKey STChest -> VPair . bimap VChestKey VChest . swap <$> genChestAndKey
  STPair l r -> VPair <$> ((,) <$> genNoOpValue l <*> genNoOpValue r)
  STOr l r -> VOr <$> Gen.choice
    [ Left <$> genNoOpValue l
    , Right <$> genNoOpValue r
    ]
  -- It's quite hard to generate proper lambda of given type, so it always returns FAILWITH.
  -- Such implementation is sufficient for now.
  STLambda{} -> pure $ VLam $ RfAlwaysFails $ Seq UNIT FAILWITH
  STMap k v ->
    VMap <$> Gen.map (Range.linear 0 100) (liftA2 (,) (genNoOpValue k) (genNoOpValue v))
  STBigMap k v ->
    VBigMap
      <$> Gen.maybe genBigMapId
      <*> Gen.map (Range.linear 0 100) (liftA2 (,) (genNoOpValue k) (genNoOpValue v))
  STInt -> genValueInt
  STNat -> VNat <$> Gen.integral (Range.linearFrom 0 0 1000)
  STString -> VString . unsafe . mkMText . fromString <$> Gen.string (Range.linear 0 100) Gen.alphaNum
  STBytes -> VBytes <$> Gen.bytes (Range.linear 0 100)
  STMutez -> genValueMutez
  STKeyHash -> genValueKeyHash
  STBls12381Fr -> VBls12381Fr <$> genBls12381Fr
  STBls12381G1 -> VBls12381G1 <$> genBls12381G1
  STBls12381G2 -> VBls12381G2 <$> genBls12381G2
  -- Note that we also have a special case for a list of BLS12 pairings
  STTimestamp -> genValueTimestamp
  STAddress -> VAddress <$> genEpAddress
  STBool -> VBool <$> Gen.bool
  STChest -> VChest . fst <$> genChestAndKey
  STChestKey -> VChestKey . snd <$> genChestAndKey
  STNever -> Gen.discard
  STSaplingState _ -> error "genValue': Cannot generate `sapling_state` value."
  STSaplingTransaction _ -> error "genValue': Cannot generate `sapling_instruction` value."
  where
    genNoOpValue
      :: (MonadGen m, GenBase m ~ Identity, WellTyped t')
      => Sing t' -> m (Value' Instr t')
    genNoOpValue st = case checkOpPresence st of
      OpAbsent -> genValue' st
      _ -> Gen.discard

    genBls12Pairing
      :: MonadGen m
      => m (Value' Instr ('TList $ 'TPair 'TBls12381G1 'TBls12381G2))
    genBls12Pairing = Gen.frequency
      [ -- random pairing (likely incorrect one)
        ( 1
        , fmap toVal $ Gen.list (Range.linear 0 10) $
            (,) <$> genBls12381G1 <*> genBls12381G2
        )
      , -- correct pairing case
        ( 1
        , do
            g1 <- genBls12381G1
            g2 <- genBls12381G2
            let pairing = [(g1, g2), (g1, BLS.negate g2)]
            Ex.assert (BLS.checkPairing pairing) $
              return (toVal pairing)
        )
      ]