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

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

module Hedgehog.Gen.Michelson.Typed
  ( genBigMap
  , genEpAddress
  , genValueBigMap
  , genValueChestAndKey
  , genValueInt
  , genValueKeyHash
  , genValueList
  , genValueMap
  , genValueMutez
  , genValueNat
  , genValuePair
  , genValueSet
  , genValueString
  , genValueTimestamp
  , genValueUnit

  , 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 (Range)
import Hedgehog.Range qualified as Range

import Hedgehog.Gen.Michelson.Untyped (genEpName)
import Hedgehog.Gen.Tezos.Address
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'(..), mkVLam, 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.Address (Constrained(..), TxRollupL2Address(..))
import Morley.Tezos.Core (Mutez, Timestamp)
import Morley.Tezos.Crypto (parseHash)
import Morley.Tezos.Crypto.BLS12381 qualified as BLS

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

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

genValueMap
  :: forall (k :: T) (v :: T) m instr.
     (MonadGen m, WellTyped k, WellTyped v, Comparable k)
  => Range Length -> m (Value' instr k) -> m (Value' instr v) -> m (Value' instr ('TMap k v))
genValueMap len genKey genVal =
  VMap <$> Gen.map (unLength <$> len) (liftA2 (,) genKey genVal)

genValueBigMap
  :: forall (k :: T) (v :: T) m instr.
     (MonadGen m, WellTyped k, WellTyped v, HasNoBigMap v, Comparable k)
  => Range Length -> m (Value' instr k) -> m (Value' instr v) -> m (Value' instr ('TBigMap k v))
genValueBigMap len genKey genVal =
  VBigMap
    <$> Gen.maybe genBigMapId
    <*> Gen.map (unLength <$> len) (liftA2 (,) genKey genVal)

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 => Range Mutez -> m (Value' instr 'TMutez)
genValueMutez = fmap VMutez . genMutez

genValueInt :: MonadGen m => Range (Value' instr 'TInt) -> m (Value' instr 'TInt)
genValueInt range = VInt <$> Gen.integral (unVal <$> range)
  where
    unVal :: Value' instr 'TInt -> Integer
    unVal (VInt val) = val

genValueNat :: MonadGen m => Range (Value' instr 'TNat) -> m (Value' instr 'TNat)
genValueNat range = VNat <$> Gen.integral (unVal <$> range)
  where
    unVal :: Value' instr 'TNat -> Natural
    unVal (VNat val) = val

genValueString :: MonadGen f => Range Length -> f (Value' instr 'TString)
genValueString rangeLen =
  VString . unsafe. mkMText . fromString <$> Gen.string (unLength <$> rangeLen) Gen.alphaNum

genValueBytes :: MonadGen f => Range Length -> f (Value' instr 'TBytes)
genValueBytes rangeLen = VBytes <$> Gen.bytes (unLength <$> rangeLen)

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

genValueSet
  :: (MonadGen m, Comparable a, SingI a)
  => Range Length -> m (Value' instr a) -> m (Value' instr ('TSet a))
genValueSet rangeLen genA = VSet <$> Gen.set (unLength <$> rangeLen) 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 => Range Timestamp -> m (Value' instr 'TTimestamp)
genValueTimestamp = fmap VTimestamp . genTimestamp

genValueTicket
  :: (MonadGen m, Comparable a)
  => Range TicketAmount -> m (Value' instr a) -> m (Value' instr ('TTicket a))
genValueTicket range genVal = VTicket . MkAddress <$> genContractAddress <*> genVal
  <*> Gen.integral (unTicketAmount <$> range)

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 -> genValueList def $ genValue' st
  STSet st -> genValueSet def $ genValue' st
  STContract STUnit -> VContract <$> genAddress <*> pure unsafeSepcCallRoot
  STContract _ -> VContract <$> (MkAddress <$> genContractAddress) <*> pure unsafeSepcCallRoot
  STTicket s -> genValueTicket def $ genValue' s
  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 $ mkVLam $ RfAlwaysFails $ Seq UNIT FAILWITH
  STMap k v -> genValueMap def (genNoOpValue k) (genNoOpValue v)
  STBigMap k v -> genValueBigMap def (genNoOpValue k) (genNoOpValue v)
  STInt -> genValueInt def
  STNat -> genValueNat def
  STString -> genValueString def
  STBytes -> genValueBytes def
  STMutez -> genValueMutez def
  STKeyHash -> genValueKeyHash
  STTxRollupL2Address ->
    VTxRollupL2Address . TxRollupL2Address
    . unsafe . parseHash <$> Gen.element
      -- TODO [#839]: implement BLS public key generation
      [ "tz4UJqedFMBS7FjAqvZojJMPNd59MLm2hkuc"
      , "tz4Y7kRVfDH2XGQtjc19ppJqejL4CBVmxHED"
      , "tz4LVHYD4P4T5NHCuwJbxQvwVURF62seE3Qa"
      , "tz4MwL5iRbyHvVxH9N69GCeDmYCqbQtewr7R"
      , "tz4AZhg8GuahEs2Uo7dFZxVZwEgNKirtYMhY"
      , "tz4D77UuwdqbmDd7Xh9VNbFRjDiqWbBeWqud"
      , "tz4SYR9zvak9GohAEENUjPk7zAQo46wo6vNE"
      , "tz4UWDSphLswG5xtBwGnodCGL7FBzN21EKSQ"
      , "tz49e42Nbrc15PuT7RgkGqC6Xi3w5jEzEzH1"
      , "tz4F1Nd91Fc3BUCxTCnFaccjnn2hH4W3Bd8X"
      , "tz4MVuy2j5GCjPZQg7cxadPiNWq2nsRA392Y"
      , "tz4StvWhTeDnVpGspXKbfuhuVXhZ1jkAp7Yq"
      ]
  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 def
  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_transaction` 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)
        )
      ]