{-# OPTIONS --safe #-}
open import Ledger.Dijkstra.Specification.Abstract
open import Ledger.Dijkstra.Specification.Transaction
module Ledger.Dijkstra.Specification.PoolReap
  (txs : _) (open TransactionStructure txs)
  (abs : AbstractFunctions txs) (open AbstractFunctions abs)
  where
open import Ledger.Prelude
open import Ledger.Dijkstra.Specification.Certs govStructure
toRewardAddress : Credential → RewardAddress
toRewardAddress x = record { net = NetworkId ; stake = x }



record PoolReapState : Type where
  constructor ⟦_,_⟧ᵖ
  field
    certState  : CertState
    acnt       : Acnt

private variable
  e          : Epoch
  prs        : PoolReapState

data _⊢_⇀⦇_,POOLREAP⦈_ : ⊤ → PoolReapState → Epoch → PoolReapState → Type where
  POOLREAP :
    let cs = PoolReapState.certState prs
        ps = CertState.pState cs
        ds = CertState.dState cs
        a  = PoolReapState.acnt prs

        -- Pool key hashes retiring this epoch
        retired : ℙ KeyHash
        retired = PState.retiring ps ⁻¹ e

        -- For each retiring pool kh with deposit c, look up its reward-account
        -- credential from PState.pools and convert to a RewardAddress.
        -- mapPartial handles the case where kh might not be in pools (invariant:
        -- retiring pools are always registered, so lookupᵐ? always succeeds).
        poolPayoutPairs : ℙ (RewardAddress × Coin)
        poolPayoutPairs =
          mapPartial
            (λ (kh , c) →
              (λ params →
                toRewardAddress (StakePoolParams.rewardAccount params) , c)
              <$> lookupᵐ? (PState.pools ps) kh)
            ((PState.deposits ps ∣ retired) ˢ)

        -- Aggregate by reward address (multiple retiring pools may share one).
        poolPayout : RewardAddress ⇀ Coin
        poolPayout = aggregate₊ (poolPayoutPairs ᶠˢ)

        -- Registered reward accounts receive a refund; unregistered amounts go
        -- to the treasury as unclaimed.  Mirrors Conway's pullbackMap pattern.
        refunds : Credential ⇀ Coin
        refunds = pullbackMap poolPayout toRewardAddress (dom (DState.rewards ds))

        unclaimed : Coin
        unclaimed = getCoin poolPayout - getCoin refunds

        -- Updated DState: credit refunds to registered reward accounts.
        dState' : DState
        dState' = record ds { rewards = DState.rewards ds ∪⁺ refunds }

        -- Updated PState: remove retired pools and their deposits.
        pState' : PState
        pState' = record ps
          { pools    = PState.pools    ps ∣ retired ᶜ
          ; fPools   = PState.fPools   ps ∣ retired ᶜ
          ; retiring = PState.retiring ps ∣ retired ᶜ
          ; deposits = PState.deposits ps ∣ retired ᶜ
          }

        certState' : CertState
        certState' = record cs { dState = dState' ; pState = pState' }

        acnt' : Acnt
        acnt' = record a { treasury = Acnt.treasury a + unclaimed }
    in
      ────────────────────────────────
      _ ⊢ prs ⇀⦇ e ,POOLREAP⦈ $\begin{pmatrix} \,\href{Ledger.Dijkstra.Specification.PoolReap.html#2723}{\htmlId{2985}{\htmlClass{Bound}{\text{certState'}}}}\, \\ \,\href{Ledger.Dijkstra.Specification.PoolReap.html#2826}{\htmlId{2998}{\htmlClass{Bound}{\text{acnt'}}}}\, \end{pmatrix}$