1use std::{future::Future, ops::Deref};
2
3use gmsol_programs::gmsol_store::{
4 accounts::{ReferralCodeV2, UserHeader},
5 client::{accounts, args},
6};
7use gmsol_solana_utils::transaction_builder::TransactionBuilder;
8use gmsol_utils::pubkey::optional_address;
9use solana_sdk::{pubkey::Pubkey, signer::Signer, system_program};
10
11use crate::{pda::ReferralCodeBytes, utils::zero_copy::ZeroCopy};
12
13pub trait UserOps<C> {
15 fn prepare_user(&self, store: &Pubkey) -> crate::Result<TransactionBuilder<C>>;
17
18 fn initialize_referral_code(
20 &self,
21 store: &Pubkey,
22 code: ReferralCodeBytes,
23 ) -> crate::Result<TransactionBuilder<C>>;
24
25 fn set_referrer(
27 &self,
28 store: &Pubkey,
29 code: ReferralCodeBytes,
30 hint_referrer: Option<Pubkey>,
31 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
32
33 fn transfer_referral_code(
35 &self,
36 store: &Pubkey,
37 receiver: &Pubkey,
38 hint_code: Option<ReferralCodeBytes>,
39 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
40
41 fn cancel_referral_code_transfer(
43 &self,
44 store: &Pubkey,
45 hint_code: Option<ReferralCodeBytes>,
46 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
47
48 fn accept_referral_code(
50 &self,
51 store: &Pubkey,
52 code: ReferralCodeBytes,
53 hint_owner: Option<Pubkey>,
54 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
55}
56
57impl<C: Deref<Target = impl Signer> + Clone> UserOps<C> for crate::Client<C> {
58 fn prepare_user(&self, store: &Pubkey) -> crate::Result<TransactionBuilder<C>> {
59 let owner = self.payer();
60 let user = self.find_user_address(store, &owner);
61 let rpc = self
62 .store_transaction()
63 .anchor_accounts(accounts::PrepareUser {
64 owner,
65 store: *store,
66 user,
67 system_program: system_program::ID,
68 })
69 .anchor_args(args::PrepareUser {});
70 Ok(rpc)
71 }
72
73 fn initialize_referral_code(
74 &self,
75 store: &Pubkey,
76 code: ReferralCodeBytes,
77 ) -> crate::Result<TransactionBuilder<C>> {
78 let owner = self.payer();
79 let referral_code = self.find_referral_code_address(store, code);
80 let user = self.find_user_address(store, &owner);
81 let rpc = self
82 .store_transaction()
83 .anchor_accounts(accounts::InitializeReferralCode {
84 owner,
85 store: *store,
86 referral_code,
87 user,
88 system_program: system_program::ID,
89 })
90 .anchor_args(args::InitializeReferralCode { code });
91 Ok(rpc)
92 }
93
94 async fn set_referrer(
95 &self,
96 store: &Pubkey,
97 code: ReferralCodeBytes,
98 hint_referrer_user: Option<Pubkey>,
99 ) -> crate::Result<TransactionBuilder<C>> {
100 let owner = self.payer();
101 let user = self.find_user_address(store, &owner);
102
103 let referral_code = self.find_referral_code_address(store, code);
104
105 let referrer_user = match hint_referrer_user {
106 Some(referrer) => referrer,
107 None => {
108 let code = self
109 .account::<ZeroCopy<ReferralCodeV2>>(&referral_code)
110 .await?
111 .ok_or(crate::Error::NotFound)?
112 .0;
113 let owner = code.owner;
114 self.find_user_address(store, &owner)
115 }
116 };
117
118 let rpc = self
119 .store_transaction()
120 .anchor_accounts(accounts::SetReferrer {
121 owner,
122 store: *store,
123 user,
124 referral_code,
125 referrer_user,
126 })
127 .anchor_args(args::SetReferrer { code });
128
129 Ok(rpc)
130 }
131
132 async fn transfer_referral_code(
133 &self,
134 store: &Pubkey,
135 receiver: &Pubkey,
136 hint_code: Option<ReferralCodeBytes>,
137 ) -> crate::Result<TransactionBuilder<C>> {
138 let owner = self.payer();
139 let user = self.find_user_address(store, &owner);
140 let receiver_user = self.find_user_address(store, receiver);
141
142 let referral_code = match hint_code {
143 Some(code) => self.find_referral_code_address(store, code),
144 None => {
145 let user = self
146 .account::<ZeroCopy<UserHeader>>(&user)
147 .await?
148 .ok_or(crate::Error::NotFound)?;
149 *optional_address(&user.0.referral.code)
150 .ok_or(crate::Error::custom("referral code is not set"))?
151 }
152 };
153
154 let rpc = self
155 .store_transaction()
156 .anchor_accounts(accounts::TransferReferralCode {
157 owner,
158 store: *store,
159 user,
160 referral_code,
161 receiver_user,
162 })
163 .anchor_args(args::TransferReferralCode {});
164
165 Ok(rpc)
166 }
167
168 async fn cancel_referral_code_transfer(
169 &self,
170 store: &Pubkey,
171 hint_code: Option<ReferralCodeBytes>,
172 ) -> crate::Result<TransactionBuilder<C>> {
173 let owner = self.payer();
174 let user = self.find_user_address(store, &owner);
175
176 let referral_code = match hint_code {
177 Some(code) => self.find_referral_code_address(store, code),
178 None => {
179 let user = self
180 .account::<ZeroCopy<UserHeader>>(&user)
181 .await?
182 .ok_or(crate::Error::NotFound)?;
183 *optional_address(&user.0.referral.code)
184 .ok_or(crate::Error::custom("referral code is not set"))?
185 }
186 };
187
188 let rpc = self
189 .store_transaction()
190 .anchor_accounts(accounts::CancelReferralCodeTransfer {
191 owner,
192 store: *store,
193 user,
194 referral_code,
195 })
196 .anchor_args(args::CancelReferralCodeTransfer {});
197
198 Ok(rpc)
199 }
200
201 async fn accept_referral_code(
202 &self,
203 store: &Pubkey,
204 code: ReferralCodeBytes,
205 hint_owner: Option<Pubkey>,
206 ) -> crate::Result<TransactionBuilder<C>> {
207 let next_owner = self.payer();
208 let receiver_user = self.find_user_address(store, &next_owner);
209 let referral_code = self.find_referral_code_address(store, code);
210
211 let owner = match hint_owner {
212 Some(owner) => owner,
213 None => {
214 let code = self
215 .account::<ZeroCopy<ReferralCodeV2>>(&referral_code)
216 .await?
217 .ok_or(crate::Error::NotFound)?
218 .0;
219 code.owner
220 }
221 };
222
223 let user = self.find_user_address(store, &owner);
224
225 let rpc = self
226 .store_transaction()
227 .anchor_accounts(accounts::AcceptReferralCode {
228 next_owner,
229 store: *store,
230 user,
231 referral_code,
232 receiver_user,
233 })
234 .anchor_args(args::AcceptReferralCode {});
235 Ok(rpc)
236 }
237}