use crate::events::FeesClaimed;
use crate::*;
use near_sdk::{env, json_types::U128, near, AccountId, PromiseError};

#[near]
impl Contract {
    #[private]
    pub fn claim_fees_callback(
        &mut self,
        #[callback_result] call_result: Result<(), PromiseError>,
        asset: Asset,
        recipient: AccountId,
        fees_claimed: U128,
        is_protocol_fee: bool,
    ) {
        match call_result {
            Ok(_) => {
                // Emit event on successful transfer
                FeesClaimed {
                    asset: &asset,
                    amount: fees_claimed,
                    recipient: &recipient,
                }
                .emit();

                if is_protocol_fee {
                    env::log_str(&format!(
                        "Claimed protocol fees, amount: {}, recipient: {}",
                        fees_claimed.0, recipient
                    ));
                } else {
                    env::log_str(&format!(
                        "Claimed fees, amount: {}, recipient: {}",
                        fees_claimed.0, recipient
                    ));
                }
            }
            Err(_) => {
                // Restore state on failure (using saturating_add since its a rollback)
                if is_protocol_fee {
                    let current_fees = self
                        .state
                        .protocol_fees_owed
                        .get(&asset)
                        .unwrap_or(&U128(0))
                        .0;
                    self.state.protocol_fees_owed.insert(
                        asset.clone(),
                        U128(current_fees.saturating_add(fees_claimed.0)),
                    );
                    env::log_str(&format!(
                        "Failed to transfer protocol fees, restored state. Protocol fees owed: {}",
                        current_fees.saturating_add(fees_claimed.0)
                    ));
                } else {
                    let current_fees = self.state.fees_owed.get(&asset).unwrap_or(&U128(0)).0;
                    self.state.fees_owed.insert(
                        asset.clone(),
                        U128(current_fees.saturating_add(fees_claimed.0)),
                    );
                    env::log_str(&format!(
                        "Failed to transfer fees, restored state. Fees owed: {}",
                        current_fees.saturating_add(fees_claimed.0)
                    ));
                }
            }
        }
    }
}