content: {- |
  Module     : Sample.Proposal.UnlockStake
  Maintainer : connor@mlabs.city
  Description: Generate sample data for testing the functionalities of unlocking stake and retracting votes
  
  Sample and utilities for testing the functionalities of unlocking stake and retracting votes
  -}
  module Sample.Proposal.Unlock (
    ParameterBundle (..),
    StakeRole (..),
    TimeRange (..),
    SignedBy (..),
    TransactionParameters (..),
    ProposalParameters (..),
    SSTOwner (..),
    StakeParameters (..),
    Validity (..),
    unlock,
    mkTestTree,
    mkValidVoterRetractVotes,
    mkValidDelegateeRetractVotes,
    mkValidVoterCreatorRetractVotes,
    mkValidCreatorRemoveLock,
    mkValidVoterRemoveLockAfterVoting,
    mkRetractVotesWhileNotVoting,
    mkUnockIrrelevantStakes,
    mkRemoveCreatorLockBeforeFinished,
    mkCreatorRetractVotes,
    mkChangeOutputStakeValue,
    mkUseFakeStakes,
    mkDisrespectCooldown,
  ) where
  
  --------------------------------------------------------------------------------
  
  import Agora.Governor (Governor (..))
  import Agora.Proposal (
    ProposalDatum (..),
    ProposalEffectGroup,
    ProposalId (..),
    ProposalRedeemer (UnlockStake),
    ProposalStatus (..),
    ProposalVotes (..),
    ResultTag (..),
   )
  import Agora.Proposal.Time (
    ProposalStartingTime (ProposalStartingTime),
    ProposalTimingConfig (..),
   )
  import Agora.SafeMoney (GTTag)
  import Agora.Stake (
    ProposalAction (Created, Voted),
    ProposalLock (..),
    StakeDatum (..),
    StakeRedeemer (RetractVotes),
   )
  import Data.Coerce (coerce)
  import Data.Default.Class (Default (def))
  import Data.Map.Strict qualified as StrictMap
  import Data.Tagged (Tagged, untag)
  import Plutarch.Context (
    input,
    normalizeValue,
    output,
    script,
    signedWith,
    timeRange,
    txId,
    withDatum,
    withRedeemer,
    withRef,
    withValue,
   )
  import Plutarch.Extra.AssetClass (assetClassValue)
  import Plutarch.Extra.ScriptContext (scriptHashToTokenName)
  import PlutusLedgerApi.V1.Value qualified as Value
  import PlutusLedgerApi.V2 (
    Credential (PubKeyCredential),
    POSIXTime,
    PubKeyHash,
    TxOutRef (..),
   )
  import Sample.Proposal.Shared (stakeTxRef)
  import Sample.Shared (
    governor,
    minAda,
    proposalAssetClass,
    proposalScriptHash,
    proposalValidator,
    stakeScriptHash,
    stakeSymbol,
    stakeValidator,
   )
  import Test.Specification (SpecificationTree, group, testValidator)
  import Test.Util (CombinableBuilder, closedBoundedInterval, mkSpending, pubKeyHashes)
  
  --------------------------------------------------------------------------------
  
  votesTemplate :: ProposalVotes
  votesTemplate =
    ProposalVotes $
      StrictMap.fromList
        [ (ResultTag 0, 0)
        , (ResultTag 1, 0)
        ]
  
  -- | Create empty effects for every result tag given the votes.
  emptyEffectFor ::
    ProposalVotes ->
    StrictMap.Map ResultTag ProposalEffectGroup
  emptyEffectFor (ProposalVotes vs) =
    StrictMap.fromList $
      map (,StrictMap.empty) (StrictMap.keys vs)
  
  -- | The default vote option that will be used by functions in this module.
  defVoteFor :: ResultTag
  defVoteFor = ResultTag 0
  
  -- | The default number of GTs the stake will have.
  defStakedGTs :: Tagged GTTag Integer
  defStakedGTs = 100000
  
  alteredStakedGTs :: Tagged GTTag Integer
  alteredStakedGTs = 100
  
  -- | Default owner of the stakes.
  defOwner :: PubKeyHash
  defOwner = pubKeyHashes !! 1
  
  defDelegatee :: PubKeyHash
  defDelegatee = pubKeyHashes !! 2
  
  defUnknown :: PubKeyHash
  defUnknown = pubKeyHashes !! 3
  
  defProposalId :: ProposalId
  defProposalId = ProposalId 0
  
  defStartingTime :: ProposalStartingTime
  defStartingTime = ProposalStartingTime 0
  
  --------------------------------------------------------------------------------
  
  data ParameterBundle = ParameterBundle
    { proposalParameters :: ProposalParameters
    , stakeParameters :: StakeParameters
    , transactionParameters :: TransactionParameters
    }
  
  data SignedBy = Owner | Delegatee | Unknown
  
  data TimeRange = WhileVoting {offset :: POSIXTime} | AfterVoting
  
  data TransactionParameters = TransactionParameters
    { signedBy :: SignedBy
    , timeRange :: TimeRange
    }
  
  data ProposalParameters = ProposalParameters
    { proposalStatus :: ProposalStatus
    , retractVotes :: Bool
    }
  
  -- | How a stake has been used on a particular proposal.
  data StakeRole
    = -- | The stake was spent to vote for a paraticular option.
      Voter
    | -- | The stake was used to create the proposal.
      Creator
    | -- | The stake was used to both create and vote for the proposal.
      Both
    | -- | The stake has nothing to do with the proposal.
      Irrelevant
    deriving stock (Bounded, Enum, Show)
  
  data SSTOwner
    = StakeValidator
    | Attacker
  
  data StakeParameters = StakeParameters
    { numStakes :: Integer
    , stakeRole :: StakeRole
    , removeVoterLock :: Bool
    , removeCreatorLock :: Bool
    , alterOutputValue :: Bool
    , sstOwner :: SSTOwner
    , votingLockCreatedAt :: POSIXTime
    }
  
  data Validity = Validity
    { forProposalValidator :: Bool
    , forStakeValidator :: Bool
    }
  
  --------------------------------------------------------------------------------
  
  mkStakeRef :: Integer -> TxOutRef
  mkStakeRef = TxOutRef stakeTxRef
  
  stakeRedeemer :: StakeRedeemer
  stakeRedeemer = RetractVotes
  
  mkStakeInputDatum :: StakeParameters -> StakeDatum
  mkStakeInputDatum ps =
    StakeDatum
      { stakedAmount = defStakedGTs
      , owner = PubKeyCredential defOwner
      , delegatedTo = Just $ PubKeyCredential defDelegatee
      , lockedBy = stakeLocks
      }
    where
      stakeLocks = mkStakeLocks' ps.stakeRole
  
      mkStakeLocks' Voter =
        [ ProposalLock defProposalId $
            Voted defVoteFor ps.votingLockCreatedAt
        ]
      mkStakeLocks' Creator = [ProposalLock defProposalId Created]
      mkStakeLocks' Both = mkStakeLocks' Voter <> mkStakeLocks' Creator
      mkStakeLocks' Irrelevant =
        let ProposalId pid = defProposalId
            ResultTag vid = defVoteFor
         in [ ProposalLock (ProposalId $ pid + 1) $
                Voted
                  (ResultTag $ vid + 1)
                  ps.votingLockCreatedAt
            , ProposalLock (ProposalId $ pid + 1) Created
            ]
  
  --------------------------------------------------------------------------------
  
  proposalRef :: TxOutRef
  proposalRef = TxOutRef stakeTxRef 0
  
  proposalRedeemer :: ProposalRedeemer
  proposalRedeemer = UnlockStake
  
  mkProposalInputDatum ::
    StakeParameters ->
    ProposalParameters ->
    ProposalDatum
  mkProposalInputDatum sps pps =
    ProposalDatum
      { proposalId = defProposalId
      , effects = emptyEffectFor votesTemplate
      , status = pps.proposalStatus
      , cosigners = [PubKeyCredential $ head pubKeyHashes]
      , thresholds = def
      , votes = updatVotes votesTemplate
      , timingConfig = def
      , startingTime = defStartingTime
      }
    where
      updatVotes (ProposalVotes vt) =
        ProposalVotes $
          StrictMap.adjust
            (+ sps.numStakes * untag defStakedGTs)
            defVoteFor
            vt
  
  --------------------------------------------------------------------------------
  
  unlock :: forall b. CombinableBuilder b => ParameterBundle -> b
  unlock ps = builder
    where
      pst = assetClassValue proposalAssetClass 1
  
      proposalInputDatum =
        mkProposalInputDatum
          ps.stakeParameters
          ps.proposalParameters
  
      proposalOutputDatum =
        if ps.proposalParameters.retractVotes
          then proposalInputDatum {votes = votesTemplate}
          else proposalInputDatum
  
      proposalValue = normalizeValue $ pst <> minAda
  
      proposalBuilder :: b
      proposalBuilder =
        mconcat
          [ input $
              mconcat
                [ script proposalScriptHash
                , withValue proposalValue
                , withDatum proposalInputDatum
                , withRef proposalRef
                , withRedeemer proposalRedeemer
                ]
          , output $
              mconcat
                [ script proposalScriptHash
                , withValue proposalValue
                , withDatum proposalOutputDatum
                ]
          ]
  
      ---
  
      sstName = case ps.stakeParameters.sstOwner of
        StakeValidator -> scriptHashToTokenName stakeScriptHash
        _ -> ""
  
      sst = Value.singleton stakeSymbol sstName 1
  
      stakeInputDatum = mkStakeInputDatum ps.stakeParameters
  
      -- TODO respect timing
      removeLocks v c =
        filter $ \(ProposalLock pid action) ->
          pid == defProposalId
            && not
              ( case action of
                  Voted _ _ -> v
                  _ -> c
              )
  
      stakeOutputDatum =
        stakeInputDatum
          { lockedBy =
              removeLocks
                ps.stakeParameters.removeVoterLock
                ps.stakeParameters.removeCreatorLock
                stakeInputDatum.lockedBy
          }
  
      mkStakeValue gt =
        normalizeValue $
          mconcat
            [ minAda
            , sst
            , assetClassValue
                governor.gtClassRef
                gt
            ]
  
      stakeInputValue = mkStakeValue defStakedGTs
  
      stakeOutputValue =
        mkStakeValue $
          if ps.stakeParameters.alterOutputValue
            then alteredStakedGTs
            else defStakedGTs
  
      stakeBuilder :: b
      stakeBuilder =
        foldMap
          ( \i ->
              mconcat
                [ input $
                    mconcat
                      [ script stakeScriptHash
                      , withValue stakeInputValue
                      , withDatum stakeInputDatum
                      , withRef $ mkStakeRef i
                      ]
                , output $
                    mconcat
                      [ script stakeScriptHash
                      , withValue stakeOutputValue
                      , withDatum stakeOutputDatum
                      ]
                ]
          )
          [1 .. ps.stakeParameters.numStakes]
  
      ---
  
      ProposalStartingTime s = defStartingTime
  
      time = case ps.transactionParameters.timeRange of
        WhileVoting offset ->
          let lb =
                ps.stakeParameters.votingLockCreatedAt
                  + offset
              ub =
                s
                  + (def :: ProposalTimingConfig).draftTime
                  + (def :: ProposalTimingConfig).votingTime
           in closedBoundedInterval (lb + 1) (ub - 1)
        AfterVoting ->
          let lb =
                s
                  + (def :: ProposalTimingConfig).draftTime
                  + (def :: ProposalTimingConfig).votingTime
              ub = lb + (def :: ProposalTimingConfig).lockingTime
           in closedBoundedInterval (lb + 1) (ub - 1)
  
      sig = case ps.transactionParameters.signedBy of
        Unknown -> defUnknown
        Owner -> defOwner
        Delegatee -> defDelegatee
  
      ---
  
      builder =
        mconcat
          [ txId "388bc0b897b3dadcd479da4c88291de4113a50b72ddbed001faf7fc03f11bc52"
          , proposalBuilder
          , stakeBuilder
          , signedWith sig
          , timeRange time
          ]
  
  --------------------------------------------------------------------------------
  
  {- | Create a test tree that runs both the stake validator and the proposal
       validator.
  -}
  mkTestTree :: String -> ParameterBundle -> Validity -> SpecificationTree
  mkTestTree name ps val = group name [stake, proposal]
    where
      spend = mkSpending unlock ps
  
      stake =
        testValidator
          val.forStakeValidator
          "stake"
          stakeValidator
          (mkStakeInputDatum ps.stakeParameters)
          stakeRedeemer
          (spend $ mkStakeRef 1)
  
      proposal =
        testValidator
          val.forProposalValidator
          "proposal"
          proposalValidator
          (mkProposalInputDatum ps.stakeParameters ps.proposalParameters)
          proposalRedeemer
          (spend proposalRef)
  
  --------------------------------------------------------------------------------
  
  mkValidVoterRetractVotes :: Integer -> ParameterBundle
  mkValidVoterRetractVotes i =
    ParameterBundle
      { proposalParameters =
          ProposalParameters
            { proposalStatus = VotingReady
            , retractVotes = True
            }
      , stakeParameters =
          StakeParameters
            { numStakes = i
            , stakeRole = Voter
            , removeVoterLock = True
            , removeCreatorLock = False
            , alterOutputValue = False
            , sstOwner = StakeValidator
            , votingLockCreatedAt =
                coerce defStartingTime
                  + (def :: ProposalTimingConfig).draftTime
                  + 1
            }
      , transactionParameters =
          TransactionParameters
            { signedBy = Owner
            , timeRange =
                WhileVoting
                  { offset =
                      coerce
                        (def :: ProposalTimingConfig).minStakeVotingTime
                        + 5
                  }
            }
      }
  
  mkValidDelegateeRetractVotes :: Integer -> ParameterBundle
  mkValidDelegateeRetractVotes i =
    let template = mkValidVoterRetractVotes i
     in template
          { transactionParameters =
              template.transactionParameters
                { signedBy = Delegatee
                }
          }
  
  mkValidVoterCreatorRetractVotes :: Integer -> ParameterBundle
  mkValidVoterCreatorRetractVotes i =
    let template = mkValidVoterRetractVotes i
     in template
          { stakeParameters =
              template.stakeParameters
                { stakeRole = Both
                }
          }
  
  mkValidCreatorRemoveLock :: Integer -> ParameterBundle
  mkValidCreatorRemoveLock i =
    let template = mkValidVoterRetractVotes i
     in template
          { proposalParameters =
              template.proposalParameters
                { proposalStatus = Finished
                , retractVotes = False
                }
          , stakeParameters =
              template.stakeParameters
                { stakeRole = Creator
                , removeCreatorLock = True
                }
          , transactionParameters =
              template.transactionParameters
                { timeRange = AfterVoting
                }
          }
  
  mkValidVoterRemoveLockAfterVoting :: Integer -> ParameterBundle
  mkValidVoterRemoveLockAfterVoting i =
    let template = mkValidVoterRetractVotes i
     in template
          { proposalParameters =
              template.proposalParameters
                { proposalStatus = Finished
                , retractVotes = False
                }
          , transactionParameters =
              template.transactionParameters
                { timeRange = AfterVoting
                }
          }
  
  mkRetractVotesWhileNotVoting :: Integer -> [ParameterBundle]
  mkRetractVotesWhileNotVoting i =
    let template = mkValidVoterRetractVotes i
     in map
          ( \s ->
              template
                { proposalParameters =
                    template.proposalParameters
                      { proposalStatus = s
                      }
                }
          )
          [Draft, Locked, Finished]
  
  mkUnockIrrelevantStakes :: Integer -> ParameterBundle
  mkUnockIrrelevantStakes i =
    let template = mkValidVoterRetractVotes i
     in template
          { stakeParameters =
              template.stakeParameters
                { stakeRole = Irrelevant
                , removeCreatorLock = True
                }
          }
  
  mkRemoveCreatorLockBeforeFinished :: Integer -> [ParameterBundle]
  mkRemoveCreatorLockBeforeFinished i =
    let template = mkValidCreatorRemoveLock i
     in map
          ( \s ->
              template
                { proposalParameters =
                    template.proposalParameters
                      { proposalStatus = s
                      }
                }
          )
          [Draft, VotingReady, Locked]
  
  mkCreatorRetractVotes :: Integer -> ParameterBundle
  mkCreatorRetractVotes i =
    let template = mkValidVoterRetractVotes i
     in template
          { proposalParameters =
              template.proposalParameters
                { proposalStatus = VotingReady
                }
          , stakeParameters =
              template.stakeParameters
                { stakeRole = Creator
                }
          }
  
  mkChangeOutputStakeValue :: Integer -> ParameterBundle
  mkChangeOutputStakeValue i =
    let template = mkValidVoterRetractVotes i
     in template
          { stakeParameters =
              template.stakeParameters
                { alterOutputValue = True
                }
          }
  
  mkUseFakeStakes :: Integer -> ParameterBundle
  mkUseFakeStakes i =
    let template = mkValidVoterCreatorRetractVotes i
     in template
          { stakeParameters =
              template.stakeParameters
                { sstOwner = Attacker
                }
          }
  
  mkDisrespectCooldown :: Integer -> ParameterBundle
  mkDisrespectCooldown i =
    let template = mkValidVoterCreatorRetractVotes i
     in template
          { transactionParameters =
              template.transactionParameters
                { timeRange =
                    WhileVoting
                      { offset =
                          coerce
                            (def :: ProposalTimingConfig).minStakeVotingTime
                            - 5
                      }
                }
          }
  